import numpy as np
import itertools
import functools


""" Set our parameters """
# Should we seed the results?
setSeed = False
seedNum = 27

""" Simulation parameters """
# Should we simulate more than once?
setSim = True
simNum = 1

""" Define our data """
# The Statespace
states = np.array(['L', 'w', 'W'])

# Possible sequences of events
transitionName = np.array([['LL', 'Lw', 'LW'],
                           ['wL', 'ww', 'wW'],
                           ['WL', 'Ww', 'WW']])

# Probabilities Matrix (transition matrix)
# Fill in
transitionMatrix = None

# Starting state
startingState = 'w'
initial_dist = None

# Steps to run
stepTime = 2
# End state you want to find probabilites of
endState = 'W'

# Get P_steps
p_steps = False

# Get Stationary Dist
stat_dist = False


# A class that implements the Markov chain to forecast the state/mood:
class markov(object):
    """simulates a markov chain given its states, current state and
    transition matrix.

    Parameters:
    states: 1-d array containing all the possible states
    transitionName: 2-d array containing a list
                   of the all possible state directions
    transitionMatrix: 2-d array containing all
                     the probabilites of moving to each state
    currentState: a string indicating the starting state
    steps: an integer determining how many steps (or times) to simulate"""

    def __init__(self, states: np.array, transitionName: np.array,
                 transitionMatrix: np.array, currentState: str,
                 steps: int):
        super(markov, self).__init__()
        self.states = states
        self.list = list
        self.transitionName = transitionName
        self.transitionMatrix = transitionMatrix
        self.currentState = currentState
        self.steps = steps

    @staticmethod
    def setSeed(num: int):
        return np.random.seed(num)

    @staticmethod
    def p_steps(transitionMatrix, initial_dist, steps):
        for _ in itertools.repeat(None, steps):
            initial_dist = transitionMatrix.T.dot(initial_dist)
        return initial_dist

    @staticmethod
    def stationary_dist(transitionMatrix, initial_dist, steps):
        for _ in itertools.repeat(None, steps):
            initial_dist = transitionMatrix.T.dot(initial_dist)
        return initial_dist

    @functools.lru_cache(maxsize=128)
    def forecast(self):
        print(f'Start state: {self.currentState}')
        # Shall store the sequence of states taken
        self.stateList = [self.currentState]
        i = 0
        # To calculate the probability of the stateList
        self.prob = 1
        while i != self.steps:
            if self.currentState == 'L':
                self.change = np.random.choice(self.transitionName[0],
                                               replace=True,
                                               p=transitionMatrix[0])
                if self.change == 'LL':
                    self.prob = self.prob * 0.8
                    self.stateList.append('L')
                    pass
                elif self.change == 'Lw':
                    self.prob = self.prob * 0.15
                    self.currentState = 'w'
                    self.stateList.append('w')
                else:
                    self.prob = self.prob * 0.05
                    self.currentState = "W"
                    self.stateList.append("W")
            elif self.currentState == "w":
                # Fill in
                pass
            elif self.currentState == "W":
                # Fill in
                pass
            i += 1
        print(f'Possible states: {self.stateList}')
        print(f'End state after {self.steps} steps: {self.currentState}')
        print(f'Probability of all the possible sequence of states:'
              f' {self.prob}\n')
        return self.stateList


def main(*args, **kwargs):
    global startingState
    try:
        simNum = kwargs['simNum']
    except KeyError:
        pass
    sumTotal = 0
    # Check validity of transitionMatrix
    for i in range(len(transitionMatrix)):
        sumTotal += sum(transitionMatrix[i])
        if i != len(states) and i == len(transitionMatrix):
            raise ValueError('Probabilities should add to 1')
    # Set the seed so we can repeat with the same results
    if setSeed:
        markov.setSeed(seedNum)
    # Save our simulations:
    list_state = []
    count = 0
    # Simulate Multiple Times
    if setSim:
        if initial_dist is not None:
            startingState = np.random.choice(states, p=initial_dist)
            for _ in itertools.repeat(None, simNum):
                markovChain = markov(states, transitionName,
                                     transitionMatrix, startingState,
                                     stepTime)
                list_state.append(markovChain.forecast())
                startingState = np.random.choice(states, p=initial_dist)
        else:
            for _ in itertools.repeat(None, simNum):
                markovChain = markov(states, transitionName,
                                     transitionMatrix, startingState,
                                     stepTime)
                list_state.append(markovChain.forecast())
    else:
        for _ in range(1, 2):
            list_state.append(markov(states, transitionName,
                                     transitionMatrix, startingState,
                                     stepTime).forecast())
    for list in list_state:
        if(list[-1] == f'{endState!s}'):
            print(True, list)
            count += 1
        else:
            print(False, list)
    if setSim is False:
        simNum = 1
    print(f'\nThe probability of starting in {startingState} and finishing'
          f' in {endState} after {stepTime} steps is {(count / simNum):.2%}')
    if p_steps:
        print(f'P_{stepTime} is '
              f'{markov.p_steps(transitionMatrix, initial_dist, stepTime)}')
    if stat_dist:
        print(f'Stat dist is {markov.stationary_dist(transitionMatrix,initial_dist, stepTime)}')


if __name__ == '__main__':
    main(simNum=simNum)