#A datastructure to encode DTMC (discrete-time Markov chanis)

import numpy as np

# Task 1
# a class for encoding a DTMC from a .txt file
class DTMC:
    def __init__(self, name, states, transitions):
        self.name = name
        self.states= states
        #self.transitionName = [] #TODO: implement transition names?
        self.transitionMatrix = transitions
    
    def __str__(self):
        return self.name
    
    def checkProbabilities(self):
        for transitions in self.transitionMatrix:
            if sum(transitions) > 1.0:
                raise ValueError(f"The sum of the probabilities {transitions} is greater than 1.")
    

# Task 2
class probability_distribution:

    def __init__(self, name, DTMC, probabilities):
        self.name = name
        self.DTMC = DTMC
        self.probabilities = probabilities # The length of the probabilities must be equal to the number of states the DTMC has

    def __str__(self):
        return self.name

# Task 3
class MarkovChain:
    # assuming you want to manage one DTMC with an unknown number of probability distributions
    def __init__(self, DTMC, *probability_distributions):
        self.DTMC = DTMC
        self.probability_distributions = list(probability_distributions)

# Task 5
    def print(self, filename=None):
        if filename is None:
            filename = str(self.DTMC) + ".txt"
        path = 'textFiles/'
        f = open(path + filename, 'w')
        
        f.write('MarkovChain ' + str(self.DTMC)+ "\n") # TODO: make DTMC return name as __str__()
        # writing the states
        for i in range(len(self.DTMC.states)):
            for j in range(len(self.DTMC.states)):
                f.write("\t" + self.DTMC.states[i] + " -> " + self.DTMC.states[j]+": " + str(self.DTMC.transitionMatrix[i][j]) + ";\n")
        
        f.write("end \n \n")

        # writing the probability distributions
        for dist in self.probability_distributions:
            f.write('ProbobailityDistribution ' + str(dist) + " of " + str(dist.DTMC) + '\n')
            for i in range(len(dist.probabilities)):
                f.write('\t' + dist.DTMC.states[i] + ': ' + str(dist.probabilities[i]) + ';\n')
            f.write('end \n \n')
        f.close

 #Task 7, 8       
    # Parse and Read state input from files    
    def readStateFromFile(filename):
          fileContext=open(filename).readlines()
          matrix = []
          for line in fileContext:
            #sep = line[:-1].replace(')','').split('(')
            sep = line[:-1].split('(')[1:]
            #print sep
            from ast import literal_eval
            b = [complex(literal_eval('(  '+i)[0],literal_eval('(  '+i)[1]) for i in sep]
            matrix.append(b)
          return np.array(matrix)
                
    # Parse and Read Probability_distributions input from files    
    def readProbability_distributionsFromFile(filename):
          fileContext=open(filename).readlines()
          matrix = []
          for line in fileContext:
            #sep = line[:-1].replace(')','').split('(')
            sep = line[:-1].split('(')[1:]
            #print sep
            from ast import literal_eval
            b = [complex(literal_eval('(  '+i)[0],literal_eval('(  '+i)[1]) for i in sep]
            matrix.append(b)
          return np.array(matrix)

#Task 10     
    # function for Calculateist Probability By Sparse Marix   
    def CalculateistProbabilitYBySparseMarix(): 
        # writing the probability distributions
        sparseMatrix = csr_matrix((3, 4), 
                          dtype = np.int8).toarray()   
        return sparseMatrix     

#Task 4
    # function for Write state to txt files    
    def WriteStateToFile(self, filename=None):
        if filename is None:
            filename = str(self.DTMC) + ".txt"
        path = 'textFiles/'
        f = open(path + filename, 'w')
        
        f.write('MarkovChain ' + str(self.DTMC)+ "\n") # TODO: make DTMC return name as __str__()
        # writing the states
        for i in range(len(self.DTMC.states)):
            for j in range(len(self.DTMC.states)):
                f.write("\t" + self.DTMC.states[i] + " -> " + self.DTMC.states[j]+": " + str(self.DTMC.transitionMatrix[i][j]) + ";\n")
        
        f.write("end \n \n")

    # function for Write DistProbabilities to txt files   
    def writeDistProbabilitiesToFile(self, filename=None): 
        # writing the probability distributions
        if filename is None:
            filename = str(self.DTMC) + ".txt"
        path = 'textFiles/'
        f = open(path + filename, 'w')
        for dist in self.probability_distributions:
            f.write('ProbobailityDistribution ' + str(dist) + " of " + str(dist.DTMC) + '\n')
            for i in range(len(dist.probabilities)):
                f.write('\t' + dist.DTMC.states[i] + ': ' + str(dist.probabilities[i]) + ';\n')
            f.write('end \n \n')
        f.close

# Task 11
                 
    def sojourn():
        fileContext=open(filename).readlines()
        matrix = []
        for line in fileContext:
          #sep = line[:-1].replace(')','').split('(')
          sep = line[:-1].split('(')[1:]
          #print sep
          from ast import literal_eval
          b = [complex(literal_eval('(  '+i)[0],literal_eval('(  '+i)[1]) for i in sep]
          matrix.append(b)
          sojourn_times = [rate for rate in sep + 1]
        print(sojourn_times)