#! /usr/bin/env python

from __future__ import print_function
import sys,os,argparse

# 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(gauge='periodic',mode='positionoperator'):

    # 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)


    # now the position operator
    dists=bp.kdists
    kpts=bp.kpnts
    

    
    # the position operator is the basis Berry connection + the
    # diagonal matrix of Wannier centers (site vectors)
    sitevectors=s.wannierCenterMatrix()
    




    with bp.openBandFile(s.dirname+'/+Rx',s.nspin,len(kpts),\
                         progress='Rkplotcalc') as fbx:
        fby=bp.openBandFile(s.dirname+'/+Ry',s.nspin,len(kpts))
        fbz=bp.openBandFile(s.dirname+'/+Rz',s.nspin,len(kpts))
        # now fb is an instance of `BandFileContext`
        for ik,k in enumerate(kpts):
            for ms in range(s.nspin):

                (Hk,Rk)=s.hamAtKPoint(k*s.kscale,ms,makebasisconnection=True,
                                      gauge=gauge)
                if mode=='positionoperator':
                    Rk[0]*=-1 
                    Rk[1]*=-1 
                    Rk[2]*=-1 
                if gauge!='periodic' and mode=='positionoperator':
                    Rk[0]+=sitevectors[0]
                    Rk[1]+=sitevectors[1]
                    Rk[2]+=sitevectors[2]
                    
                (EV,C)=LA.eigh(Rk[0])
                fbx.write(ms,dists[ik],k,EV)
                (EV,C)=LA.eigh(Rk[1])
                fby.write(ms,dists[ik],k,EV)
                (EV,C)=LA.eigh(Rk[2])
                fbz.write(ms,dists[ik],k,EV)

        fby.close()
        fbz.close()



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


if __name__ == '__main__':


    parser=argparse.ArgumentParser(description='',
                                   conflict_handler='resolve',
                                   epilog='')
    parser.add_argument('--gauge', action='store',dest='gauge',
                        default='periodic')
    parser.add_argument('--mode', action='store',dest='mode',
                        default='positionoperator')


    args=parser.parse_args()
    
    if args.gauge=='relative': args.gauge='force'+args.gauge
    

    work(gauge=args.gauge,mode=args.mode)

