import os
import itertools
import numpy as np
from ufss.HLG import DiagonalizeHamiltonian, LiouvillianConstructor, DiagonalizeLiouvillian
import yaml

class manualH:
    def __init__(self,H,mu_ket_up,*,savedir='',RWA = True,for_RKE=False,
                 flourescence_yield=np.array([])):
        """Takes as input a Hamiltonian (H), and dipole operator (mu_ket_up) 
            and a dissipation super-operator D
"""
        self.H = self.convert_H_to_dict(H)
        self.mu = self.convert_mu_to_dict(mu_ket_up)
        if flourescence_yield.size == 0:
            fy_flag = False
        else:
            fy_flag = True
            self.fy = self.convert_fy_to_dict(flourescence_yield)
        self.check_mu_shape()
        if RWA:
            self.check_mu()

        self.for_RKE = for_RKE
        self.base_path = savedir
        self.save_path = os.path.join(self.base_path,'closed')
        os.makedirs(self.save_path,exist_ok=True)

        self.save_H()
        self.save_mu()
        self.make_params_file()
        if fy_flag:
            self.save_fy()

        self.run_closed()

    def make_params_file(self):
        filename = os.path.join(self.base_path,'params.yaml')
        params = {'ManualH':True}
        with open(filename,'w+') as yamlstream:
            yamlstream.write('### This file was generated by UFSS, do not modify unless you know what you are doing \n')
            yaml.dump(params,yamlstream)

    def convert_fy_to_dict(self,fy):
        if isinstance(fy,np.ndarray):
            fy_dict = {'all_manifolds':fy}
        elif isinstance(H,dict):
            fy_dict = fy
        else:
            raise Exception('Input H is not in a recognized format')

        return fy_dict

    def convert_H_to_dict(self,H):
        if isinstance(H,np.ndarray):
            Hdict = {'all_manifolds':H}
        elif isinstance(H,dict):
            Hdict = H
        else:
            raise Exception('Input H is not in a recognized format')

        return Hdict

    def convert_mu_to_dict(self,mu):
        if isinstance(mu,np.ndarray):
            mu_dict = {'up':mu}
        elif isinstance(mu,dict):
            mu_dict = mu
        else:
            raise Exception('Input mu is not in a recognized format')

        return mu_dict

    def check_mu_shape(self):
        for key in self.mu.keys():
            shape = self.mu[key].shape
            if len(shape) == 3:
                pass
            else:
                raise Exception("""Each component of mu must be a 3D ndarray,
                    where the third axis is the molecular frame cartesian 
                    axis""")

            if shape[-1] == 3:
                pass
            else:
                raise Exception("""Each component of mu must be a 3D ndarray,
                    where the third axis is the molecular frame cartesian 
                    axis, and must have 3 entries for x, y, z components""")

    def save_mu(self):
        np.savez(os.path.join(self.save_path,'mu_original_H_basis.npz'),
                 **self.mu)

    def save_H(self):
        np.savez(os.path.join(self.save_path,'H.npz'),**self.H)

    def save_fy(self):
        np.savez(os.path.join(self.save_path,'fluorescence_yield.npz'),**self.fy)
            
    def check_mu(self):
        if 'all_manifolds' not in self.H.keys():
            check = True
        else:
            for key in self.mu.keys():
                for i in range(3):
                    mu_i = self.mu[key][...,i]
                    check = np.allclose(mu_i,np.tril(mu_i))
                    if check:
                        pass
                    else:
                        break
        if not check:
            raise Exception("""dipole operator must be lower triangular 
                in order to ensure that the rotating wave approximation 
                is followed. If you know what you are doing, you can turn
                off this error by calling manual_H_and_D with RWA = 
                False""")

    def run_closed(self):
        if self.for_RKE:
            pass
        else:
            dh = DiagonalizeHamiltonian(self.base_path)

