"""
Function and classes for running kMC

Author: Zeyu Deng
Email: dengzeyu@gmail.com
"""

import random
import numpy as np
from numba.typed import List
import numpy as np
import pandas as pd
from copy import copy
import json
from tracker import Tracker
from event import Event
import pickle as pkl 
from pymatgen.io.vasp import Poscar
import time 

class KMC:
    """
    main function of kinetic monte carlo
    """
    def __init__(self,**kwargs):
        if "api" in kwargs:
            self.api=kwargs["api"]
        else:
            self.api=1
        pass

    def initialization(self,**kwargs):
        """initialization function with capability on dealing with multiple versions of input
        
        This function try to grab the api keyword in kwargs. If fail, then set api=1
        
        Then it will call the initialization function of specified version, pass all kwargs to the initialization1 or initialization2... except the api kwargs (of course)
        
        Notice that the initialization functions have default value for the input! It is recommended to set the parameters explicitly. It is better to overdefining the parameters rather than underdefining parameters. Do not worry passing too much argument. For example, pass immutable_sites=["Zr","O"] to initialization1, The initialization1 will receive it as kwargs, do nothing on it and forget it.
        
        Raises:
            NotImplementedError: if the api parameter is strange

        Returns:
            list: list of kmc.event.Event object, generated by the different version of initialization function
        """
        if "api" in kwargs:
            self.api=kwargs["api"]

            

        return self.initialization3(**kwargs)

 
    def initialization3(self,
    occ='initial_occupation.pkl',
    prim_fname='kmc_supercell_reference_poscar.vasp',
    event_fname="events_for_one_supercell.pkl",
    site_to_event_list_name = 'site_to_event_list.pkl',
    supercell_shape=[2,1,1],
    v=5000000000000,
    T=298,
    convert_to_primitive_cell=False,
    **kwargs):
        """
        XIEWEIHANG 220608
        
        the 3rd version of initialization process. 
        
        
        Args:
            occ (np.array, optional): this is the chebyshev occupation of sites, representing the initial state of the model. Defaults to np.array([-1, -1, -1, -1, -1, -1, -1, -1, 1, -1, -1, -1, -1, -1, -1, 1, -1, 1, -1, -1, 1, -1, -1, -1, -1, -1, -1, -1]).
            prim_fname (str, optional): the primitive cell of the input structure. Defaults to './inputs/prim.cif'.
            fitting_results (str, optional): This is the fitting results matrix, related to the activation barrier E_kra. Defaults to './inputs/fitting_results.json'.
            fitting_results_site (str, optional):  This is the fitting results matrix, related to the site energy E_end Defaults to './inputs/fitting_results_site.json'.
            event_fname (str, optional): record all possible event. Defaults to "./inputs/events.json".
            supercell_shape (list, optional): _description_. Defaults to [2,1,1].
            v (int, optional): vibration frequency. Defaults to 5000000000000.
            T (int, optional): temperature. Defaults to 298.
            lce_fname (str, optional): This contains all lce orbit related to the activation barrier E_kra. Defaults to "./inputs/lce.json".
            lce_site_fname (str, optional): This contains all lce orbit related to the activation barrier E_end. Defaults to "./inputs/lce_site.json".
            immutable_sites (list, optional):the sites that do not participate in the monte carlo process. For example,. in NaSICON, the Zr and O do not participate, the Na/Vac and P/S pairs are considered. Defaults to ["Zr","O"].
            verbose (bool, optional): trying to enable the verbose output. Defaults to False.
            convert_to_primitive_cell(bool): whether or not convert input cif file to primitive cell
            **kwargs: other parameters that are not used in this function. This is for the compatibility of different versions of initialization function.
        Returns:
            list: a list of kmc.event.Event object
        """



        print('Initializing kMC calculations with pirm.cif at',prim_fname,'...')
        # self.structure=StructureKMCpy.from_cif(prim_fname,primitive=convert_to_primitive_cell)
        # supercell_shape_matrix = np.diag(supercell_shape)
        # print('Supercell Shape:\n',supercell_shape_matrix)

        # self.structure.make_supercell(supercell_shape_matrix)

        self.structure = Poscar.from_file(prim_fname).structure
        


        print('Loading occupation:',occ)
        if type(occ)==str:
            with open(occ, 'rb') as h:
                self.occ_global = pkl.load(h)
                self.occ_global = np.array(self.occ_global)
        else:
            self.occ_global=np.array(occ)
    

        if type(event_fname)==str:
            print('Loading events from:',event_fname)
            with open(event_fname,'rb') as handle:
                events_intial = pkl.load(handle)
        else:
            events_intial = event_fname

        
        print('Initializing correlation matrix and E_kra for all events ...')
        events = []
        for event in events_intial:
            event.update_event(self.occ_global,v,T)
            events.append(event)
            
        print('Initializing hopping probabilities ...')
        self.prob_list = [event.probability for event in events]
        self.prob_cum_list = np.cumsum(self.prob_list)
        
        with open(site_to_event_list_name, 'rb') as h: 
            self.site_to_event_list = pkl.load(h)
        return events
    
    
    def propose(self,events,random_seed=123456, rng=np.random.default_rng(),**kwargs): # propose() will propose an event to be updated by update()
        """propose a new event to be updated by update()

        Args:
            events (list): list of event
            random_seed (int or bool, optional): random seed, if None, then no randomseed. This only make sense if use_numpy_random_kernel. Defaults to 114514.
            use_numpy_random_kernel (bool, optional): whether to Use numpy random number generator. Defaults to True.
            api (int, optional): api version. Usage of numpy random kernel is introduced since api=2. Defaults to 2.
            rng (np.random.generator object, optional): a random number generator object. Defaults to np.random.default_rng(). Theoratically this function will receive a random number generator object as a input

        Raises:
            NotImplementedError: for API=1, "propose" function will use the python internal random function. For API=2 and use_numpy_random_kernel=True," propose" function will use the numpy generator class to generate the random numbers. Other functions are not implemented.

        Returns:
            event and time_change: what event is chosed by the random, and the time for this event to occur
        """
    
        random_seed = rng.random()
        random_seed_2 = rng.random()
        proposed_event_index = np.searchsorted(self.prob_cum_list/(self.prob_cum_list[-1]),random_seed,side='right')
        time_change = (-1.0/self.prob_cum_list[-1])*np.log(random_seed_2)

        event = events[proposed_event_index]

        return event, time_change


   
    def update(self,event,events):
        self.occ_global[event.mobile_ion_specie_1_index]*=(-1)
        self.occ_global[event.mobile_ion_specie_2_index]*=(-1)

        events_to_be_updated = np.unique(copy(self.site_to_event_list[event.mobile_ion_specie_1_index]) + copy(self.site_to_event_list[event.mobile_ion_specie_2_index]))# event_to_be_updated= list, include the indices, of the event that need to be updated. 

        events_to_be_updated = np.sort(events_to_be_updated)
        for e_index in events_to_be_updated:
            old_probability = copy(self.prob_list[e_index])
            events[e_index].update_event(self.occ_global,self.v,self.T)
            new_probability = events[e_index].probability
            self.prob_list[e_index] = new_probability

            prob_difference = new_probability-old_probability
            self.prob_cum_list[e_index:]+=prob_difference

        # start_cumsum = time.time()
        # self.prob_cum_list = np.cumsum(self.prob_list)
        # end_cumsum = time.time()
        # print('cum_sum_time: ',end_cumsum-start_cumsum)
   
            

    def run3(self,kmc_pass=1000,
    equ_pass=1,v=1E13,
    T=298,
    events="./inputs/events.json",
    comp=1,
    random_seed=114514,
    use_numpy_random_kernel=True,
    verbose=False,
    mobile_ion_specie='Li',
    q=1.0,
    dimension=3,
    structure_idx=1,
    elem_hop_distance=2.7,
    **kwargs):
        """kmc main function version 3. Previous versions are removed for readability

        Args:
            kmc_pass (int, optional): number of pass to run. Defaults to 1000.
            equ_pass (int, optional): _description_. Defaults to 1.
            v (int, optional): refer to paper. Defaults to 5000000000000.
            T (int, optional): refer to paper, temperature. Defaults to 298.
            events (str, optional): path to event.json. Defaults to "./inputs/events.json".
            comp (int, optional): composition, refer to paper. Defaults to 1.
            random_seed (int, optional): random seed. Should only be provided when test the KMC function. IF not testing as developer ,strictly set this to None. Defaults to 114514.
            use_numpy_random_kernel (bool, optional): whether use numpy random kernel which enable manually assign the random seed. If not testing as a developer, strictly set this to False. Defaults to True.
            verbose (bool, optional): verbose output. Defaults to False.
            mobile_ion_specie (str, optional): mobile ion specie to track. Defaults to 'Na'.
            q (float, optional): charge of mobile ion specie. Defaults to 1.0.
            dimension (int, optional): dimension of migration. For LiCoO2 it is 2 (2D migration). for NaSICON it is 3. Defaults to 3.
            structure_idx (int, optional): structure index . Defaults to 1.
            elem_hop_distance (float, optional): the hopping distance for the mobile ion specie, for NasiCON, this is the distance between Na1 and nearest Na2 site, for LiCoO2, this is the distance between 2 Li site. Defaults to 3.47782.

        Returns:
            kmcpy.tracker.Tracker: return the tracker for testing unit to assert some functions.
        """
        api=3
        
        if use_numpy_random_kernel==True:
            self.rng=np.random.default_rng(seed=random_seed)
        else:
            self.rng=False
        
        print('Simulation condition: v =',v,'T = ',T)
        self.v = v
        self.T = T
        pass_length = len([el.symbol for el in self.structure.species if mobile_ion_specie in el.symbol])

        for current_pass in np.arange(equ_pass):
            for this_kmc in np.arange(pass_length):
                event,time_change = self.propose(events,random_seed=random_seed,rng=self.rng,verbose=verbose)
                self.update(event,events)

        print('Start running kMC ...')
        tracker = Tracker(api=3)

        tracker.initialization(occ_initial=self.occ_global,structure=self.structure,T=self.T,v=self.v,q=q,mobile_ion_specie=mobile_ion_specie,dimension=dimension,elem_hop_distance=elem_hop_distance,**kwargs) #
        # print('Pass\tTime\t\tMSD\t\tD_J\t\tD_tracer\tConductivity\tH_R\t\tf') # < ------- Change this please
        for current_pass in np.arange(kmc_pass):
            print('pass length: ', pass_length)
            start = time.time()
            for this_kmc in np.arange(pass_length):
                # start_propose = time.time()
                event,time_change = self.propose(events,random_seed=random_seed,use_numpy_random_kernel=use_numpy_random_kernel,api=api,rng=self.rng,verbose=verbose) # < ------- Change this please

                tracker.update(event,self.occ_global,time_change)
           
                self.update(event,events)
            _ = tracker.summary()
           

            stop = time.time()
            print('time for one pass_length: ', stop-start)

            # previous_conduct = tracker.summary()
            # tracker.show_current_info(comp,current_pass)

        tracker.write_results(comp,structure_idx,current_pass,self.occ_global)
        if verbose:
            print("verbose information: kmc.KMC.run is called. API version 3")
        return tracker
                        

    # def run(self,kmc_pass,equ_pass,v,T,events):
    #     print('Simulation condition: v =',v,'T = ',T)
    #     self.v = v
    #     self.T = T
    #     pass_length = len([el.symbol for el in self.structure.species if 'Na' in el.symbol])
    #     print('============================================================')
    #     print('Start running kMC ... ')
    #     print('\nInitial occ_global, prob_list and prob_cum_list')
    #     print('Starting Equilbrium ...')
    #     for current_pass in np.arange(equ_pass):
    #         for this_kmc in np.arange(pass_length):
    #             event,time_change = self.propose(events)
    #             self.update(event,events)

    #     print('Start running kMC ...')
    #     tracker = Tracker()
    #     tracker.initialization(self.occ_global,self.structure,self.T,self.v)
    #     print('Pass\tTime\t\tMSD\t\tD_J\t\tD_tracer\tConductivity\tH_R\t\tf\t\tOccNa(1)\tOccNa(2)')
    #     for current_pass in np.arange(kmc_pass):
    #         for this_kmc in np.arange(pass_length):
    #             event,time_change = self.propose(events)
    #             tracker.update(event,self.occ_global,time_change)
    #             self.update(event,events)

    #         previous_conduct = tracker.summary(comp,current_pass)
    #         tracker.show_current_info(current_pass)

    #     tracker.write_results(None,None,current_pass,self.occ_global)

    def as_dict(self):
        d = {"@module":self.__class__.__module__,
        "@class": self.__class__.__name__,
        "structure":self.structure.as_dict(),
        "keci":self.keci,
        "empty_cluster":self.empty_cluster,
        "keci_site":self.keci_site,
        "empty_cluster_site":self.empty_cluster_site,
        "occ_global":self.occ_global}
        return d
    
    def to_json(self,fname):
        print('Saving:',fname)
        with open(fname,'w') as fhandle:
            d = self.as_dict()
            jsonStr = json.dumps(d,indent=4,default=convert) # to get rid of errors of int64
            fhandle.write(jsonStr)
    
    @classmethod
    def from_json(self,fname):
        print('Loading:',fname)
        with open(fname,'rb') as fhandle:
            objDict = json.load(fhandle)
        obj = KMC()
        obj.__dict__ = objDict
        print("load complete")
        return obj

def _convert_list(list_to_convert):
    converted_list = List([List(List(j) for j in i ) for i in list_to_convert])
    return converted_list

