"""

:mod:`Inelastica.BandStruct`
============================

Bandstructure and Fermi surface calculator

.. currentmodule:: Inelastica.BandStruct

"""
from __future__ import print_function

import Inelastica.Symmetry as SYM
import Inelastica.io.siesta as SIO
import Inelastica.MakeGeom as MG
import Inelastica.math as MM
import numpy as N
import scipy.linalg as SLA
import sys
import glob
import os
from optparse import OptionParser, OptionGroup
import Inelastica.io.xmgrace as XMGR

mm = MM.mm
outerAdd = MM.outerAdd
dist = MM.dist
mysqrt = MM.mysqrt

# Dummy classes for global variables
# general, HS, basis and geom are a global variables

########################################################
##################### Main routine #####################
########################################################


def main():
    #setupParameters()

    readxv()
    readHS()
    readbasis()
    for ispin in range(HS.nspin):
        calcBands(ispin)
    for ispin in range(HS.nspin):
        calcFS(ispin)

########################################################


def readxv():
    global geom
    # Read geometry from first .XV file found in dir
    fns = glob.glob('*.XV')

    if len(fns) > 1:
        print("ERROR: BandStruct: More than one .XV file ... which geometry to choose???")
        sys.exit(1)
    elif len(fns) < 1:
        print("ERROR: BandStruct: Error ... No .XV file found!")
        sys.exit(1)

    print(('Reading geometry from "%s" file' % fns[0]))
    geom = MG.Geom(fns[0])
    geom.sym = SYM.Symmetry(fns[0], onlyLatticeSym=True)
    if geom.sym.NNbasis != geom.natoms:
        print("ERROR: Siesta cell does not contain one unit cell")
        sys.exit(1)

########################################################


def readbasis():
    global basis
    fn = glob.glob('RUN.fdf')
    basis = SIO.BuildBasis(fn[0], 1, geom.sym.NN, HS.lasto)

########################################################


def readHS():
    # Setup H, S and self-energies
    global HS
    fn = glob.glob('*.TSHS')
    if len(fn) > 1:
        print("ERROR: BandStruct: More than one .TSHS file ... which to choose???")
        sys.exit(1)
    if len(fn) < 1:
        print("ERROR: BandStruct: No .TSHS file ???")
        sys.exit(1)
    HS = SIO.HS(fn=fn[0])


########################################################
def calcFS(ispin):
    # Calculate Fermi-surface
    NNk = 31
    bands = N.zeros((NNk, NNk, NNk, HS.N), N.float64)
    for ix in range(NNk):
        for iy in range(NNk):
            for iz in range(NNk):
                # "unitless" k-vect
                kpnt = N.array([ix, iy, iz], N.float64)/float(NNk-1)
                HS.setkpoint(kpnt, verbose=False)
                eival = SLA.eigh(HS.H[ispin], HS.S, eigvals_only=True)
                ipiv = N.argsort(eival)
                bands[ix, iy, iz, :] = eival[ipiv]
        SIO.printDone(ix, NNk, 'Fermi Surface: ')
    writeFS(ispin, NNk, bands)


def writeFS(ispin, NNk, bands):
    # Write output
    if HS.nspin > 1:
        sspin = ['.up', '.down'][ispin]
    else:
        sspin = ''

    # How many bands to include ?
    indx = []
    for ii in range(HS.N):
        if N.sum((bands[:, :, :, ii].reshape((-1,)) < general.eMax)*\
                 (bands[:, :, :, ii].reshape((-1,)) > general.eMin)):
            indx += [ii]

    f = open(general.DestDir+'/'+'FermiSurface'+sspin+'.BXSF', 'w')
    f.write("""
BEGIN_INFO
   # Band-XCRYSDEN-Structure-File
   # Generated by Inelastica.BandStruct
   # Launch as: xcrysden --bxsf example.bxsf
   #
   Fermi Energy: 0.0000
END_INFO
BEGIN_BLOCK_BANDGRID_3D
   What_is_this_text
   BEGIN_BANDGRID_3D_simple_example
     %i"""%len(indx))
    f.write("\n     %i %i %i\n"%(NNk, NNk, NNk))
    f.write("\n     0.0 0.0 0.0\n")
    rlat = [geom.sym.b1, geom.sym.b2, geom.sym.b3]
    f.write("\n     %f %f %f\n"%(rlat[0][0], rlat[0][1], rlat[0][2]))
    f.write("\n     %f %f %f\n"%(rlat[1][0], rlat[1][1], rlat[1][2]))
    f.write("\n     %f %f %f\n"%(rlat[2][0], rlat[2][1], rlat[2][2]))
    for ii in indx:
        f.write("\n   BAND:  %i\n"%(ii))
        for ix in range(NNk):
            for iy in range(NNk):
                for iz in range(NNk):
                    f.write('%f '%bands[ix, iy, iz, ii])
                f.write('\n')
            f.write('\n')
    f.write("""
   END_BANDGRID_3D
END_BLOCK_BANDGRID_3D
""")


