#!/usr/bin/env python3
"""Entropy and information theory related calculations.

**Author: Jonathan Delgado**

"""

######################## Imports ########################


import numpy as np

import stp


######################## Helper functions ########################


def _eps_filter(x):
    """ Checks if the value is within machine-epsilon of zero and maps it to 
        zero if it is the case. Useful for removing negative values in entropies that should otherwise be zero.
        
        Args:
            x (float): value to be checked.
    
        Returns:
            (float): x if the value is not within machine epsilon of zero, 0 otherwise.
    
    """
    return x if not np.isclose(x, 0, atol=9*10E-15) else 0
    

######################## Entropy calculations ########################


def entropy(p):
    """ Calculates the Shannon entropy for a marginal distribution.

        Args:
            p (np.ndarray): the marginal distribution.

        Returns:
            (float): the entropy of p

    """
    # Since zeros do not contribute to the Shannon entropy by definition, we 
        # ignore them to avoid any errors/warnings.
    p = p[p != 0]

    H = -np.dot(p, np.log(p))
    # Filter against machine epsilon
    return _eps_filter(H)


def delta_entropy(R, p):
    """ Calculates the discrete time change in entropy using the entropy of p 
        evolved with R, minus the entropy of p.
    
        Args:
            R (np.ndarray): the transition matrix.

            p (np.ndarray): the marginal distribution.
    
        Returns:
            (float): the change in entropy
    
    """
    return entropy(step(R, p)) - entropy(p)


def relative_entropy(p, q):
    """ Calculates the Kullback-Leibler divergence, which is nonnegative and 
        vanishes if and only if the distributions coincide.
    
        Args:
            p, q (np.ndarray): the probability distributions.
    
        Returns:
            (float): the relative entropy.
    
    """
    if p.shape[0] != q.shape[0]:
        print('Dimensions of vectors are not equal. Cannot find relative entropy.')
        sys.exit()

    # Any values where p is zero are defined to be zero and hence do not
        # contribute to the relative entropy
    # By masking q as well we automatically skip the values that were supposed
        # to vanish with p avoiding any misalignment issues
    # Note that by masking q only where p is zero doesn't remove
        # any mismatching meaning it will still be infinite (as it should be)
        # in the case where q has a zero that p does not.
    p_filtered = p[p != 0]
    log_ratio = np.log(p_filtered / q[p != 0])

    return np.dot(p_filtered, log_ratio)


def entropy_production(matrix, p, discrete=True):
    """ Calculates the entropy production for either discrete or continuous 
        time.
    
        Args:
            matrix (np.ndarray): the stochastic matrix, either a discrete time transition matrix or a continuous time rate matrix.

            p (np.ndarray): the marginal distribution

        Kwargs:
            discrete (bool): True if we are calculating the discrete time entropy production (nats), False if we are calculating it in continuous time (nats/time).
    
        Returns:
            (float/np.inf): the entropy production
    
    """
    log_product = matrix * np.log( matrix / matrix.T )
    # The entropy term only exists in the case of discrete time
        # it vanishes when we calculate the continuous time EP,
        # by multiplying by the boolean we include it only when
        # necessary
    EP = np.dot(log_product.sum(axis=0), p) - (entropy(p) * discrete) \
        - np.dot(stp.step(matrix, p), np.log(p))
    return EP


def entropy_flow(R, p):
    """ Calculates the discrete time entropy flow. This has not been    
        generalized to handle the continuous time entropy flow yet.
    
        Args:
            R (np.ndarray): the discrete time transition matrix
            
            p (np.ndarray): the marginal distribution
    
        Returns:
            (float): the entropy flow
    
    """
    # Vectorized calculation
    log_product = R * np.log( R / R.T )
    p_step = step(R, p)
    EF = -np.dot(log_product.sum(axis=0), p) + entropy(p_step) \
        + np.dot(p_step, np.log(p))
    return EF


######################## Entropy rates ########################


def entropy_rate(R):
    """ Calculates the asymptotic entropy rate for the provided transition 
        matrix. If the matrix is time-inhomogeneous then we return a function that generates the entropy_rate as a function of n by calculating the systems limiting distribution for each n.
        
        Args:
            R (np.ndarray/function): the transition matrix.
    
        Returns:
            (float/function): the entropy velocity.
    
    """
    if callable(R):
        return lambda n : entropy_rate(R(n))

    pst = stp.get_stationary_distribution(R, discrete=True)
    RProduct = (R * np.log(R)).sum(axis=0)

    return -np.dot(pst, RProduct)


