#!/usr/bin/env python
import sys,os,shutil,math
from w2k_utils import *
from struct_utils import *
from io_utils import *
from constants import * 


def f_Help_Info():
  myname = sys.argv[0]
  print "\n"+myname + ": a Python script to extract the work function \n\
(ionization potential for insulating systems) results from a wien2k slab calculation \n"
  print "Usage: " + os.path.basename(myname) + " [options]"\
  + """
  Options:  
    -h               # print out this information 
    -D               # set debug mode 
    -b <bulk>        # the prefix for the bulk calculation (including both path and the case name)  
                        if not defined, then the bulk correction is not taken into account 
    -c <core>        # which core state is used as the reference ( the default 1S ) 
    -f <case_name>   # case name for the slab calculation 
    
    -gw <igw, gwcase, typegw>  # the method to apply GW correction  
                     #  igw == 1 -- GW-VBM scheme 
                     #  igw == 2 -- GW-BGC scheme
                     #  gwdir: the directory+casename where the .outgw file is saved
                     #  typegw: the GW scheme of the correction
                     #    typegw == 1 --  G0W0
                     #    typegw == 2 --  GW0
    -irb <ir_bulk>   # the index of the reference atom in the bulk
    -irs <ir_slab>   # the index of the reference atom in the slab 
    -iz <iz=3>       # the normal direction of the slab (1/2/3 for x/y/z)
                     #  using -3/-2/-1 if the slab is not centered at 0.5, but at 0.0  
    -r <ref_atom>    # the atom whose deepest core is used as the reference 
    -nr <nat_ref>    # the number of atoms used for the reference 
  """
  sys.exit(0)

if f_Getopt('-h',0,False): f_Help_Info()

gw_def = [0, '', 0]
gwval = 0.0

debug      = f_Getopt('-D',   0, False)
case_name  = f_Getopt('-f',   1,  '' )
bulk_name  = f_Getopt('-b',   1, ''   ) 
core       = f_Getopt('-c',   1, '1S' )
gw         = f_Getopt('-gw',  3, gw_def ) 
iz         = f_Getopt('-iz',  1, 3   ) 
ir_bulk    = f_Getopt('-irb', 1, 0    ) 
ir_slab    = f_Getopt('-irs', 1, 0    ) 
ref_atom   = f_Getopt('-r',   1, ''   ) 
nat_ref    = f_Getopt('-nr',  1, 1    ) 

case_name = f_Check_Name(case_name)

fn_scf  = case_name+".scf"
fn_scfc = case_name+".scfc"

(igw, gwcase, typegw ) = ( gw[0], gw[1], gw[2]) 

(latt_type,latt,struct_slab) = f_Read_Struct_w2k(case_name,mode=1)
print "Structure of the slab model:"
print "  Lattice constants (a,b,c): (%8.3f,%8.3f,%8.3f)"%(latt[0],latt[1],latt[2])
print " #   atom   mult    atomic positions:"
nat = len(struct_slab)
for ia in range(nat):
  atom = struct_slab[ia][0]
  neq = len(struct_slab[ia][1])
  info = "%4d %4s %4d "%(ia+1,atom,neq) 
  for ieq in range(neq):
    xyz = struct_slab[ia][1][ieq]
    info = info+"(%6.4f, %6.4f, %6.4f) "%(xyz[0],xyz[1],xyz[2])
  print info 

print ''

# get v_vac
if iz == 3:
  v_vac = float( io_grep_lines( fn_scf,':VZERO',-1,3) )
elif iz == -3:
  v_vac = float( io_grep_lines( fn_scf,':VZERO',-1,7) )
elif iz == 2:
  v_vac = float( io_grep_lines( fn_scf,':VZERY',-1,3) )
elif iz == -2:
  v_vac = float( io_grep_lines( fn_scf,':VZERY',-1,7) )
elif iz == 1:
  v_vac = float( io_grep_lines( fn_scf,':VZERX',-1,3) )
elif iz == -1:
  v_vac = float( io_grep_lines( fn_scf,':VZERX',-1,7) )

if debug: print " Read V_vac=%12.6f Ry"%(v_vac)

# get evbm_slab
evbm_slab = float( io_grep_lines( fn_scf, ':FER ',-1, 10))

# get egap_slab
egap_slab = w2k_get_egap(case_name) 

if debug: print " Read e_VBM=%12.6f"%(evbm_slab)

IP1 = (v_vac-evbm_slab)* Ry2eV  
EBGC = IP1 - 0.5*egap_slab

# get gwval from gwdir
# GW-VBM correction mode
if igw == 1:
  print "GW correction: using VBM scheme"
  gwdir = os.path.abspath(gwcase+'.struct')[:-len(os.path.basename(gwcase+'.struct'))]
  nvel = int(io_grep_lines(gwdir+'gw.log', 'nvel ',1,3))
  nvbm = math.ceil(nvel/2)
  if typegw == 0:
    typegw = 1   # default is G0W0 correction
  if typegw == 1:
    fn_gw = gwcase+'.eqpeV_GW'
  elif typegw == 2:
    fn_gw = gwcase+'.epqeV_GW0'
  with open(fn_gw,'r') as hd_fngw:
    while True:
      gwline = hd_fngw.readline()
      if gwline == '':
        break
      if gwline.strip().startswith('#'):
        continue
      (ik,ie,eks,egw,dqp) = (gwline.split()[i] for i in [0,1,2,3,8]) # iband,E_KS,E_GW,DE_GW
      if int(ie) != nvbm:
        continue
