#! /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():
    printKresolvedChernNumber(homo=10,
                              mirrorindex=36, # m(z)=m(001)
                              nx=100,ny=100,
                              xaxis=[1.,0.,0.],
                              yaxis=[0.,1.,0.],
                              # In units of s.kscale.
                              # Since it is fcc the interval for one unit cell
                              # is twice as long: [-1,1] \equiv [-pi,pi]
                              xinterval=[-1,1],
                              yinterval=[-1,1])

    sq2=np.sqrt(2.)
    printKresolvedChernNumber(homo=10,
                              mirrorindex=45, # m(xy)=m(110)
                              nx=int(100*sq2),ny=100,
                              xaxis=[-1.,1.,0.],
                              yaxis=[ 0.,0.,1.],
                              # In units of s.kscale.
                              # Since it is fcc the interval for one unit cell
                              # is twice as long: [-1,1] \equiv [-pi,pi]
                              xinterval=[0,sq2*2],
                              yinterval=[-1,1])
    return
# ===================================================================
# 
# ===================================================================

def printKresolvedChernNumber(homo,mirrorindex,nx,ny,xaxis,yaxis,
                              xinterval,yinterval):

    
    # prepare Slabify
    hamdata='../../+hamdata'
    
    s=sla.Slabify()
    s.object='3d'
    s.printStructureSettings()
    s.prepare(hamdata)
    
    # put here the mirror operation indices which are wanted
    opindices=[mirrorindex] 


    box=sla.BoxMesh()
    # put the mirror plane vectors as printed by calculateMirrorChernNumbers
    xaxis=np.array(xaxis,dtype=float)
    yaxis=np.array(yaxis,dtype=float)
    
    zaxis=np.cross(xaxis,yaxis)
    zaxis/=LA.norm(zaxis)


    
    box.setBox(xaxis=xaxis,yaxis=yaxis,origin=[0,0,0])

    box.setMesh(nx,ny,1,xinterval,yinterval,[0,0])

    mesh=box.mesh(s.kscale,True)

    
    for ms in range(s.nspin):

        fname='f_mirror{}_homo{}.spin{}'.format(mirrorindex,homo,ms+1)
        with open(fname,'w') as f:

            f.write('#xfbp file type: xynz\n')
            for kx in box.xmesh():
                f.write('{}\n'.format(kx*np.pi))
            f.write('\n')
            for ky in box.ymesh():
                f.write('{}\n'.format(ky*np.pi))
            f.write('\n')

            x=3
            sub=int(len(mesh)/int(100./x))
            for ik,k in enumerate(mesh):
                #print( 'k: ',(100.*ik)/len(mesh),k,' spin: ',ms)
                if (ik % sub)==0 or ik==0:
                    print( 'progress: ',(100.*ik)/len(mesh),"%")

                # get Hamiltonian and symmetry reps
                (Hk,dhk,WF)=s.hamAtKPoint(k,ms,opindices=opindices,
                                          makedhk=True,makewfsymops=True)

                (EV,C)=LA.eigh(Hk)

                w=WF[0]
                if ms==0 and ik==0: print( w)
                if w.timerev: raise RuntimeError('mirror is combined with time reversal')

                # make CZ diagonalizes H and U=C^{+} Dk C
                (EU,CZ)=s.coDiagonalize(EV,C,w.Dk,check=True)

                # make an array which has element=1 if imaginary part
                # of mirror eigenvalue is positive else element=0
                subspace=list(map(lambda x: 1 if x>0 else 0,EU.imag))
                # get subspace Berry curvature using CZ, which diagonalizes
                # Hk and U=C^{+} Dk C
                
                (Fp,notused)=s.berryCurvature(EV,CZ,dhk,subspace)

                # the other subspace
                subspace=list(map(lambda x: 1 if x<0 else 0,EU.imag))
                (Fm,notused)=s.berryCurvature(EV,CZ,dhk,subspace)

                ff=0
                for i in range(homo):
                    ff+=-(Fp[:,i]-Fm[:,i]).dot(zaxis)*0.5

                f.write('{}\n'.format(ff))
                    


        txt='''
# -- mode: PYTHON --
killall()
g=G[1]
g.read('xynz','{1}')
g.autoscale()

g.title.text='homo {0}'
g.xaxislabel.text='kx [-$~p,$~p]'
g.yaxislabel.text='ky [-$~p,$~p]'
g.legend.off()

arrange(paperwidth=442,Nx=1,leftgap=0.18,rightgap=0.12,
    hgap=0.1,Ny=1,topgap=0.166,bottomgap=0.166,vgap=0.1,
    aspectratio=1, commonxaxis=True,commonyaxis=True,
    commontitle=True)
'''.format(homo,fname)
        with open(fname+'.xpy','w') as f:
            f.write(txt)
        os.system('xfbp {}'.format(fname+'.xpy'))
    print( 'done')
    return


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


if __name__ == '__main__':

    work()

