# Author: Samuel Ponc\'e
# Date: 30/04/2013
# Script to compute the spectral function

import sys
import os
from rf_mods import system
try:
  import numpy as N
except ImportError:
  import warnings
  warnings.warn("The numpy module is missing!")
  raise
try:
  import netCDF4 as nc
except ImportError:
  import warnings
  warnings.warn("The netCDF4 module is missing!")
  raise

#import matplotlib.pyplot as plt

#############
# Constants #
#############
# If you want to hardcode the weight of the k-points you can do it here:
# wtq = [0.00, 0.125, 0.5,0.375]

tol6 = 1E-6
tol8 = 1E-8
Ha2eV = 27.21138386
kb_HaK = 3.1668154267112283e-06

# Interaction with the user
print '\n################################################'
print '# Spectral function of the dynamical self-energy #'
print '##################################################'
print '\nThis script compute the zero-point motion and the temperature dependence \n\
of eigenenergies due to electron-phonon interaction. This script can \n\
only compute Q-points with the same weight. If you want symmetry you must hack the script.\n\
WARNING: The first Q-point MUST be the Gamma point\n'

# Define the output file name
user_input = raw_input('Enter name of the output file\n')
output = user_input

# Get the value of the smearing parameter (in eV)
user_input = raw_input('Enter value of the smearing parameter (in eV)\n')
smearing = N.float(user_input)
smearing = smearing/Ha2eV

# Frequency range?
user_input = raw_input('What is the range of frequencies you want to compute your spectral function upon (in eV)\n\
  [start end steps]. e.g. -15 10 0.5 for an energy step every 0.5 eV between -15 eV and 10 eV. \n')
freq = N.arange(N.float(user_input.split()[0]),N.float(user_input.split()[1]),N.float(user_input.split()[2]))
freq = freq/Ha2eV # From eV to Ha

# Temperature
user_input = raw_input('Enter the temperature for the spectral function A_nk(omega,T) [in K]\n')
temperature = N.float(user_input)

# Band index
user_input = raw_input('Enter the band index for the spectral function A_nk(omega,T)\n')
try:
  band = N.int(user_input)
except ValueError:
  raise Exception('The value you enter is not an integer!')

# Get the nb of random Q-points from user 
user_input = raw_input('Enter the number of random Q-points you have\n')
try:
  nbQ = int(user_input)
except ValueError:
  raise Exception('The value you enter is not an integer!')

# Get the path of the DDB files from user
user_input = raw_input('Enter the name of the %s DDB files separated by a space\n' %nbQ)
if len(user_input.split()) != nbQ:
  raise Exception("You sould provide %s DDB files" %nbQ)
else:
  DDB_files = user_input.split()

# Test if the first file is at the Gamma point
DDBtmp = system(directory='.',filename=DDB_files[0])
if N.allclose(DDBtmp.iqpt,[0.0,0.0,0.0]) == False:
  raise Exception('The first Q-point is not Gamma!')

# Choose a k-point in the list below:
print 'Choose a k-point number in the list below for A_nk(omega,T)\n'
for ii in N.arange(DDBtmp.nkpt):
  print '%s) %s' % (ii,DDBtmp.kpt[ii,:])
user_input = raw_input('Enter the number of the k-point you want to analyse\n')
try:
  kpt = N.int(user_input)
except ValueError:
  raise Exception('The value you enter is not an integer!')

# Take the EIG at Gamma
user_input = raw_input('Enter the name of the unperturbed EIG.nc file at Gamma\n')
if len(user_input.split()) != 1:
  raise Exception("You sould only provide 1 file")
else:
  eig0 = system(directory='.',filename=user_input)

N.arange# Find the degenerate eigenstates
DDB = system(directory='.',filename=DDB_files[0])
degen =  N.zeros((DDB.nkpt,DDB.nband),dtype=int)
for ikpt in N.arange(DDB.nkpt):
  count = 0
  for iband in N.arange(DDB.nband):
    if iband != DDB.nband-1:
      if N.allclose(eig0.EIG[0,ikpt,iband+1], eig0.EIG[0,ikpt,iband]):
        degen[ikpt,iband] = count
      else:
        degen[ikpt,iband] = count
        count += 1
        continue
    else:
      if N.allclose(eig0.EIG[0,ikpt,iband-1], eig0.EIG[0,ikpt,iband]):
        degen[ikpt,iband] = count   
    if iband != 0:
      if N.allclose(eig0.EIG[0,ikpt,iband-1], eig0.EIG[0,ikpt,iband]):
        degen[ikpt,iband] = count
    else:
      if N.allclose(eig0.EIG[0,ikpt,iband+1], eig0.EIG[0,ikpt,iband]):
        degen[ikpt,iband] = count

