import pandas
import numpy

'''
Interprets a csv file with symmetric labels and data as a transition probability matrix
'''
class TransitionMatrix:

    def __init__(self,filename):
        '''
        Initializes TransitionMatrix object from data in 'filename' csv file
        Stores as DataFrame
        Copies labels of transition states
        '''
        self.transitions = pandas.read_csv(filename, index_col = 0)
        self.states = self.transitions.index
        
    def get_dataframe(self):
        '''
        Returns data as a DataFrame object with appropriate labels
        '''
        return self.transitions
    
    def get_matrix(self):
        '''
        Returns data as a numpy matrix (no labels)
        '''
        return numpy.matrix(self.transitions)
    
    def future_transition(self, years):
        '''
        Calculates the transition matrix corresponding to the cumulative probability
        of transition from/to each state with a specified number of years
        Does this by finding the nth power of the matrix
        '''
        matrix = self.get_matrix()
        for i in range(years - 1):
            matrix = matrix*self.get_matrix()            
        matrix = pandas.DataFrame(matrix, index = self.states, columns = self.states)
        return matrix
    
    def future_transition_to(self, label, years):
        '''
        Calculates a future transition matrix and pulls the column corresponding to 
        the label; useful for calculating default probabilities
        '''
        matrix = self.future_transition(years)
        return matrix[label]