#!/usr/bin/env python

"""
Mock-up implementation of a general model for discrete character
evolution.
"""

class ParameterValue(object):
    """
    Generic rich value used when specifying model parameters.
    Can be composed of one or more of the following: Function, Distribution, Expression, Fixed Value, etc.
    Can have constrained range (to particular data types or range of values).
    Can have priors.
    """
    def __init__(self):
        pass
    def evaluate(self):
        pass

class CharacterStateDictionary(dictionary):
    """
    Describes the states in our character state transition matrix.
    Determines the number of states, and maps each state label to
    the corresponding matrix index.
    e.g., SimpleNucleotides = {'A', 'C', 'G', 'T'},
          SimpleCodons = {'AGT', 'AGG' ... etc.}
          CanthusRostralis = {'truncate', 'rounded', 'pointed', ... etc.}
    Also tracks/manages frequencies of each character state (with values as ParameterValue objects).
    Should be iterable.
    """
    def __init__(self, labels=[]):
        self.labels = labels
        self.frequencies = {}

class MarkovMatrix(object):
    """
    Encapsulates the fundamental and universal operations available to
    a n-by-n Markov matrix.  Very generic: applicable to any 
    """
    def __init__(self, nstates):
        self.nstates = nstates

class CharacterStateTransitionMatrix(object):
    """
    A rich Markov model describing the probabilities of state transitions.
    """
    def __init__(self, state_dictionary=None):
        self.__state_dictionary = None   # state names/labels + frequencies
        self.__rmatrix = None            # instantaneous rate matrix (do we need this?)
        self.__qmatrix = None            # calculated based on above, plus any other model-specific parameters
        self.__pmatrix = None            # calculated  
        self.__rate_multiplier = None    # a ParameterValue() specifying, for example,
                                         # a discretized gamma distribution of a particular
                                         # shape parameter and number of categories.
                                         # Note that site-specific rates are not modelled here:
                                         # that will be part of the character position mapping.
    def _get_state_dictionary(self):
        return self.__state_dictionary
    def _set_state_dictionary(self, sd):
        self.__state_dictionary=sd
        self.rmatrix=MarkovMatrix(len(self.__state_dictionary))
        self.qmatrix=self.foo1()  # some function that populates this according to the rules of the model
        self.pmatrix=self.foo2()  # ditto

class Edge(object):
    """
    Encapsulates/managers:
        - edge length (ParameterValue)
        - a vector of CharacterStateTransitionMatrix objects that apply to this edge
        - a vector of probabilities to be applied to each CSTM (allows implementation of mixture models)
        - HMM can be incorporated here, perhaps?
          e.g, self_cstm = [cstm1, cstm2, cstm3]
               self_cstm_probs = [ ParameterValue("estimate"), ParameterValue("estimate"), ParameterValue("estimate") ]
    """
    pass

class Graph/Tree(object):
    """
    Branching history of sequences: composed of a set of Edges, and
    rules specifying their connection.
    """
    pass

class CharacterEvolutionModel(object):
    """
    Wraps everything up. Like the Mesquite "bundle".
    """

class DataMap(object):
    """
    For each position, specifies a vector of the CharacterEvolutionModel of
    evolutionary history for that position, with another vector of
    probabilities of each CEM model. Can also implement HMM's here
    maybe?

    Also would incorporate a position x position covariance matrix,
    which describes positional interactions.
    """
    pass