class manualD:
    def __init__(self,D,*,savedir='',H_eigenbasis=True,for_RKE=False,
                 conserve_memory=True):
        """Takes as input a Hamiltonian (H), and dipole operator (mu_ket_up) 
            and a dissipation super-operator D
"""
        self.base_path = savedir
        self.load_path = os.path.join(self.base_path,'closed')
        self.save_path = os.path.join(self.base_path,'open')
        os.makedirs(self.save_path,exist_ok=True)

        self.load_H(H_eigenbasis)

        self.D = self.convert_D_to_dict(D)
        self.D_operators_to_superoperator()
        self.flatten_all_Ds()

        self.set_L()
        self.save_L()
        self.modify_params_file(H_eigenbasis)
        if not conserve_memory:
            self.load_mu(H_eigenbasis)

            self.set_mu_Liouville_space()
            self.save_mu()
        if not for_RKE:
            dl = DiagonalizeLiouvillian(self.base_path,
                                        conserve_memory=conserve_memory)
            
    def modify_params_file(self,H_eigenbasis):
        filename = os.path.join(self.base_path,'params.yaml')
        params = {'ManualD':True,
                  'H_eigenbasis':H_eigenbasis}
        with open(filename,'a') as yamlstream:
            yaml.dump(params,yamlstream)

    def convert_D_to_dict(self,D):
        if isinstance(D,np.ndarray):
            Ddict = {'all_manifolds':D}
        elif isinstance(D,list):
            Ddict = {'all_manifolds':D}
        elif isinstance(D,dict):
            Ddict = D
        else:
            raise Exception('Input H is not in a recognized format')

        Dkeys = Ddict.keys()

        if 'all_manifolds' in Dkeys and len(Dkeys)==1:
            pass
        elif '00' in Dkeys:
            pass
        else:
            raise Exception('D dictionary keys are not recognized')

        return Ddict

    def D_operators_to_superoperator(self):
        for key in self.D.keys():
            if isinstance(self.D[key],np.ndarray):
                pass
            elif isinstance(self.D[key],list):
                LC = LiouvillianConstructor
                self.D[key] = LC.make_Liouvillian(self.D[key])
            else:
                raise Exception('D is not in a supported format')

    def flatten_all_Ds(self):
        for key in self.D.keys():
            if len(self.D[key].shape) == 4:
                self.D[key] = self.flatten_D(self.D[key])
            elif len(self.D[key].shape) == 2:
                pass
            else:
                raise Exception('D must be either a rank 2 or rank 4 tensor')
                

    def flatten_D(self,D):
        new_shape = (D.shape[0] * D.shape[1], D.shape[2] * D.shape[3])
        new_D = np.zeros(new_shape)

        for k,l in itertools.product(range(D.shape[2]),range(D.shape[3])):
            n = k*D.shape[3] + l
            new_rho = D[:,:,k,l]
            new_D[:,n] += new_rho.flatten()

        return new_D

    def load_H(self,H_eigenbasis):
        self.H = {}
        if H_eigenbasis:
            eigvals_path = os.path.join(self.load_path,'eigenvalues.npz')
            with np.load(eigvals_path) as arch:
                for key in arch.keys():
                    self.H[key] = np.diag(arch[key])
        else:
            H_path = os.path.join(self.load_path,'H.npz')
            with np.load(H_path) as arch:
                for key in arch.keys():
                    self.H[key] = arch[key]

        self.manifolds = list(self.H.keys())
                    
    def load_mu(self,H_eigenbasis):
        self.mu = {}
        if H_eigenbasis:
            mu_path = os.path.join(self.load_path,'mu.npz')
        else:
            mu_path = os.path.join(self.load_path,'mu_original_H_basis.npz')

        with np.load(mu_path) as arch:
            for key in arch.keys():
                self.mu[key] = arch[key]

    def make_U(self,ket_manifold,bra_manifold):
        H_ket = self.H[ket_manifold]
        H_bra = self.H[bra_manifold]
        LC = LiouvillianConstructor
        U_ins = LC.make_commutator_instructions2(-1j*H_ket,-1j*H_bra)
        U = LC.make_Liouvillian(U_ins,sparse=True)
        return U

    def set_L(self):
        self.L = {}
        if self.manifolds[0] == 'all_manifolds':
            L_key = 'all_manifolds'
            U = self.make_U(L_key,L_key)
            self.L[L_key] = U + self.D[L_key]
        else:
            for ket_man,bra_man in itertools.product(self.manifolds,
                                                     self.manifolds):
                L_key = ket_man + bra_man
                U = self.make_U(ket_man,bra_man)
                self.L[L_key] = U + self.D[L_key]

    def save_L(self):
        np.savez(os.path.join(self.save_path,'L.npz'),**self.L)

    def mu_key_to_manifold_keys(self,key):
        if self.manifolds[0] == 'all_manifolds':
            starting_key = 'all_manifolds'
            ending_key = 'all_manifolds'
        else:
            starting_key, ending_key = key.split('_to_')
        return starting_key, ending_key


    def make_mu_by_manifold_ket(self,old_manifold,change):
        i,j = old_manifold
        i2 = i + change
        j2 = j
        
        if i2 >= 0 and i2 <= self.PolyVib.maximum_manifold:
            pass
        else:
            return None, None

        if i2 > i:
            mu_key = str(i) + '_to_' + str(i2)
            mu = self.mu[mu_key]
        else:
            mu_key = str(i2) + '_to_' + str(i)
            mu = self.mu[mu_key].transpose(1,0,2)

        j_size = self.H[str(j)].shape[0]
        i_size = self.H[str(i)].shape[0]
        i2_size = self.H[str(i2)].shape[0]
        
        bra_eye = np.eye(j_size)
        old_key = str(i) + str(j)
        new_key = str(i2) + str(j2)
        
        mu_shape = (i2_size*j_size,i_size*j_size,3)
        new_mu = np.zeros(mu_shape,dtype='complex')
        for i in range(3):
            mu_i = np.kron(mu[:,:,i],bra_eye)
            new_mu[:,:,i] = mu_i

        if np.allclose(np.imag(new_mu),0):
            new_mu = np.real(new_mu)
        mu_key = old_key + '_to_' + new_key
        return mu_key, new_mu

    def make_mu_by_manifold_bra(self,old_manifold,change):
        i,j = old_manifold
        i2 = i
        j2 = j + change
        
        if j2 >= 0 and j2 <= self.PolyVib.maximum_manifold:
            pass
        else:
            return None, None

        if j2 > j:
            mu_key = str(j) + '_to_' + str(j2)
            mu = self.mu[mu_key].transpose(1,0,2)
        else:
            mu_key = str(j2) + '_to_' + str(j)
            mu = self.mu[mu_key]

        i_size = self.H[str(i)].shape[0]
        j_size = self.H[str(j)].shape[0]
        j2_size = self.H[str(j2)].shape[0]
        
        ket_eye = np.eye(i_size)
        old_key = str(i) + str(j)
        new_key = str(i2) + str(j2)
        
        mu_shape = (i_size*j2_size,i_size*j_size,3)
        new_mu = np.zeros(mu_shape,dtype='complex')
        for i in range(3):
            mu_i = np.kron(ket_eye,mu[:,:,i].T)
            new_mu[:,:,i] = mu_i

        if np.allclose(np.imag(new_mu),0):
            new_mu = np.real(new_mu)
        mu_key = old_key + '_to_' + new_key
        return mu_key, new_mu

    def make_mu_unseparable_manifolds(self,change,ket_flag):
        mu = self.mu['up']
        if ket_flag:
            mu_key = 'ket'
        else:
            mu = mu.transpose(1,0,2)
            mu_key = 'bra'
        if change == 1:
            mu_key += '_up'
        elif change == -1:
            mu = mu.transpose(1,0,2)
            mu_key += '_down'
        else:
            raise Exception('change must be either +1 or -1')
        H_size = self.H['all_manifolds'].shape[0]
        
        H_eye = np.eye(H_size)
        
        mu_shape = (H_size**2,H_size**2,3)
        new_mu = np.zeros(mu_shape,dtype='complex')
        for i in range(3):
            if ket_flag:
                mu_i = np.kron(mu[:,:,i],H_eye.T)
            else:
                mu_i = np.kron(H_eye,mu[:,:,i].T)
                
            new_mu[:,:,i] = mu_i

        if np.allclose(np.imag(new_mu),0):
            new_mu = np.real(new_mu)
            
        return mu_key, new_mu

    def append_mu_by_manifold(self,old_manifold,change,ket_flag):
        if ket_flag:
            f = self.make_mu_by_manifold_ket
        else:
            f = self.make_mu_by_manifold_bra
        key, mu = f(old_manifold,change)
        if key == None:
            pass
        else:
            self.mu_L_basis[key] = mu

    def set_mu_Liouville_space_unseparable_manifolds(self):
        self.mu_L_basis = dict()
        changes = [-1,1]
        for change,ket_flag in itertools.product(changes,[True,False]):
            mu_key, mu = self.make_mu_unseparable_manifolds(change,ket_flag)
            self.mu_L_basis[mu_key] = mu

    def set_mu_Liouville_space_separable_manifolds(self):
        self.mu_L_basis = dict()
        for i_key in self.manifolds:
            for j_key in self.manifolds:
                manifold = (int(i_key),int(j_key))
                self.append_mu_by_manifold(manifold,1,True)
                self.append_mu_by_manifold(manifold,-1,True)
                self.append_mu_by_manifold(manifold,1,False)
                self.append_mu_by_manifold(manifold,-1,False)

    def set_mu_Liouville_space(self):
        if self.manifolds[0] == 'all_manifolds':
            self.set_mu_Liouville_space_unseparable_manifolds()
        else:
            self.set_mu_Liouville_space_separable_manifolds()
                
    def save_mu(self):
        np.savez(os.path.join(self.save_path,'mu_original_L_basis.npz'),**self.mu_L_basis)