########################################################
def calcBands(ispin):
    # Calculate bandstructure along high symmetry directions
    # what = [["text",k-direction, k-origin, Num k-points], ...]
    what = geom.sym.what()

    bands = []
    for ii in range(len(what)-1):
        #txt = what[ii][1]+'-'+what[ii+1][1]
        f, t = what[ii][0]/(2.0*N.pi), what[ii+1][0]/(2.0*N.pi)
        korig, kdir = f, t-f
        Nk = general.NNk

        ev = N.zeros((Nk, HS.N), N.float64)
        for ii in range(Nk):
            kpnt = korig + kdir*(ii/float(Nk-1))
            # Change to "unitless" k-vect
            kpnt2 = mm(N.array([geom.sym.a1, geom.sym.a2, geom.sym.a3]), kpnt)

            HS.setkpoint(kpnt2)
            eival = SLA.eigh(HS.H[ispin], HS.S, eigvals_only=True)
            ipiv = N.argsort(eival)
            ev[ii, :] = eival[ipiv]
        bands += [ev]

    writeBands(ispin, what, bands)


########################################################
def writeBands(ispin, what, bands):
    # Write output
    if HS.nspin > 1:
        sspin = ['.up', '.down'][ispin]
    else:
        sspin = ''

    Graphs = []
    for jj in range(len(what)-1):
        #for jj, elem in enumerate(what):
        txt = what[jj][1]+"-"+what[jj+1][1]
        Nk = general.NNk

        f = open(general.DestDir+'/'+txt+sspin+'.dat', 'w')
        xx = N.array(list(range(Nk)), N.float64)/(Nk-1.0)
        iColor, Datasets = 1, []
        for ii in range(len(bands[jj][0, :])):
            # Choose bands within +-5 eV from Ef
            if N.sum((bands[jj][:, ii] < general.eMax)*\
                     (bands[jj][:, ii] > general.eMin)) > 1e-5:
                f.write("\n# Band %i \n"%(ii))
                for kk, data in enumerate(bands[jj][:, ii]):
                    f.write("%i %e\n"%(kk, data))
                Datasets += [XMGR.XYset(xx, bands[jj][:, ii], Lcolor=iColor)]
                iColor += 1
        f.close()

        g = XMGR.Graph()
        for data in Datasets:
            g.AddDatasets(data)
        g.SetSubtitle(txt)

        g.SetXaxis(label='', majorUnit=0.5, minorUnit=0.1, vmax=1, vmin=0)
        if jj == 0:
            g.SetYaxis(label='E-Ef [eV]', majorUnit=1, minorUnit=0.2,
                       vmax=general.eMax, vmin=general.eMin)
        else:
            g.SetYaxis(label='', majorUnit=1e10, minorUnit=0.2,
                       vmax=general.eMax, vmin=general.eMin)
        Graphs += [g]

    p = XMGR.Plot(general.DestDir+'/BandStruct'+sspin+'.agr', Graphs[0])

    for ii in range(1, len(Graphs)):
        p.AddGraphs(Graphs[ii])
    p.ArrangeGraphs(nx=len(Graphs), ny=1, hspace=0.0, vspace=0.0)

    # Finally, write the plot file
    p.ShowTimestamp()
    p.WriteFile()
    p.Print2File(general.DestDir+'/BandStruct.eps')


###########################################################
def setupParameters():
    global general

    usage = "usage: %prog [options] DestinationDirectory"
    description = """
Calculation of Fermi surface.

For help use --help!
"""
    parser = OptionParser(usage, description=description)
    EC = OptionGroup(parser, "Options for Fermi surface")
    EC.add_option("-N", "--NumPoints", dest='NNk', default=20, type='int',
                  help="Number of k-points for Fermi-surface calculation [%default]")
    EC.add_option("-f", "--fdf", dest='fdfFile', default='./RUN.fdf', type='string',
                  help="fdf file used for transiesta calculation [%default]")
    EC.add_option("-a", "--min", dest='eMin', default=-5.0, type='float',
                  help="min of energy range to plot [%default]")
    EC.add_option("-b", "--max", dest='eMax', default=5.0, type='float',
                  help="max of energy range to plot [%default]")

    parser.add_option_group(EC)

    (general, args) = parser.parse_args()
    print(description)

    if not os.path.exists(general.fdfFile):
        parser.error("No input fdf file found, specify with --fdf=file.fdf (default RUN.fdf)")

    print(args)
    if len(args) != 1:
        parser.error('ERROR: You need to specify destination directory')
    general.DestDir = args[0]
    if not os.path.isdir(general.DestDir):
        print('\nBandStruct : Creating folder %s' %general.DestDir)
        os.mkdir(general.DestDir)
    else:
        parser.error('ERROR: destination directory %s already exist!'%general.DestDir)

    def myprint(arg, iofile):
        # Save in parameter file
        print(arg)
        iofile.write(arg+'\n')

    class myopen(object):
        # Double stdout to RUN.out and stdout

        def write(self, x):
            self.stdout.write(x)
            self.file.write(x)

    fo = myopen()
    fo.stdout, fo.file = sys.stdout, open(general.DestDir+'/RUN.out', 'w', 0)
    sys.stdout = fo

    iofile = open(general.DestDir+'/Parameters', 'w')
    argv = ""
    for ii in sys.argv:
        argv += " "+ii
    myprint(argv, iofile)
    myprint('##################################################################################', iofile)
    myprint('## Band structure ', iofile)
    myprint('fdfFile            : %s'%general.fdfFile, iofile)
    myprint('Number of k-points : %f'%general.NNk, iofile)
    myprint('##################################################################################', iofile)
    iofile.close()


##################### Start main routine #####################

if __name__ == '__main__':
    main()
