try:
    from gpaw.lcao.tightbinding import TightBinding
    from gpaw import restart, GPAW
    from gpaw.lcao.tools import get_lead_lcao_hamiltonian, get_lcao_hamiltonian, get_bf_centers, get_bfi
    _has_gpaw=True
except:
    _has_gpaw=False    

#from banddownfolder.scdm.downfolder import BandDownfolder
import numpy as np
from scipy.linalg import eigh
from ase.dft.kpoints import monkhorst_pack
from ase.units import Ha
import pickle

class GPAWWrapper():
    """
    https://gitlab.com/gpaw/gpaw/-/blob/master/gpaw/lcao/tightbinding.py
    """
    def __init__(self, atoms=None, calc=None, gpw_fname=None):
        if not _has_gpaw:
            raise ImportError("GPAW was not imported. Please install gpaw.")
        if atoms is None or calc is None:
            atoms, calc=restart(gpw_fname, fixdensity=False,  txt='nscf.txt')
        self.atoms=atoms
        self.calc=calc
        #self.wfs=self.calc.wfs
        #self.h=self.calc.hamiltonian
        #self.kpt_u = self.wfs.kpt_u
        self.cell = self.atoms.get_cell()
        #self.wann_centers=get_bf_centers(self.atoms)
        self.positions=self.cell.scaled_positions(get_bf_centers(self.atoms))
        #if calc.get_spin_polarized():
        #    self.positions=np.vstack([self.positions, self.positions])
        self.efermi=self.calc.get_fermi_level()
        self.R2kfactor= -2.0j* np.pi
        #self.nbasis=len(self.positions)
        self.norb=len(self.positions)
        self.nbasis=self.norb*2
        self._name='GPAW'


    #def gen_ham(self, k):
    #    H_MM=self.wfs.eigensolver.calculate_hamiltonian_matrix(self.h, self.wfs, k)
    #    tri2full(H_MM*Ha)
    #    return H_MM

    #def Sk(self, k):
    #    S_MM=self.wfs.S_qMM[kpt.q]
    #    tri2full(S_MM)
    #    return S_MM


    #def solve(self, k):
    #    return eigh(self.hamk(k), self.S(k))

    def get_kpts(self):
        return calc.get_ibz_k_points()

    def HS_and_eigen(self, kpts=None, convention=2):
        if kpts is not None:
            self.calc.set(kpts=kpts)
            self.atoms.calc=self.calc
            self.atoms.get_total_energy()

        self.calc.set(symmetry='off', fixdensity=True)
        self.atoms.get_total_energy()
        #self.kpt_u = self.wfs.kpt_u
        H, S= get_lcao_hamiltonian(self.calc)
        np.save('kpts.npy', kpts)
        np.save('H.npy', H)
        np.save('S.npy', S)

        nspin, nkpt, norb, _ = H.shape
        nbasis=nspin*norb
        evals=np.zeros((nkpt, nbasis), dtype=float)
        evecs=np.zeros((nkpt, nbasis, nbasis),dtype=complex)
        if self.calc.get_spin_polarized():
            H2=np.zeros((nkpt, nbasis, nbasis),dtype=complex)
            # spin up
            H2[:, :norb, :norb]=H[0]
            # spin down
            H2[:, norb:, norb:]=H[1]
            S2=np.zeros((nkpt, nbasis, nbasis),dtype=complex)
            S2[:, :norb, :norb]=S
            S2[:, norb:, norb:]=S

            for ikpt, k in enumerate(self.calc.get_ibz_k_points()):
                evals0, evecs0 = eigh(H[0, ikpt,:,:], S[ikpt,:,:])
                evals1, evecs1 = eigh(H[1, ikpt,:,:], S[ikpt,:,:])
                evals[ikpt, :norb]=evals0
                evals[ikpt, norb:]=evals1
                evecs[ikpt, :norb, :norb]=evecs0
                evecs[ikpt, norb:, norb:]=evecs1
        else:
            H2=H[0]
            for ikpt, k in self.kpt_u:
                evals[ikpt], evecs[ikpt]= eigh(H[0, ikpt], S[ikpt])
        np.save('kpts.npy', kpts)
        np.save('H2.npy', H2)
        np.save('S2.npy', S2)
        np.save('evals.npy', evals)
        np.save('evecs.npy', evecs)
        return H2, S2, evals, evecs

    def solve_all(self, kpts, convention=2):
        H, evals, evecs=self.H_and_eigen(kpts=kpts, convention=convention)
        return evals, evecs