######################## Information Space Objects ########################


class InfoSpace:
    """ Information space. Holds collections of paths that traverse states in a 
        state space as a matrix, and the probability of each of those paths. 
        
        Provides functionality on this path space such as providing path entropies.
        
        Attributes:
            paths: the matrix of paths.

            probabilities: a list of probabilities each path.

            num_paths: the number of paths considered.

            path_length: the length of the paths considered.

            probabilities: a matrix where the (i,j)th element is the probability of observing the first j states of the ith path.

            entropies: a list of path entropies for each path

            total_probability: the sum of the probabilities of each path.

    """

    def __init__(self, paths, p_matrix):
        """ Initializes the InfoSpace object.
            
            Args:
                paths (np.ndarray): a matrix of paths where the (i,j)th element corresponds to the jth symbol of the ith path.

                p_matrix (np.ndarray): a matrix of probabilities where the (i,j)th element corresponds to the probability of observing the ith path for the first j+1 (zero-indexing) symbols.
        
        """
        self._paths = np.array(paths)
        
        # Matrix of probabilities corresponding to the probability for the path
            # at each moment.
        self._p_matrix = np.array(p_matrix)
        
        if self._p_matrix.size != 0:
            # The information space is not empty
            self._probabilities = self._p_matrix[:, -1]
        else:
            # There is zero probability here.
            self._probabilities = 0


    #------------- Properties -------------#


    @property
    def paths(self):
        return self._paths


    @property
    def num_paths(self):
        return self.paths.shape[0]


    @property
    def path_length(self):
        return self.paths.shape[1]
    

    @property
    def probabilities(self):
        return self._probabilities


    @property
    def entropies(self):
        """ Returns a list of path entropies for each corresponding path 
            probability.

        """
        try:
            return self._entropies
        except AttributeError:
            # It's never been calculated before
            self._entropies = -np.log(self.probabilities)

        return self._entropies


    @property
    def total_probability(self):
        try:
            return self.probabilities.sum()
        except AttributeError:
            # Space is empty
            return 0
    

    #------------- Static methods -------------#


    @staticmethod
    def shorten(infospace, path_length, return_index=False):
        """ Takes an Information Space and shortens it. Since unique paths of 
            length n, may be degenerate when truncated to paths of length m < n, we need to check for degeneracies and filter them out in both paths and probabilities.
            
            Args:
                infospace (InfoSpace): the information space to shorten.

                path_length (int): the path length the information space should be shortened to.

            Kwargs:
                return_index (bool): returns the indices of the non-degenerate paths for the given path length using the original matrix. Useful for filtering other quantities of interest that may not be attached to this object.      
        
            Returns:
                (InfoSpace): the shortened InfoSpace.
        
        """
        if path_length < 1:
            raise ValueError(f'Invalid path length: {path_length}. Path length must be an integer greater than 0.')
        elif path_length > infospace.path_length:
            raise ValueError(f'Cannot shorten an InformationSpace from length: {infospace.path_length} -> {path_length}.')

        if infospace.paths.size == 0:
            # This is an empty information space
            return infospace if not return_index else (infospace, [])

        # Truncate the path matrix
        paths = infospace.paths[:, :path_length]
        # Return index will provide the path indices of the non-degenerate paths
        _, indices = np.unique(paths, axis=0, return_index=True)
        # Sort the indices
        indices = sorted(indices)
        # Filter out the paths. Not taken from np.unique to ensure the correct
            # ordering.
        paths = paths[indices, :]
        # Truncate the probability matrix
        p_matrix = infospace._p_matrix[:, :path_length]
        # Filter the probabilities matrix
        p_matrix = p_matrix[indices, :]

        infospace = InfoSpace(paths, p_matrix)
        return infospace if not return_index else infospace, indices


