from .Random import random
import numpy as np
from .logging import logging
from abc import ABC

class sim_base(ABC):
    def __init__(self, events:list, system, verbose=False, name="base",log=True, debug=False):
        # random number generator
        self.rand   = random()
        self.events = events
        self.num_events = len(self.events)
        self.verbose= verbose

        # This should ideally be off
        self.debug     = debug

        # initialize time
        self.t         = 0
        self.system    = system

        # stores the indices of the possible events on each site 
        self.ev_ind    = [[] for _ in range(self.num_events)]

        # All the rate of the events as well as the cumulative sum
        self.ev_R      = np.zeros((self.num_events,))
        self.ev_cumsum = np.zeros((self.num_events,)) 

        # Total rate 
        self.total_rate= 0.0
        self.name      = name
        self.log       = log

        if self.log:
            self.logger       = logging(name)
            self.logger_state = logging("state")

        # track number of times each event is performed
        self.num_times = np.zeros((self.num_events, ))

        if self.log:
            self.init()

    def init(self):
        """
        Function that initializes the logger and anything necessary
        """
        print("Init not implemented.")

    def update(self):
        """
        Update the system and see which new events became possible
        """
        self.ev_ind     = [[] for _ in range(len(self.events))]
        self.total_rate = 0.0
        self.ev_R       = np.zeros((len(self.events),))
        self.ev_cumsum  = np.zeros((len(self.events),)) 

        # update the indices of the system
        for i in range(self.system.lattice.Nx):
            for j in range(self.system.lattice.Ny):
                index     = np.array([i,j])
                for k, ev in enumerate(self.events):
                    if ev.is_possible(self.system, index):
                        rate          = self.events[k].get_rate(self.system, index)
                        self.ev_R[k] += rate
                        self.ev_ind[k].append(index)

        # update the rate of the system
        self.ev_cumsum     = np.cumsum(self.ev_R)
        self.total_rate    = self.ev_R.sum() 

        assert self.total_rate == self.ev_cumsum[-1], \
            "The total rate does not match the largest of the cumsum, something is wrong"

        # Divide cumulative sum by total rate 
        self.ev_cumsum     /= self.total_rate
        
    def step(self):
        """
        Take a Monte Carlo step
        """
        choose_ev     = self.rand.uniform()
        time          = self.rand.exponential(self.total_rate)

        # we just want the number while digitize returns a np.ndarray
        order_        = np.digitize(choose_ev,self.ev_cumsum)[0] 

        # Update the count of the event by 1
        self.num_times[order_] += 1

        # Check the number of sites that are possible for this event
        num_sites     = len(self.ev_ind[order_])
        site_num      = np.random.choice(np.arange(num_sites))
        site_idx      = self.ev_ind[order_][site_num]

        if self.verbose:
            print("Time: {} ms".format(self.t + time[0]))
            for i in range(self.num_events):
                print("The rate for event {} is {}".format(self.events[i].name, self.ev_R[i]))
            print("The event carried out is {}".format(self.events[order_].name))
        
        # perform the event chosen, this will actually do another check (if_possible)
        self.events[order_].do_event(self.system, site_idx)
        self.t       += time[0]

        # Write to data
        if self.log:
            self.log_data()
            self.log_state()

        # check debug 
        if self.debug:
            self.check_if_debug()
 
    def log_data(self):
        """
        Function that writes data to file every time step
        """
        print("log_data not implemented")
    
    def log_state(self):
        """
        This logs the state of the system
        """
        print("log_state not implemented.")
         
    def check_if_debug(self):
        """
        This is a helper function that checks whether or not 2 rectangles overlap if we are in the debug mode 
        """
        print("Debug not implemented.") 
from .Random import random
import numpy as np
from .logging import logging
from abc import ABC

