#! /usr/bin/env python

from __future__ import print_function
import sys,os

# If your pyfplo is not found you could also
# explicitly specify the pyfplo version path:
#sys.path.insert(0,"/home/magru/FPLO/FPLO22.00-62/PYTHON/doc");

import numpy as np
import numpy.linalg as LA
import pyfplo.slabify as sla
import pyfplo.fploio as fploio
import pyfplo.common as com


print( '\npyfplo version=: {0}\nfrom: {1}\n'.format(sla.version,sla.__file__))
# protect against wrong version
#if fedit.version!='22.00': raise RuntimeError('pyfplo version is incorrect.')


# ===================================================================
# 
# ===================================================================

def work():

    # read bandplot points from =.in
    p=fploio.INParser()
    p.parseFile('../../=.in')
    d=p()('special_sympoints')
    l=[]
    for i in range(d.size()):
        l.append([ d[i]('label').S,d[i]('kpoint').listD])
    

    
    # prepare Slabify
    hamdata='../../+hamdata'
    
    s=sla.Slabify()
    s.object='3d'
    s.printStructureSettings()
    s.prepare(hamdata)
    
    # prepare BandPlot
    bp=sla.BandPlot()
    bp.points=l
    bp.calculateBandPlotMesh(s.dirname)

    # use Slabify routine to create the band structure
    s.calculateBandStructure(bp,suffix='');

    # now hand made
    dists=bp.kdists
    kpts=bp.kpnts
    
    # convert velocities from eV*aB to 10^6 m/s
    # 
    # first convert Bohr raddi into meters then divide by hbar in J*s
    # and then convert eV into J and divide by 10^6
    units=(com.c_abtoang*com.c_angstroem_m)\
        /com.c_hbar_Js*com.c_echarge_C/10**(6)
    

    # First just diagonalize the Hamiltonian using numpy and write band structure.
    with bp.openBandFile(s.dirname+'/+b',s.nspin,len(kpts),\
                         progress='bandplotcalc') as fb:
        # now fb is an instance of `BandFileContext`
        for ik,k in enumerate(kpts):
            for ms in range(s.nspin):
                Hk=s.hamAtKPoint(k*s.kscale,ms)
                (EV,C)=LA.eigh(Hk)
                fb.write(ms,dists[ik],k,EV)



    # now also create band weights containing the absolute values of
    # the Fermi velocities 
    with bp.openBandFile(s.dirname+'/+w',1,len(kpts),progress='bandplot',
                         weightlabels=['vx','vy','vz']) as fb:
        for dk,k in zip(dists,kpts):
            (Hk,dHk)=s.hamAtKPoint(k*s.kscale,0,makedhk=True)

            (EV,C)=LA.eigh(Hk) # diagonalize
            v=np.zeros(shape=(3,EV.shape[0]))
            for n in range(C.shape[1]):
                for i in range(3):
                    a=dHk[i].dot(C[:,n])
                    v[i,n]+=np.abs(np.vdot(C[:,n],a))*units
            
            fb.write(0,dk,k,EV,v) # write band file
                
    # some on the fly shell+xfbp scripting
    os.system('''
    xfbp -p - <<EOL 2>&1 >/dev/null &
kill all
read band "{0}/+b"
g1.gr1 line width 3
read band "{0}/+band_sf"
g1.gr2 line width 1
line color 0xfff00
autoscale
title "compare automatically and manually created bands"
EOL
    '''.format(s.dirname))

    os.system('''
    xfbp -p - <<EOL 2>&1 >/dev/null &
kill all
read bandweights "{0}/+w"
g1.gr1.w2 on
w3 on
autoscale
    title "|vx|, |vy| and |vz| in 10$^6$. m/s"
EOL
    '''.format(s.dirname))    



# ===================================================================
# 
# ===================================================================


if __name__ == '__main__':

    work()