class PartitionedInfoSpace(InfoSpace):
    """ Partitioned Information Space. Constructs a typical set on an 
        information space to partition it into a typical information space and an atypical one. 

        Holds path probabilities, typical paths, atypical paths, atypical path probabilities and more. This object will use a provided (often sampled) path space to partition the space into a collection of typical and atypical paths depending on the dynamics provided. Will also track other quantities of interest such as the upper and lower bounds on the path probabilities required for the paths to be considered typical.
    
        Attributes:
            paths: the matrix of paths.

            probabilities: a list of probabilities each path.

            num_paths: the number of paths considered.

            path_length: the length of the paths considered.

            probabilities: a matrix where the (i,j)th element is the probability of observing the first j states of the ith path.

            entropies: a list of path entropies for each path.

            entropy_rates: a list of the entropy rates for each various path length. This will be the center of the epsilon-neighborhood for path entropies to qualify paths as typical for.

            epsilon: the widths of the neighborhood used for paths to be considered typical for each path length.

            upper/lower: the upper/lower bounds as measured in nats. This means that a path is typical if and only if its path entropy rate is within these bounds.

            typicalities: a matrix where the (i,j)th element is a boolean determining whether the ith path is typical after j+1 steps.

            ts: the typical set.

            ats: the atypical set.

    """

    def __init__(self, entropy_rates, epsilon, paths=None, p_matrix=None, typical_space=None, atypical_space=None):
        """ Generates the PartitionedInfoSpace.
            
            Args:
                entropy_rates (np.ndarray): a list of the entropy rates for each various path length. This will be the center of the epsilon-neighborhood for path entropies to qualify paths as typical for.

                epsilon (np.ndarray): the widths of the neighborhood used for paths to be considered typical for each path length.

            Kwargs:
                paths (np.ndarray/None): the entire sampled path space, the union of the typical and atypical spaces. If not provided these spaces will be merged to generate it.

                p_matrix (np.ndarray/None): the entire matrix of probabilities for each path and each path length. If not provided, this will be generated by merging the p_matrix of the typical and atypical spaces.

                typical_space (InfoSpace/None): the typical set on this space. If None, partitions the provided path space.

                atypical_space (InfoSpace): the atypical set on this space. If None, partitions the provided path space.
        
        """
        # Bool if the space simply needs to be partitioned
        must_partition = (paths is None) or (p_matrix is None)
        # Bool if the space simply needs to be merged since it's already been 
            # partitioned into a typical and atypical space
        must_union = (typical_space is None) or (atypical_space is None)

        if must_partition and must_union:
            # We need either the paths AND the p_matrix or the tupical/atypical 
                # spaces to partition/union the spaces respectively.
            raise TypeError('In sufficient information provided to partition/union the Information Space. We need either paths with their probabilities or the already partitioned spaces.')


        if must_partition:
            # Partition the paths and probability matrix into a typical and
                # atypical space

            # Need to generate the upper/lower bounds for the partitioning
                # of the spaces
            self._lower = entropy_rates - epsilon
            self._upper = entropy_rates + epsilon

            ts_paths = []; ts_p_matrix = []
            ats_paths = []; ats_p_matrix = []

            for path, path_index in enumerate(paths):
                path_prob = p_matrix[path_index]
                # The path entropy rate for direct comparison with the
                    # upper/lower bounds
                path_entropy_rate = -np.log(path_prob[-1]) / path_length

                is_typical = (
                    (self.lower[-1] <= path_entropy_rate)
                    and (path_entropy_rate <= self._upper)
                )

                if is_typical:
                    ts_paths.append(path)
                    ts_p_matrix.append(path_prob)
                else:
                    ats_paths.append(path)
                    ats_p_matrix.append(path_prob)

                typical_space = InfoSpace(ts_paths, ts_p_matrix)
                atypical_space = InfoSpace(ats_paths, ats_p_matrix)

        elif must_union:
            # Union the path data
            ts_empty = (typical_space.paths.size == 0)
            ats_empty = (atypical_space.paths.size == 0)

            if not ts_empty and not ats_empty:
                # Both are nonempty
                paths = np.vstack( (typical_space.paths, atypical_space.paths) )
                p_matrix = np.vstack(
                    (typical_space._p_matrix, atypical_space._p_matrix)
                )
            elif ts_empty:
                # Only the typical_space is empty
                paths = atypical_space.paths
                p_matrix = atypical_space._p_matrix
            else:
                # Only the atypical_space is empty
                paths = typical_space.paths
                p_matrix = typical_space._p_matrix

        ### Storing properties ###
        self._paths = paths
        self._p_matrix = p_matrix

        self._probabilities = self._p_matrix[:, -1]

        self._entropy_rates = entropy_rates

        # Generalize the epsilon to a path_length dependent epsilon for
            # potential generalizations in child classes.
        if isinstance(epsilon, list):
            epsilon = np.array(epsilon)
        if not isinstance(epsilon, np.ndarray):
            # We were only provided a float
            epsilon = np.full(self.path_length, epsilon)
        self._epsilon = epsilon

        self._ts = typical_space
        self._ats = atypical_space


    #------------- Properties -------------#


    @property
    def entropy_rates(self):
        return self._entropy_rates
    

    @property
    def epsilon(self):
        return self._epsilon


    @property
    def upper(self):
        try:
            return self._upper
        except AttributeError:
            # It's never been calculated before
            self._upper = self.entropy_rates + self.epsilon        

        return self._upper


    @property
    def lower(self):
        try:
            return self._lower
        except AttributeError:
            # It's never been calculated before.
            self._lower = self.entropy_rates - self.epsilon

        return self._lower


    @property
    def typicalities(self):
        """ Returns the matrix of typicalities. """
        try:
            return self._typicalities
        except AttributeError:
            # It's never been calculated before
            typicalities = []
            ns = np.arange(1, self.path_length + 1)

            path_entropy_rates = -np.log(self._p_matrix) / ns

            self._typicalities = (
                (self.lower <= path_entropy_rates)
                & (path_entropy_rates <= self.upper)
            )

        return self._typicalities


    @property
    def ats(self):
        return self._ats


    @property
    def ts(self):
        return self._ts
    
    
    #------------- Static methods -------------#


    @staticmethod
    def shorten(pinfospace, path_length, return_index=False):
        """ Takes a PartitionedInformationSpace and shortens it. Since unique 
            paths of length n, may be degenerate when truncated to paths of length m < n, we need to check for degeneracies and filter them out in both paths and probabilities.
            
            Args:
                pinfospace (PartitionedInfoSpace): the partitioned information space to shorten.

                path_length (int): the path length the information space should be shortened to.    

            Kwargs:
                return_index (bool): returns the indices of the non-degenerate paths for the given path length using the original matrix. Useful for filtering other quantities of interest that may not be attached to this object.    
        
            Returns:
                (PartitionedInfoSpace): the shortened PartitionedInfoSpace.
        
        """
        # Hold the current information space to access properties
        old_pinfospace = pinfospace
        # Call parent method
        # Paths and p_matrix will be handled here along with any other 
            # properties shared with parent. Sorted indices of non-degenerate 
            # paths will be calculated here too.
        pinfospace, indices = InfoSpace.shorten(old_pinfospace, path_length, return_index=True)
        
        # Finish the rest of this object's specific properties

        # Truncate the entropy_rates
        entropy_rates = old_pinfospace.entropy_rates[:path_length]

        # Truncate the epsilon
        epsilon = old_pinfospace.epsilon[:path_length]

        # Truncate the typicalities matrix
        # Necessary to re-partition the space.
        # Filter out the typicalities matrix
        typicalities = old_pinfospace.typicalities[indices, :path_length]

        ### Partitioning ###
        ts_paths, ts_p_matrix = [], []
        ats_paths, ats_p_matrix = [], []

        paths = pinfospace.paths
        p_matrix = pinfospace._p_matrix
        for path_index, is_typical in enumerate(typicalities[:, -1]):
            path = paths[path_index]
            probs = p_matrix[path_index]

            if is_typical:
                ts_paths.append(path)
                ts_p_matrix.append(probs)
            else:
                ats_paths.append(path)
                ats_p_matrix.append(probs)

        # The partitioned spaces
        ts = InfoSpace(ts_paths, ts_p_matrix)
        ats = InfoSpace(ats_paths, ats_p_matrix)

        pinfospace = PartitionedInfoSpace(entropy_rates=entropy_rates, epsilon=epsilon, paths=paths, p_matrix=p_matrix, typical_space=ts, atypical_space=ats)
        
        # Save the pre-generated property
        pinfospace._typicalities = typicalities

        return pinfospace if not return_index else pinfospace, indices


    @staticmethod
    def partition_space(R, p, paths, epsilon=0.5, return_p=False):
        """ Partitions a path space using the dynamics provided.

            Args:
                R (np.ndarray/function): the transition matrix, time-dependent if provided as a function.

                p (np.ndarray): the initial marginal distribution.

                paths (np.ndarray): the portion of the path space to use.
        
            Kwargs:
                epsilon (float/np.ndarray): the radius/radii of the epsilon neighborhood to consider paths to be typical within.

                return_p (bool): False, return only the PartitionedInfoSpace, True returns both the PartitionedInfoSpace and a list of the marginal vs time.
        
            Returns:
                (ParitionedInfoSpace/2-tuple): the PartitionedInfoSpace (PIS) or the PIS and a list of the marginal versus observation step if return_p is True.

        """

        #------------- Data preparation -------------#

        # Convert the transition matrix to add time-dependence as a constant 
            # matrix if a constant matrix was provided
        if not callable(R):
            # Not being saved as an attribute since this is not easily
                # recoverable by being saved to a file.
            # Emphasize saving properties that can be saved/loaded.
            oldR = R
            R = lambda n : oldR

        num_paths, path_length = paths.shape

        p_matrix = np.zeros(paths.shape)
        # Initialize the marginal distribution data
        for x, path in enumerate(paths):
            # Just equal to the initial marginal
            p_matrix[x, 0] = p[path[0]]
        
        # Used for the bounds
        entropy_rates = np.array([
            entropy_rate(R(i))
            for i in range(path_length)
        ])

        # The marginal versus time
        if return_p: p_vs_time = [p]

        #------------- Data gathering -------------#

        # bar = gui.ProgressBar(path_length * num_paths, width=300, title='Gathering data...')

        ### Quantities versus time ###
        for current_path_length in range(2, path_length + 1):
            # The data index
            i = current_path_length - 1
            # Since the marginals are zero-indexed as are the paths
            step_index = current_path_length - 2

            currentR = R(current_path_length - 1)
            # Propagate the marginal one step and save it separately
                # for quantities like the temporal coarse graining term
            pstep = stp.step(currentR, p)

            ### Path probability calculations ###
            for x, path in enumerate(paths):

                current_state = path[step_index]
                jump_state = path[step_index + 1]

                # Forward calculations
                # Recursive calculation to save time
                last_joint = p_matrix[x, i - 1]
                jump_prob = currentR[jump_state, current_state]
                p_matrix[x, i] = last_joint * jump_prob

            # If updated in each iteration, slows down the simulation 
                # drastically
            # bar.update(amount=num_paths)

            if return_p: p_vs_time.append(pstep)
            # Finished data gathering for this iteration, propagate marginal
                # forward in time
            p = pstep

        # bar.finish()

        ### Partitioning ###
        upper = entropy_rates + epsilon
        lower = entropy_rates - epsilon

        ts_paths, ts_p_matrix = [], []
        ats_paths, ats_p_matrix = [], []

        path_entropy_rates = -np.log(p_matrix[:, -1]) / path_length

        # Identify the paths that are typical and atypical
        for path_index, path_entropy_rate in enumerate(path_entropy_rates):
            # Can't create typicality matrix since partitioning it will
                # break the ordering
            # Determines whether this path is ultimately typical
            is_typical = (
                (lower[-1] <= path_entropy_rate)
                and (path_entropy_rate <= upper[-1])
            )

            probs = p_matrix[path_index]

            if is_typical:
                ts_paths.append(path)
                ts_p_matrix.append(probs)
            else:
                ats_paths.append(path)
                ats_p_matrix.append(probs)

        # The partitioned spaces
        ts = InfoSpace(ts_paths, ts_p_matrix)
        ats = InfoSpace(ats_paths, ats_p_matrix)

        pinfospace = PartitionedInfoSpace(
            entropy_rates=entropy_rates,
            epsilon=epsilon,
            paths=paths,
            p_matrix=p_matrix,
            typical_space=ts,
            atypical_space=ats
        )

        # Set pre-calculated properties
        pinfospace._upper = upper
        pinfospace._lower = lower

        return (pinfospace, p_vs_time) if return_p else pinfospace


######################## Entry ########################


def main():
    print('info.py')
    
    ### Testing ###
    p = stp.rand_p(3)
    R = stp.self_assembly_transition_matrix()
    paths = stp.complete_path_space(3, 4)
    pinfospace = PartitionedInfoSpace.partition_space(R, p, paths)
    print( f'pinfospace.total_probability: {pinfospace.total_probability}' )
    print(pinfospace.ats.num_paths)


if __name__ == '__main__':
    main()