# Create the random Q-integration (wtq=1/nqpt):
wtq = N.ones((nbQ))
wtq = wtq*(1.0/nbQ)


# Compute phonon freq. and eigenvector for each Q-point 
# from each DDB (1 qpt per DDB file)
fan_tot = N.zeros((len(freq)),dtype=complex)
ddw_tot = N.zeros((),dtype=complex)
self_energy = N.zeros((len(freq)),dtype=complex)
spectral = N.zeros((len(freq)))
iiqpt = 0
DDB = system()
FANterm = system()
EIGR2D = system()
eigq = system()
#Compute the spectral function
for ii in DDB_files:
  DDB.__init__(directory='.',filename=ii)

# Calcul of gprimd from rprimd
  rprimd = DDB.rprim*DDB.acell
  gprimd = N.linalg.inv(N.matrix(rprimd))

# Transform from 2nd-order matrix (non-cartesian coordinates, 
# masses not included, asr not included ) from DDB to
# dynamical matrix, in cartesian coordinates, asr not imposed.
  IFC_cart = N.zeros((3,DDB.natom,3,DDB.natom),dtype=complex)
  for ii in N.arange(DDB.natom):
    for jj in N.arange(DDB.natom):
      for dir1 in N.arange(3):
        for dir2 in N.arange(3):
          for dir3 in N.arange(3):
            for dir4 in N.arange(3):
              IFC_cart[dir1,ii,dir2,jj] += gprimd[dir1,dir3]*DDB.IFC[dir3,ii,dir4,jj] \
            *gprimd[dir2,dir4]

# Reduce the 4 dimensional IFC_cart matrice to 2 dimensional Dynamical matrice.
  ipert1 = 0
  Dyn_mat = N.zeros((3*DDB.natom,3*DDB.natom),dtype=complex)
  while ipert1 < 3*DDB.natom:
    for ii in N.arange(DDB.natom):
      for dir1 in N.arange(3):
        ipert2 = 0
        while ipert2 < 3*DDB.natom:
          for jj in N.arange(DDB.natom):
            for dir2 in N.arange(3):
              Dyn_mat[ipert1,ipert2] = IFC_cart[dir1,ii,dir2,jj]
              ipert2 += 1
        ipert1 += 1

# Hermitianize the dynamical matrix
  dynmat = N.matrix(Dyn_mat)
  dynmat = 0.5*(dynmat + dynmat.transpose().conjugate())

# Solve the eigenvalue problem with linear algebra (Diagonalize the matrix)
  [eigval,eigvect]=N.linalg.eigh(Dyn_mat)

# Orthonormality relation 
  eigvect = (eigvect)*N.sqrt(5.4857990965007152E-4/float(DDB.amu[0]))

# Phonon frequency (5.4857990946E-4 = 1 au of electron mass)
  omega = N.sqrt((eigval*5.4857990965007152E-4)/float(DDB.amu[0]))

# Now read the EIGq, EIGR2D and FAN
  user_input = raw_input('Enter the name of the the _EIG.nc that contain\n\
 the %s Q-points. \n' %DDB.iqpt)
  if len(user_input.split()) != 1:
    raise Exception("You sould provide only 1 ***_EIG.nc file" )
  else:
    eigq.__init__(directory='.',filename=user_input.split()[0])

  user_input = raw_input('Enter the name of the the EIGR2D that contain\n\
 the %s Q-points. \n' %DDB.iqpt)
  if len(user_input.split()) != 1:
    raise Exception("You sould provide only 1 ***_EIGR2D file" )
  else:
    EIGR2D.__init__(directory='.',filename=user_input.split()[0])

  user_input = raw_input('Enter the name of the the _FAN that contain\n\
 the %s Q-points. \n' %DDB.iqpt)
  if len(user_input.split()) != 1:
    raise Exception("You sould provide only 1 ***_FAN file" )
  else:
    FANterm.__init__(directory='.',filename=user_input.split()[0])