class GPAWTBWrapper():
    def __init__(self, calc=None, atoms=None,  gpw_fname=None, pickle_fname=None):
        if not _has_gpaw:
            raise ImportError("GPAW was not imported. Please install gpaw.")
        if pickle_fname is None:
            if gpw_fname is not None:
                calc=GPAW(gpw_fname, fixdensity=True, symmetry='off', txt='TB2J_wrapper.log', basis='dzp', mode='lcao')
                atoms=calc.atoms
                atoms.calc=calc
                E = atoms.get_potential_energy()
            tb=TightBinding(atoms, calc)
            self.H_NMM, self.S_NMM=tb.h_and_s()
            self.Rlist=tb.R_cN.T
        else:
            with open(fname, 'rb') as myfile:
                self.H_NMM, self.S_NMM, self.Rlist=pickle.load(myfile)
        self.nR, self.nbasis, _ = self.H_NMM.shape
        self.positions=np.zeros((self.nbasis, 3))

    def save_pickle(self, fname):
        with open(fname, 'wb') as myfile:
            pickle.dump([self.H_NMM, self.S_NMM, self.Rlist], myfile)

    def solve(self, k):
        self.Hk=np.zeros((self.nbasis, self.nbasis), dtype='complex')
        self.Sk=np.zeros((self.nbasis, self.nbasis), dtype='complex')
        for iR, R in enumerate(self.Rlist):
            phase=np.exp(-2j*np.pi* np.dot(k, R))
            self.Hk+=self.H_NMM[iR]*phase
            self.Sk+=self.S_NMM[iR]*phase
        self.Sk=np.real((self.Sk+self.Sk.T)/2)
        return eigh(self.Hk, self.Sk)
        try:
            return eigh(self.Hk, self.Sk)
        except:
            return np.zeros(self.nbasis), np.zeros((self.nbasis, self.nbasis))

    def solve_all(self, kpts):
        evals = []
        evecs = []
        for ik, k in enumerate(kpts):
            evalue, evec = self.solve(k)
            evals.append(evalue)
            evecs.append(evec)
        return np.array(evals, dtype=float), np.array(evecs,
                                                      dtype=complex,
                                                      order='C')



def test_Ham():
    calc=GPAW('/home/hexu/projects/gpaw/bccFe/atoms_nscf.gpw', fixdensity=True, symmetry='off',  txt='nscf.txt')
    atoms=calc.atoms
    atoms.calc=calc
    calc.set(kpts=(3,3,3))
    E = atoms.get_potential_energy()
    tb=TightBinding(atoms, calc)
    kpath=monkhorst_pack([3,3,3])
    evals, evecs=tb.band_structure(kpath, blochstates=True)


def test():
    g=GPAWWrapper(fname='STO_nscf.gpw')
    g.H_and_eigen(kpts=monkhorst_pack(2,2,2))
    #g.save_pickle('model.pickle')
    #df=GPAWDownfolder('model.pickle')
    #df.downfold(nwann=1, anchor_kpt=[0,0,0] )
    #df.plot_band_fitting(npoints=100)

#test_Ham()

def test():
    g=GPAWWrapper(gpw_fname='/home/hexu/projects/gpaw/bccFe/atoms_nscf.gpw')
    H, E, V=g.H_and_eigen(kpts=monkhorst_pack([2,2,2]))
    print(H.shape)
    print(E.shape)
    print(E)
    print(V.shape)


