import numpy as np
import os
import glob
import fermiEnergy

class dos_files:

    def __init__(self, path_dos, prefix, atoms, orbitals = {'s':1,'p':3,'d':5,'f':7}):
        self.path_dos = path_dos
        self.prefix = prefix
        self.atoms = atoms
        self.orbitals = orbitals

        os.chdir(path_dos)

        self.atom_orbt = []
        self.files = []
        self.all_sumpdos = []

        for current_atom in self.atoms: #maps through the atoms
            for current_orbital in self.orbitals: #maps through the orbitals
                current_orbital_files = glob.glob(f'{prefix}*({current_atom})*({current_orbital})*')

                num_files = len(current_orbital_files) #number of files in current orbital

                self.atom_orbt.append([current_atom, current_orbital, num_files])
                self.files.append(current_orbital_files)
                self.all_sumpdos.append([])

                #print('-'*50)
                #print(current_orbital, current_orbital_files)

                if num_files: #if there are files in current orbital
                    file_content = np.genfromtxt(current_orbital_files[0])
                    lenthE = len(file_content)

                    self.energy = file_content[:, 0] #file column 1

                    sumpdos = np.zeros((orbitals[current_orbital], lenthE))

                    for current_file in current_orbital_files:

                        file_content = np.genfromtxt(current_file)

                        #ldos = file_content[:, 1] #column 2

                        pdos = [] #other columns
                        for pdos_column in range(2,self.orbitals[current_orbital]+2):
                            pdos.append(file_content[:, pdos_column])

                        pdos = np.array(pdos)
                        sumpdos += pdos

                    self.all_sumpdos[-1] = sumpdos
                else:
                    self.files[-1] = f'There are no files in {current_orbital} orbital'
                    self.all_sumpdos[-1] = f'There are no files in {current_orbital} orbital'

    def return_values(self):
        return self.files, self.all_sumpdos

    def pdos_files(self):
        column_names = {'s':{0:'s', 1:'total'}, 'p':{0:'pz', 1:'px', 2:'py', 3:'total'}, 'd':{0:'dz2', 1:'dzx', 2:'dzy', 3:'dx2-y2', 4:'dxy', 5:'total'}, 'f':{0:0, 1:1, 2:2, 3:3, 4:4, 5:5, 6:6, 7:'total'}}

        for current_atom_orbt in self.atom_orbt:
            current_atom_orbt_index = self.atom_orbt.index(current_atom_orbt)
            current_sumpdos = self.all_sumpdos[current_atom_orbt_index]
            name = f'{current_atom_orbt[0]}_{current_atom_orbt[1]}'
            num_files = current_atom_orbt[2]

            current_columns_names = column_names[current_atom_orbt[1]]

            if num_files:
                with open(f'{name}.dos','w') as wrt_file: #open writable file

                    wrt_file.write('#E           ')
                    for column in range(len(current_sumpdos)+1):
                        wrt_file.write(f'{current_columns_names[column]}           ')
                    wrt_file.write('\n')

                    for line in range(len(self.energy)): #maps through all values of energy
                        pdos_totalsum = 0
                        wrt_file.write(f'{self.energy[line]:.6f}     ')

                        for column in current_sumpdos: #maps trhough all columns of sumpdos
                            current_value = column[line]

                            wrt_file.write(f'{current_value:.6f}     ')
                            pdos_totalsum += current_value

                        wrt_file.write(f'{pdos_totalsum:.6f}')
                        wrt_file.write('\n')


    def pdos_graphs(self):
        pass




#test area
if __name__ == '__main__':
    dirs = glob.glob('S*')
    dirs = sorted(dirs)
    print(dirs)

    prefix = 'MnZrTe2'
    atoms = ['Zr','Te','Mn']
    orbitals = {'s':2,'p':4,'d':6}
    orbital_names = {'s':{0:'', 1:'_total'}, 'p':{0:'_pz', 1:'_px', 2:'_py', 3:'_total'}, 'd':{0:'_dz2', 1:'_dzx', 2:'_dzy', 3:'_dx2-y2', 4:'_dxy', 5:'_total'}, 'f':{0:0, 1:1, 2:2, 3:3, 4:4, 5:5, 6:6, 7:'total'}}


    with open('dos.dat','w') as arq:
        arq.write('#clusters           ')

        for currentAtom in atoms:
            for currentOrbital in orbitals:
                currentOrbitalName = orbital_names[currentOrbital]
                for i in range(orbitals[currentOrbital]):
                    arq.write(f'{currentAtom}_{currentOrbital}{currentOrbitalName[i]}          ')
        arq.write('total\n')

        for currentDir in dirs:
            os.chdir(currentDir)

            clusters = glob.glob('MnZrTe2*')
            clusters = sorted(clusters)
            for currentCluster in clusters:
                print(currentCluster)
                arq.write(f'{currentCluster}    ')

                os.chdir(currentCluster)
                print(os.getcwd())

                path = './'
                sistema = dos_files(path, prefix, atoms)
                #values = sistema.return_values()
                sistema.pdos_files() #save files

                Efermi = fermiEnergy.fermiEnergy(path, fileName='scf.out')
                print(Efermi)
                filesExtension='.dos'
                dosFiles, allDOS = fermiEnergy.calculateDosFermi(path, Efermi, filesExtension=filesExtension)

                print(dosFiles)
                print(allDOS)
                soma = 0
                for currentAtom in atoms:
                    for currentOrbital in orbitals:
                        try:
                            currentFile = f'{currentAtom}_{currentOrbital}{filesExtension}'
                            fileIndex = dosFiles.index(currentFile)
                            dosCurrentFile = allDOS[fileIndex]
                            for i in range(orbitals[currentOrbital]):
                                if i!=orbitals[currentOrbital]-1:
                                    soma += dosCurrentFile[i]
                                arq.write(f'{dosCurrentFile[i]:.5f}   ')
                        except:
                            for i in range(orbitals[currentOrbital]):
                                arq.write('0.00000  ')
                        arq.write(f'{soma}')

                print(os.getcwd(), 'done')
                os.chdir('..')

                arq.write('\n')
            os.chdir('..')