class ALD:  
    
    def __init__(self, 
                 outdir = './'
                 ):

        '''Give path to output folder, all the files should be as same as files generated by ALD code, 
        all the returning lists will be in the order of harmonic_properties.dat, unless it is mentioned, like in case of cummulative lists'''
        
        from ase import Atoms
        from ase.io.trajectory import Trajectory
        import numpy as np
        import re
        
        self.outdir = outdir
        self.primcell = self.load_data(f'{self.outdir}/cell_primitive.dat',head=1)
        self.harmprop = self.load_data(f"{self.outdir}/harmonic_properties.dat", head=1)
        self.natoms = len(self.get_scaled_positions())
        self.nqpts = int((len(self.harmprop))/int(3*self.natoms))
        
        f = open(f'{self.outdir}/parameters.dat', 'r')
        lines = f.readlines()
        f.close()

        for i in lines:
            if 'nlength' in i:
                nlength = float(i.split(':')[-1])*1e10
                break
        self.nlength = nlength
    
    def load_data(self, filename, head=0):
        import numpy as np
        f = open(filename, 'r')
        lines = [i.split() for i in f.readlines()[head:]]
        f.close()

        for n, i in enumerate(lines):
            for o, j in enumerate(i):
                try:
                    lines[n][o] = float(j)
                except:
                    lines[n][o] = j

        length = max(map(len, lines))
        y=np.array([xi+[None]*(length-len(xi)) for xi in lines])

        return y

    
    def flatten(self, list_of_lists):
            
        if len(list_of_lists) == 0:
            return list_of_lists
        if isinstance(list_of_lists[0], list):
            return self.flatten(list_of_lists[0]) + self.flatten(list_of_lists[1:])
        return list_of_lists[:1] + self.flatten(list_of_lists[1:])

    def get_cell(self):
        '''
        Returns cell parameters as per primitive.dat '''
        return self.primcell[0:3, 0:3]*self.nlength

    def get_direct_qpts(self):

        '''
        Returns directs qpts from harmonic_properties.dat, in the same order as in file '''
        import numpy as np
        qpts = []
        for i in self.harmprop:
            if i[0] != len(qpts)-1 or len(qpts) == 0:
                qpts.append(i[2:5])

        return np.array(qpts, dtype=float)
    
    def get_cartesian_qpts(self):
        '''
        Returns crtesian qpts from harmonic_properties.dat, in the same order as in file '''
        import numpy as np
        qpts = []
        for i in self.harmprop:
            if i[0] != len(qpts)-1 or len(qpts) == 0:
                qpts.append(i[5:8])

        return np.array(qpts, dtype=float)

    def get_elements(self):
        '''
        Returns list of elements from primitive.dat'''
        return self.primcell[3:,1]

    def get_scaled_positions(self):
        return self.primcell[3:, 2:5]
    
    def get_cartesian_positions(self):
        return self.primcell[3:, 5:8]
    
    def get_eigvecs(self):
        
        import re
        import numpy as np
        
        eigvecs = []
        eigvecdata1 = self.load_data(f'{self.outdir}/eigenVector.dat', head=1)        
        eigvecdata2 = np.array([self.flatten([re.findall(r"[-+]?\d*\.\d+|\d+", str(s)) for s in s1]) for s1 in eigvecdata1], dtype=float)
        
        real = np.array([[[0]*self.natoms*3]*self.natoms*3]*self.nqpts, dtype=float)
        imaginary = np.array([[[0]*self.natoms*3]*self.natoms*3]*self.nqpts, dtype=float)
        
        for i in eigvecdata2:
            qid = int(float(i[0]))
            modeid = int(float(i[1]))
            r = [i[j] for j in range(3, len(i)) if j%2 == 1]
            im = [i[j] for j in range(3, len(i)) if j%2 == 0]
            real[qid][modeid] = r
            imaginary[qid][modeid] = im    
        
        return real, imaginary
    
    def get_frequencies(self):
        return self.harmprop[:,8]
    
    def get_heat_capacities(self):
        sps = self.harmprop[:,9]
        import numpy as np
        sps[np.isnan(sps)]=0
        return sps
    
    def get_group_velocities(self):
        import numpy as np
        vgs = self.harmprop[:,10:13]
        vgs[np.isnan(vgs)] = 0
        return vgs
    
    def get_gruneissen_parameters(self):
        return self.harmprop[:,14:]
    
    def get_mean_free_paths(self):
        import numpy as np
        return np.array([np.sqrt(np.sum(i**2)) for i in self.get_group_velocities()])*self.get_life_times()
                             
    def get_life_times(self, index=-1):
        import numpy as np
        rta = self.load_data(f'{self.outdir}/phonon_RTA.dat', head=1)
        rta[np.isinf(rta)]=0
        return rta[:,index]   
    
    def get_band_structure(self):
        
        import numpy as np
        data = self.load_data(f'{self.outdir}/directional_harmonic_properties.dat', head=1)
        distances = np.array([[0]*int(len(data)/3/self.natoms)]*3*self.natoms, dtype=float)
        bands =np.array([[0]*int(len(data)/3/self.natoms)]*3*self.natoms, dtype=float)
        
        for i in data:
            bands[int(float(i[1]))][int(float(i[0]))] = i[9]
            distances[int(float(i[1]))][int(float(i[0]))] = i[8]
        return distances, bands
    
    def plot_bands(self):
        import matplotlib.pyplot as plt
        fig ,ax = plt.subplots()
        dists, bands = self.get_band_structure()
        for i in range(len(bands)):
            ax.plot(dists[i], bands[i], 'k-') 
            ax.set_xlim([0,max(dists[0])]) 
            ax.set_ylim([0, max(bands[-1])+1])
        return fig, ax  
    
    def write_modes(self, qid=0, mode=0, nframes=100, repeat=(1,1,1), factor=1):

        '''
        Writes *.traj file for given phonon mode which can be visualized with ase-gui'''
        from ase import Atoms
        from ase.io.trajectory import Trajectory
        import numpy as np
        import re
        
        
        eigvecs = np.array(self.flatten([re.findall(r"[-+]?\d*\.\d+|\d+", str(s)) for s in self.eigdata[(qid)*3*self.natoms+mode]]), dtype=float)
        r = [eigvecs[j] for j in range(3, len(eigvecs)) if j%2 == 1]
        im = [eigvecs[j] for j in range(3, len(eigvecs)) if j%2 == 0]
        cmplxeigvecs = np.array([0]*3*self.natoms, dtype=complex)
        
        for i in range(len(r)):
            cmplxeigvecs[i] = np.complex(r[i], im[i])        
        cmplxeigvecs = cmplxeigvecs.reshape(self.natoms,-1)
        
        cell = self.get_cell()
        elements = self.get_elements()
        positions = self.get_cartesian_positions()
        
        atoms = Atoms(symbols=elements, positions=positions, cell=cell, pbc=True)
        traj = Trajectory(f'{self.outdir}/q{qid}_m{mode}.traj', 'w')
        
        for x in np.linspace(0, 2*np.pi, nframes, endpoint=False):
            disps = ((np.exp(1.j * x)*cmplxeigvecs).real)*factor
            new_pos = atoms.get_positions() + disps
            atoms.set_positions(new_pos)
            atoms *= repeat
            traj.write(atoms)
        atoms = Atoms(symbols=elements, positions=positions, cell=cell, pbc=True)
        traj.close()
        
        return None
        
    def get_pdf(self, x, w, bandwidth=0.05, xmin=0, xmax=None, npts=1000):
        
        from scipy.stats import gaussian_kde
        import numpy as np

        gk = gaussian_kde(x, weights=w)
        gk.covariance_factor = lambda : bandwidth
        gk._compute_covariance()

        if xmax == None:
            xmax = np.max(x)

        f = np.linspace(xmin, xmax, npts)
        g = gk(f)
        g[-1] = 0

        return f, gk(f)
    
    def get_pdos(self, bandwidth=0.05, xmin=0, xmax=None, npts=1000):
        
        import numpy as np
    
        natoms = self.natoms
        freqs = self.get_frequencies()
        spdos = [[] for i in range(natoms)]      

        pdos = self.load_data(f'{self.outdir}/pdos.dat', head=1)[:,1:natoms+1].T

        for i in range(len(spdos)):
            sfreqs, spdos[i] = self.get_pdf(freqs, pdos[i], xmin=xmin, xmax=xmax, npts=npts, bandwidth=bandwidth) 
            
        tdos = []
        for i in range(len(spdos[0])):
            a = 0
            for j in spdos:
                a += j[i]
            tdos.append(a)

        return np.array(sfreqs) , np.array(spdos), np.array(tdos)

    def get_scattering_phase_space3(self):

        phase3 = self.load_data(f"{self.outdir}/scattering_phase_space_3.dat", head=0)
        return phase3[:,2]

    def get_scattering_phase_space4(self):
        
        phase4 = self.load_data(f"{self.outdir}/scattering_phase_space_4.dat", head=0)
        return phase4[:,2]

    def get_kappas(self, indices=[1,1]):
        
        vels = self.get_group_velocities()
        times = self.get_life_times()
        heats = self.get_heat_capacities()

        kappas = heats*times*vels[:,(indices[0]-1)]*vels[:,(indices[1]-1)]
        return kappas


    def get_cummulative_conductivity_with_frequency(self, indices=[1,1]):
        
        import numpy as np

        freqs = self.get_frequencies()
        kk = self.get_kappas()
        sorted_freqs = freqs[freqs.argsort()]
        sorted_kk = kk[freqs.argsort()]
        cum_k = [0]
        for i in sorted_kk:
            cum_k.append(cum_k[-1]+i)

        return sorted_freqs, np.array(cum_k[1:])
    
    def get_cummulative_conductivity_with_mean_free_paths(self, indices=[1,1]):

        import numpy as np

        mpaths = self.get_mean_free_paths()
        kk = self.get_kappas()
        sorted_mpaths = mpaths[mpaths.argsort()]
        sorted_kk = kk[mpaths.argsort()]
        cum_k = [0]
        for i in sorted_kk:
            cum_k.append(cum_k[-1]+i)

        return sorted_mpaths, np.array(cum_k[1:])

    def get_cummulative_conductivity_with_life_times(self, indices=[1,1]):

        import numpy as np

        ltimes = self.get_life_times()
        kk = self.get_kappas()
        sorted_ltimes = ltimes[ltimes.argsort()]
        sorted_kk = kk[ltimes.argsort()]
        cum_k = [0]
        for i in sorted_kk:
            cum_k.append(cum_k[-1]+i)

        return sorted_ltimes, np.array(cum_k[1:])

    def get_force_constants_3(self):
        fcs = self.load_data(f"{self.outdir}/red_3.dat", head=1)
        return fcs[:,-1]

    def get_force_constants_4(self):
        fcs = self.load_data(f"{self.outdir}/red_4.dat", head=1)
        return fcs[:,-1]

    def get_thermal_conductivity(self, indices=[1,1]):
        import numpy as np
        kappas = self.get_kappas(indices=indices)
        return np.sum(kappas)

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

def load_data(filename, head=0):
        import numpy as np
        f = open(filename, 'r')
        lines = [i.split() for i in f.readlines()[head:]]
        f.close()

        for n, i in enumerate(lines):
            for o, j in enumerate(i):
                try:
                    lines[n][o] = float(j)
                except:
                    lines[n][o] = j

        length = max(map(len, lines))
        y=np.array([xi+[None]*(length-len(xi)) for xi in lines])

        return y

def get_pdf(x, w, bandwidth=0.05, xmin=0, xmax=None, npts=1000):
        
        from scipy.stats import gaussian_kde
        import numpy as np

        gk = gaussian_kde(x, weights=w)
        gk.covariance_factor = lambda : bandwidth
        gk._compute_covariance()

        if xmax == None:
            xmax = np.max(x)

        f = np.linspace(xmin, xmax, npts)
        g = gk(f)
        g[-1] = 0

        return f, gk(f)



        



            
        
