import numpy as np
from scipy import linalg
from scipy import sparse
from .markov_chain import markov_chain
import math

class dtmc(markov_chain):
    """
    Implements a finite discrete-time Markov chain (CTMC) 
    
    The chain is defined by its number of states, states, and a 
    transition matrix. The class provides methods to compute both
    stationary and transient metrics, as well as to check the 
    chain properties (aperiodicity and ergodicity). 
    """

    # number of states in string array form
    n_states:int=1

    # states in string array form   
    states:np.array = np.array([1])

    # transition matrix as 2d numpy array
    transition_matrix:np.array
    
    # empty initializer 
    # chain with a single state
    def __init__(self):
        """
        Creates a continuous-time Markov chain from its transition matrix
        """
        self.n_states = 1
        self.transition_matrix = np.array([1])

    # initializer with only the number of states
    def __init__(self,n:int):
        self.n_states = n
        self.transition_matrix = np.eye(n, dtype=float)

    # initializer with a transition matrix
    def __init__(self,transition_matrix:np.array):
        if not self._check_transition_matrix(transition_matrix):#Lets check if transition matrix is logical (i.e the rows sum 1)
            raise ValueError("the rows of transition matrix do not sum 1 or has non positive values")
        self.n_states=transition_matrix.shape[0]
        self.transition_matrix = transition_matrix

    def _check_transition_matrix(self, M:np.array):
        """
        Checks that a matrix is a Markov chain transition matrix
         
        Checks that all entries are non-negative and all row sums are equal to one
        """
        #Check if a given transition matrix has the condition that for every row the sum of all elements is equal to 1
        #Check if a given transition matrix has the condition that every element of the matrix is between 0 and 1
        if np.all(np.isclose(np.sum(M,axis=1),1,1e-5)) and np.all((M>=0) & (M<=1)):
            return True
        else:
            return False
        

    def steady_state(self) -> np.ndarray:
        """
        Computes the steady state distribution of the discrete-time Markov chain

        Computes the steady state probability distribution by replacing one of the 
        matrix equations with a normalizing equation that ensures the result is a 
        probability distribution. 

        Returns the stationary probability distribution in array form
        """
        # computes the steady state distribution
        if self.transition_matrix.any():
            # sets A to be I-P and replace first column with ones for the normalizing equation
            A = self.transition_matrix - np.eye(self.n_states, dtype=float)
            A[:, 0] =  1
            # sets b as the right hand vector with a 1 for the normalizing equation
            b = np.zeros(self.n_states, dtype=float)
            b[0] = 1
            # solves transposed system as pi is a row vector
            res = linalg.solve(A,b,transposed=True)
            return res
        else:    
            print("Empty transition matrix")
        return 0
    

    def transient_probabilities(self,n,alpha):
        """
        Computes the transient distribution at transition n with initial state alpha

        Computes alpha*(P^n)*ones to obtain the probability distribution
        at time/transition n given the initial probability distribution alpha
        """
         
        #First, lets verify n is integer
        if not isinstance(n,int):
            raise ValueError("The number of transitions n must be integer")
        #Now lets verify the dimensions of the vector alpha:
        try:
            shape=alpha.shape
        except:
            print("alpha vector must be a numpy array")
            return None
        if shape[0]!=self.n_states:
            raise ValueError("The dimensions of alpha vector are incorrect. It must be a vector 1xn_states")
        
        #Computes transient_matrix**n
        matrix_n=np.linalg.matrix_power(self.transition_matrix,n)
        vector=alpha@matrix_n
        return vector


    def first_passage_time(self, target:str):
        """
        Computes the expected first passage time to a target state from a start state.

        This method calculates the expected number of steps required for the Markov chain to reach
        the specified target state from the start state by creating a sub-matrix of the transition matrix with the target removed.

        Returns the expected steps to reach the target state from any start state (except target) in array form        
        """
        #The transition matrix and the number of states are brought
        n=self.n_states 
        M=self.transition_matrix.copy()
        # The identity matrix is created with size n-1 states
        i=np.identity(n-1)
        # Matrix of ones with n-1 rows is created
        u=np.full([n-1,1],1)
        # The column and row corresponding to the target state are eliminated from the transition matrix
        p=np.delete(M,target, axis=0)
        p=np.delete(p,target, axis=1)
        
        t=np.matmul(np.linalg.inv(i-p),u)
        return t 

    def occupation_time(self, n:int):
        """
        Computes the expected occupation time in n steps.

        This method calculates the expected time that the discrete-time Markov chain will remain in each state, from all starting states.


        Returns the occupation matrix in array form  
        """
        # Create an identity matrix with the same shape of transition matrix
        occupation=np.eye(self.transition_matrix.shape[0],self.transition_matrix.shape[1])
        # Iterate for each step
        for i in range(1,n+1):
            # Calculate the occupation time for each step, and update the occupation matrix
            occupation+=np.linalg.matrix_power(self.transition_matrix,i)
        return occupation
    
    def period(self):
        """
        Computes the period of a discrete-time Markov Chain.
        
        This method calculates the period of a discrete-time Markov Chain by iterating a total of N steps, where N represents
        the number of states that the chain has. In each k-th iteration, the chain is elevated to the k-th power to compute the
        transient probabilities in the k-th step. If one of the elementsin in the diagonal of the resulting matrix is greater than 0,
        the chain can potentially have a period k. The period of the matrix will be the Greatest Common Divisor (GCD) between the
        potential new period of the matrix in the k-th iteration and the current GCD before the k-th iteration.

        Returns the period of the chain as an int
        """
        
        #Initializes the transition matrix and the greatest common divisor for the periodicity calculations
        P_k = self.transition_matrix
        gcd = 0
        
        #Loop to elevate the transition matrix to the k th power and evaluate its periodicity:
        for k in range(2, (self.n_states)+1):
            P_k = np.dot(self.transition_matrix, P_k)
            #If the transition matrix to the k-th power has a non-zero element in its diagonal, the chain could have a period k:
            if np.any(np.diagonal(P_k) != 0):
                #GCD between the current GCD for the period of the matrix and the potential new GCD
                gcd = math.gcd(gcd, k)
            #If the GCD is 1, there is no need to keep looking for the period, the chain is aperiodic:
            if gcd == 1:
                break
        return gcd
    
    def is_irreducible(self):
        """
        Checks if the given discrete-time Markov Chain is irreducible.
        
        This method determines if the discrete-time Markov Chain is irreducible by checking if, starting in
        any state, it is possible to reach any other state in a sequence of transitions.

        Returns a boolean
        """
        # Brings a copy of the transition matrix
        P = np.copy(self.transition_matrix)
        # Checks if all the states in the Markov chain are strongly connected
        if sparse.csgraph.connected_components(P, directed=True,connection='strong',return_labels=False)==1:
            return True
        else:
            return False

    def is_ergodic(self):
        """
        Checks if a given discrete-time Markov Chain is ergodic.
        
        Given that a finite discrete-time Markov Chain is ergodic if it is aperiodic and irreducible, this
        method checks if both conditions are met to determine if the provided chain is ergodic or not.
        
        Returns a boolean
        """
        # Checks if the chain is irreducible and if it has a period of 1
        if self.period() == 1 & self.is_irreducible() == True:
            return True
        else:
            return False
    