class sim_base(ABC):
    def __init__(self, events:list, system, verbose=False, name="base",log=True, debug=False):
        # random number generator
        self.rand   = random()
        self.events = events
        self.num_events = len(self.events)
        self.verbose= verbose

        # This should ideally be off
        self.debug     = debug

        # initialize time
        self.t         = 0
        self.system    = system

        # stores the indices of the possible events on each site 
        self.ev_ind    = [[] for _ in range(self.num_events)]

        # All the rate of the events as well as the cumulative sum
        self.ev_R      = np.zeros((self.num_events,))
        self.ev_cumsum = np.zeros((self.num_events,)) 

        # Total rate 
        self.total_rate= 0.0
        self.name      = name
        self.log       = log

        if self.log:
            self.logger       = logging(name)
            self.logger_state = logging("state")

        # track number of times each event is performed
        self.num_times = np.zeros((self.num_events, ))

        if self.log:
            self.init()

    def init(self):
        """
        Function that initializes the logger and anything necessary
        """
        print("Init not implemented.")

    def update(self):
        """
        Update the system and see which new events became possible
        """
        self.ev_ind     = [[] for _ in range(len(self.events))]
        self.total_rate = 0.0
        self.ev_R       = np.zeros((len(self.events),))
        self.ev_cumsum  = np.zeros((len(self.events),)) 

        # update the indices of the system
        for i in range(self.system.lattice.Nx):
            for j in range(self.system.lattice.Ny):
                index     = np.array([i,j])
                for k, ev in enumerate(self.events):
                    if ev.is_possible(self.system, index):
                        rate          = self.events[k].get_rate(self.system, index)
                        self.ev_R[k] += rate
                        self.ev_ind[k].append(index)

        # update the rate of the system
        self.ev_cumsum     = np.cumsum(self.ev_R)
        self.total_rate    = self.ev_R.sum() 

        assert self.total_rate == self.ev_cumsum[-1], \
            "The total rate does not match the largest of the cumsum, something is wrong"

        # Divide cumulative sum by total rate 
        self.ev_cumsum     /= self.total_rate
        
    def step(self):
        """
        Take a Monte Carlo step
        """
        choose_ev     = self.rand.uniform()
        time          = self.rand.exponential(self.total_rate)

        # we just want the number while digitize returns a np.ndarray
        order_        = np.digitize(choose_ev,self.ev_cumsum)[0] 

        # Update the count of the event by 1
        self.num_times[order_] += 1

        # Check the number of sites that are possible for this event
        num_sites     = len(self.ev_ind[order_])
        site_num      = np.random.choice(np.arange(num_sites))
        site_idx      = self.ev_ind[order_][site_num]

        if self.verbose:
            print("Time: {} ms".format(self.t + time[0]))
            for i in range(self.num_events):
                print("The rate for event {} is {}".format(self.events[i].name, self.ev_R[i]))
            print("The event carried out is {}".format(self.events[order_].name))
        
        # perform the event chosen, this will actually do another check (if_possible)
        self.events[order_].do_event(self.system, site_idx)
        self.t       += time[0]

        # Write to data
        if self.log:
            self.log_data()
            self.log_state()

        # check debug 
        if self.debug:
            self.check_if_debug()
 
    def log_data(self):
        """
        Function that writes data to file every time step
        """
        print("log_data not implemented")
    
    def log_state(self):
        """
        This logs the state of the system
        """
        print("log_state not implemented.")
         
    def check_if_debug(self):
        """
        This is a helper function that checks whether or not 2 rectangles overlap if we are in the debug mode 
        """
        print("Debug not implemented.") 
# random number generator # This should ideally be off # initialize time # stores the indices of the possible events on each site # All the rate of the events as well as the cumulative sum # Total rate # track number of times each event is performed Function that initializes the logger and anything necessary Update the system and see which new events became possible # update the indices of the system # update the rate of the system # Divide cumulative sum by total rate Take a Monte Carlo step # we just want the number while digitize returns a np.ndarray # Update the count of the event by 1 # Check the number of sites that are possible for this event # perform the event chosen, this will actually do another check (if_possible) # Write to data # check debug Function that writes data to file every time step This logs the state of the system This is a helper function that checks whether or not 2 rectangles overlap if we are in the debug mode