# Compute the displacement = eigenvectors of the DDB. 
# Due to metric problem in reduce coordinate we have to work in cartesian
# but then go back to reduce because our EIGR2D matrix elements are in reduced coord.
  displ_FAN =  N.zeros((3,3),dtype=complex)
  displ_DDW =  N.zeros((3,3),dtype=complex)
  if N.allclose(EIGR2D.iqpt,[0.0,0.0,0.0]):
    ddw_save = N.zeros((3,EIGR2D.natom,3,EIGR2D.natom),dtype=complex)
  fan_add = N.zeros((len(freq)),dtype=complex)
  ddw_corr = N.zeros((),dtype=complex)

  for imode in N.arange(3*EIGR2D.natom): #Loop on perturbation (6 for 2 atoms)
    if omega[imode].real > tol6:
      for iatom1 in N.arange(EIGR2D.natom):
        for iatom2 in N.arange(EIGR2D.natom):
          for idir1 in N.arange(0,3):
            for idir2 in N.arange(0,3):
              displ_DDW[idir1,idir2] = (eigvect[3*iatom2+idir2,imode].conj()\
                 *eigvect[3*iatom2+idir1,imode]+eigvect[3*iatom1+idir2,imode].conj()\
                 *eigvect[3*iatom1+idir1,imode])/(4.0*omega[imode].real)
          # Now switch to reduced coordinates in 2 steps (more efficient)
          tmp_displ_DDW = N.zeros((3,3),dtype=complex)
          for idir1 in N.arange(3):
            for idir2 in N.arange(3):
              tmp_displ_DDW[:,idir1] = tmp_displ_DDW[:,idir1]+displ_DDW[:,idir2]*gprimd[idir2,idir1]
          displ_red_DDW = N.zeros((3,3),dtype=complex)
          for idir1 in N.arange(3):
            for idir2 in N.arange(3):
              displ_red_DDW[idir1,:] = displ_red_DDW[idir1,:] + tmp_displ_DDW[idir2,:]*gprimd[idir2,idir1]
          # Now compute the T=0 shift due to this q point
          for idir1 in N.arange(3):
            for idir2 in N.arange(3):
              # DDW matrix only computed at Gamma  
              if N.allclose(EIGR2D.iqpt,[0.0,0.0,0.0]):
                ddw_save[idir1,iatom1,idir2,iatom2] = EIGR2D.EIG2D[kpt,band-1,idir1,iatom1,idir2,iatom2]
                ddw_corr += ddw_save[idir1,iatom1,idir2,iatom2]*\
                    displ_red_DDW[idir1,idir2]
              else:
                ddw_corr += ddw_save[idir1,iatom1,idir2,iatom2]*\
                    displ_red_DDW[idir1,idir2]
      if temperature < tol6:
        bose = 0
      else:
        bose = 1.0/(N.exp(omega[imode].real/(kb_HaK*temperature))-1)
      ddw_corr = ddw_corr*(2*bose+1.0)
      for jband in N.arange(EIGR2D.nband):
        index = 0
        for ifreq in freq:
          delta_E = ifreq - eigq.EIG[0,ikpt,jband] + smearing*1j
          fan_add[index] = fan_add[index] + FANterm.FAN[kpt,band-1,imode,jband]*(\
                                  (bose+0.5)*(2*delta_E/(delta_E**2-(omega[imode].real)**2)) \
                                - (1-EIGR2D.occ[iband])*(omega[imode].real/(delta_E**2-(omega[imode].real)**2)))\
                                 /(2.0*omega[imode].real)
          index += 1

  fan_tot += fan_add*wtq[iiqpt] # for each freq.
  ddw_tot += ddw_corr*wtq[iiqpt]
  iiqpt +=1


print 'eig0.EIG[0,kpt,band-1]',eig0.EIG[0,kpt,band-1]*Ha2eV
# Computation of the self-energy and the spectral function at a given temperature.
index = 0
with open(output,"w") as O:
  O.write('# Spectral function at T = '+str(temperature)+' for band = '+str(band)+' and kpt = '+str(DDBtmp.kpt[kpt,:])+'\n')
  O.write('# Omega [eV]  Spectral function\n')
  for ifreq in freq:
    self_energy[index] = 1.0/(ifreq-eig0.EIG[0,kpt,band-1]-fan_tot[index]-ddw_tot)
    spectral[index]= (1.0/N.pi)*N.abs((self_energy[index]).imag)
    O.write(str(ifreq*Ha2eV)+' '+str(spectral[index])+'\n')
    index +=1

# Make a matplotlibplot
#plt.plot(freq*Ha2eV,spectral)
#plt.ylabel('Spectral function')
#plt.xlabel('Energy [eV]')
#plt.show()















