#!/usr/bin/python
# Find the top of the valence band and the bottom of the conduction band from EIGENVAL  
# Usage: chmod +x gbandedges
#        gbandedges [path] [ef]
#        ef: trial Fermi energy
#        ef read from OUTCAR if not given
# Last modified on Jun.21,2014

from sys import argv,exit
from os import path

from tklib.tkapplication import tkApplication
from tklib.tkutils import terminate, pint, pfloat, getarg, getintarg, getfloatarg


app    = tkApplication()
logfile = 'gbandedges-out.txt'
print(f"Open logfile [{logfile}]")
app.redirect(targets = ["stdout", logfile], mode = 'w')


ef = 0.0
ief = 0
if len(argv) == 1:
    dir = '.'
elif len(argv) == 2:
    try:
        ef = float(argv[1])
        dir = '.'
        ief = 1
    except: 
        dir = argv[1].rstrip('/')
elif len(argv) == 3:
    dir = argv[1].rstrip('/')
    ef  = float(argv[2])
    ief = 1
else:
    print("Usage: gbandedges [path] [Ef]")
    print("       ef: trial Fermi energy")
    print("       ef will be taken from OUTCAR if not given")
    exit(0)

fn_outcar = path.join(dir, 'OUTCAR')
fn_eig    = path.join(dir, 'EIGENVAL')
print(f"ef: {ef}")
print(f"OUTCAR: {fn_outcar}")
print(f"EIGENVAL: {fn_eig}")

if not ief:
    if path.exists(fn_outcar):
        for line in open(fn_outcar, 'r'):
            if line.find('E-fermi') > -1:
                ef = float(line.split()[2])
    else:
        print('OUTCAR does not exist!')
        input("Press ENTER to exit>> ")
        exit(0)

if not path.exists(fn_eig):
    print('EIGENVAL does not exist!') 
    input("Press ENTER to exit>> ")
    exit(0)

f_eig = open(fn_eig)
lines = f_eig.readlines()
f_eig.close()

nspin   = int(lines[0].split()[3])
nkpoint = int(lines[5].split()[1])
nband   = int(lines[5].split()[2])
print(f"nspin  : {nspin}")
print(f"nkpoint: {nkpoint}")
print(f"nband  : {nband}")
print("nlines:", len(lines))

eup   = ef + 99.0
edown = ef - 99.0
k = []
for i in range(nkpoint):
    iline = 7 + i * (nband + 2)
#    print(f"lines[{iline}]: {lines[iline]}", end = '')
    k.append(lines[iline].split()[0:3])
    for j in range(nband):
        for s in range(nspin):
            eval = float(lines[8 + i * (nband + 2) + j].split()[s + 1])
            d    = eval - ef 
            if (d > 0):
                if (eval < eup):
                    eup = eval
                    kup = i
                    bup = j
                    sup = s
            if (d < 0):
                if (eval > edown):
                    edown = eval
                    kdown = i
                    bdown = j
                    sdown = s

print("")
print(f'Efermi  : {ef:.3f}')
print(f'N band  : {nband}')
print(f'N kpoint: {nkpoint}')
print(f'N spin  : {nspin}')
print(f'Band gap: {eup-edown:.3f}')
print(f'Bottom of conduction band:')
print(f'  Energy : {eup:.3f}')
print(f'  K point: #{kup+1} k={k[kup]}')
print(f'  Band   : {bup + 1}')
print(f'  Spin   : {sup + 1}')
print(f'Top of valence band:')
print(f'  Energy : {edown:.3f}')
print(f'  K point: #{kdown + 1} k={k[kdown]}')
print(f'  Band   : {bdown + 1}')
print(f'  Spin   : {sdown + 1}')


app.terminate(pause = True)