# if the vbm states are the same in KS and GW
      if float(eks) == 0.0 and float(egw) == 0.0:
        gwval = float(dqp)
        print "VBM correction (band %d at ik %d): %10.4f eV" % (int(ie),int(ik),gwval)
        break
# if not
      elif float(egw) == 0.0:
        gwval = float(eks) + float(egw)
        print "VBM correction (band %d at ik %d): %10.4f eV" % (int(ie),int(ik),gwval)
        break
elif igw == 2:
  print "GW-BGC correction not implemented yet. Pass"
  igw = 0
else:
  print "Unsupported tag for GW correction. Pass"
  igw = 0

# add GW corrections
if igw == 1: 
  IP1gw = IP1 - gwval 
elif igw == 2 :
  IP1gw = IP1 + 0.5*(gwval - egap_slab)

if bulk_name == '':  # no bulk corrections
  if igw == 0 :
    print "Ionization potential from the slab model only\n%12s %12s %8s %8s %8s" %("V_vac/Ry","e_VBM/Ry","IP1/eV", "-E_BGC","Eg(slab)")
    print "%12.6f %12.6f %8.3f %8.3f %8.3f"%(v_vac, evbm_slab, IP1, EBGC, egap_slab)
  else:
    print "Ionization potential from the slab model only\n%12s %12s %8s %8s %8s %8s" %("V_vac/Ry","e_VBM/Ry","IP1/eV", "-E_BGC","Eg(slab)","IP1GW")
    print "%12.6f %12.6f %8.3f %8.3f %8.3f %8.3f"%(v_vac, evbm_slab, IP1, EBGC, egap_slab, IP1gw)

else:  # with bulk corrections 
  if ref_atom != '':
    ir_bulk = 0 
    struct_bulk = f_Read_Struct_w2k(bulk_name,mode=1)[2]

    # get the index of the reference atom in the bulk structure 
    for ia in range(len(struct_bulk)):
      if struct_bulk[ia][0] == ref_atom: 
        ir_bulk = ia + 1
        break 
    if ir_bulk == 0: 
      print "ERROR: can not find the target atom in the bulk structure"
      sys.exit(1)
    
    # get the index of the reference atom in the slab structure 

    ir_slab = [] 
    for ir in range(nat_ref):
      ir_slab.append(0)

    for ir in range(nat_ref): 
      dz_min = 100.0 
      for ia in range(nat):
        atom = struct_slab[ia][0]
        xyz = struct_slab[ia][1][0]
        if atom != ref_atom: continue 

        if   iz == 3:     dz = abs(xyz[2] - 0.5 )
        elif iz == -3:    dz = abs(xyz[2]) 
        elif iz == 2:     dz = abs(xyz[1] - 0.5 )
        elif iz == -2:    dz = abs(xyz[1])
        elif iz == 1:     dz = abs(xyz[0] - 0.5 )
        elif iz == -1:    dz = abs(xyz[0])

        if iz < 0 and dz > 0.5: 
          dz = 1.0 - dz  

        if dz < dz_min and  not (ia+1 in ir_slab):
          ir_slab[ir] = ia + 1 
          dz_min = dz 

    if ir_bulk == 0: 
      print "ERROR: can not find the target atom in the bulk structure"
      sys.exit(1)

    print "The reference atom:",ref_atom,"\n the index in the bulk:",ir_bulk, "\n the index in the slab:",ir_slab 
    for ir in range(nat_ref):
      xyz = struct_slab[ir_slab[ir]-1][1][0]
      print " the coordinate of the reference atom in the slab (%8.4f %8.4f %8.4f)"%(xyz[0],xyz[1],xyz[2])

  evbm_bulk = float( io_grep_lines(bulk_name+".scf", ':FER ',-1, 10) )
  egap_bulk = w2k_get_egap(bulk_name) 

  ec_slab = 0.0 
  for ir in range(nat_ref):
    ec_slab += w2k_get_ecore(case_name,ir_slab[ir],core)
  ec_slab /= nat_ref

  ec_bulk = w2k_get_ecore(bulk_name, ir_bulk, core)
  IP2= ((v_vac-ec_slab)-(evbm_bulk-ec_bulk))*Ry2eV
  print "\nData from slab calculations (all in units of Ry)"
  print "%16s %16s %16s %16s %16s"%("V_vac(slab)","e_VBM(slab)","e_c(slab)","e_vbm(bulk)","e_c(bulk)")
  print "%16.6f %16.6f %16.6f %16.6f %16.6f"%(v_vac, evbm_slab, ec_slab, evbm_bulk, ec_bulk)

  if igw > 0 :
    if igw == 1:
      IP2gw = IP2 - gwval
    else:
      IP2gw = IP2 + 0.5*(gwval - egap_bulk)

    print "\nIonization potential with GW correction (igw=%d), all in units of eV"%(igw) 
    print "\nIonization potential all in units of eV\n%12s %12s %12s %12s %12s %12s %12s" %("Eg(bulk)", "Eg(slab)","IP1","IP2","-E_BGC","IP1GW","IP2GW")   
    print "%12.3f %12.3f %12.3f %12.3f %12.3f %12.3f %12.3f"%(egap_bulk, egap_slab,  IP1, IP2,EBGC, IP1gw, IP2gw)

  else:    
    print "\nIonization potential all in units of eV\n%12s %12s %12s %12s %12s" %("Eg(bulk)", "Eg(slab)","IP1","IP2","-E_BGC")   
    print "%12.3f %12.3f %12.3f %12.3f %12.3f"%(egap_bulk, egap_slab, IP1, IP2,EBGC)
