"""
Visualization of Limiting Matrix in Transition Matrix

By
Gurpinder Singh
20 May 2020
"""

import numpy as np
import matplotlib.pyplot as plt


def ChainGenerator(initialMatrix, numGen):
    """
    Input : a transition Matrix (Numpy array) and number of Generations
    
    Takes the nth dot product of the initial matrix with its n-1th dot product,
    stores the results in a 3D array called results.
    
    OUTPUT = Results
    """
    Results = np.zeros([len(initialMatrix), len(initialMatrix[0]), numGen]) #initializing the 3D array
    CurrentGen = initialMatrix[:] #making a copy of initial array
    for i in range(len(CurrentGen)): #appending the Gen0 values to Results
            for j in range(len(CurrentGen)):
                Results[i][j][0] = CurrentGen[i][j]
    
    for G in range(1, numGen): #looping over generations
        CurrentGen = np.dot(initialMatrix, CurrentGen) #taking dot products
        for i in range(len(CurrentGen)): #appending each dot product to results
            for j in range(len(CurrentGen)):
                Results[i][j][G] = CurrentGen[i][j]
                
    return Results


def plotMarkov(Results, labels, title="none", figureSize=(10,10)):
    """
    Plotting function to the results matrix generated by Chain Generator.
    Inputs:
        Labels is the array for row/column labels of the probabilities matrix.
        len(label) should be equal to len(Results)
    """
    plt.figure(figsize = figureSize)
    plt.title(title)
    GX = range(len(Results[0][0]))
    
    for i in range(len(Results)):
        for j in range(len(Results)):
            if labels[i] == labels[j]: #plots that show self Transition
                plt.plot(GX, Results[i][j], "--", label = labels[i] + " to " + labels[j], linewidth = 2.0)
            else: #other transitions are solid lines
                plt.plot(GX, Results[i][j], label = labels[i] + " to " + labels[j], linewidth = 2.0)
    plt.legend(loc = "best")
    plt.xlabel("Generations")
    plt.ylabel("Probabilities")
    plt.savefig(title.replace(" ", "_")+".jpg") #save the plot as .jpg            

# #running the code! (on Assignment questions ofcourse)
# """
# Question Number 1, part 1
# """
# Mat1 = np.array([[.8, 0, .2],
#                 [.2, .7, .1],
#                 [.3, .3, .4]])

# labels1 = np.array(["Harvard", "Darthmouth", "Yale"])
# plotMarkov(ChainGenerator(Mat1, 30), labels1, "Student Transitions over 30 Generations")
# """
# Question Number 1, part 2
# """
# Mat2 = np.array([[1.0, 0, 0],
#                [.2, .7, .1],
#                [.3, .3, .4]])

# labels1 = np.array(["Harvard", "Darthmouth", "Yale"])
# plotMarkov(ChainGenerator(Mat2, 40), labels1, "Student Transitions over 40 Generations, with Harvard to Harvard = 1.0")

# """
# Question Number 2
# """
# Mat3 = np.array([[.8, .1, .1],
#                [.2, .6, .2],
#                [.25, .25, .5]])

# labels2 = np.array(["Professional", "Skilled worker", "Unskilled Worker"])
# plotMarkov(ChainGenerator(Mat3, 30), labels2, "Worker Transitions over 30 Generations")

"""
Absorbing Matrix
"""
Mat = np.array([[1, 0, 0, 0],
                [0, 1, 0, 0],
                [.1, 0, .8, .1],
                [.1, .1, .4, .4]])



labels = np.array(["F", "B", "G", "A"])
plotMarkov(ChainGenerator(Mat, 100), labels, "Absorbing Markov Chain", (12, 12))


