#! /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 spin polarized band structure
    s.calculateBandStructure(bp,suffix='');

    # now hand made with removeal of the polarization
    # using the xc-field
    dists=bp.kdists
    kpts=bp.kpnts
    

    gauge='relative'

    with bp.openBandFile(s.dirname+'/+band_nsp',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,B)=s.hamAtKPoint(k*s.kscale,ms,gauge=gauge,
                                     makesigma=True,makexcfield=True)
                # subtract S.B from H
                Hk-=(B[0].dot(S[0])+B[1].dot(S[1])+B[2].dot(S[2]))
                (EV,C)=LA.eigh(Hk)
                fb.write(ms,dists[ik],k,EV)



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


if __name__ == '__main__':

    work()

