##     _____  ____  _    _ _____   _____  ____  _____  
##   / ____|/ __ \| |  | |  __ \ / ____|/ __ \|  __ \ 
##  | (___ | |  | | |  | | |__) | (___ | |  | | |__) |
##   \___ \| |  | | |  | |  _  / \___ \| |  | |  ___/ 
##   ____) | |__| | |__| | | \ \ ____) | |__| | |     
##  |_____/ \____/ \____/|_|  \_\_____/ \____/|_|     

## Alex Holehouse (Pappu Lab and Holehouse Lab) and Jared Lalmansing (Pappu lab)
## Simulation analysis package
## Copyright 2014 - 2022
##


"""
ssmutualinformation contains all the functions associated with computing things related to mutual
information. This file is not a class, but represents a set of stand alone functions, rather than 
including this function in any one other class.
"""

import numpy as np
import scipy
from .ssexceptions import SSWarning, SSException


# ........................................................................
#
def calc_MI(X, Y, bins, weights=False):
    """
    Calculate the mutual information given two observables, X and Y, 
    and the histogram generated by binning those observables across
    the bins as 1D and as 2D. Can pass a set of weights for the observables
    (i.e. each element in X and Y). If weights provided - requires len(X) == len(Y)
    == len(weights). Note that we could in the future extend this to provide custom
    weights numbers but in general it's likely X and Y will be same length.

    Parameters
    --------------
    X : np.array
        An array or list of numerical observables. Must be same length as Y.

    Y : np.array
        An array or list of numerical observables. Must be same length as X.

    bins : np.array
        A valid set of bins must be uniformly and monotonically increasing in value
        a must expand across the full data range passed in X and in Y.

    weights : np.array
        An array or list of numerical values. If provided means, X,Y and weights must
        all be the same length, and these weights are passed when generating histograms
        from the data

    Returns
    ------------
    Mutual information
        Returns the mutual information. Larger values - greater mutual information but note
        this intrinsically depends on bin-size 

    """

    if len(X) != len(Y):
        raise SSException(f'Error: X and Y vectors must be the same length')

    if np.min(bins) > np.min(X) or np.min(bins) > np.min(Y):
        raise SSException(f'Error: Bins passed to calc_MI in ssmutualinformation() do not straddle the full data range. Bin minimum {np.min(bins)} is bigger than one/both of data minima: X={np.min(X)}, Y={np.min(Y)}')

    if np.max(bins) < np.max(X) or np.max(bins) < np.max(Y):
        raise SSException(f'Error: Bins passed to calc_MI in ssmutualinformation() do not straddle the full data range. Bin max {np.max(bins)} is smaller than one/both of data maxima: X={np.max(X)}, Y={np.max(Y)}')
   
    if weights:
        c_XY = np.histogram2d(X,Y,bins,weights=weights)[0]
        c_X = np.histogram(X,bins,weights=weights)[0]
        c_Y = np.histogram(Y,bins,weights=weights)[0]
    else:
        c_XY = np.histogram2d(X,Y,bins)[0]
        c_X = np.histogram(X,bins)[0]
        c_Y = np.histogram(Y,bins)[0]
              
    # note c_XY is a 2D distribution (i.e. a matrix)
    # while c_X and c_Y are 1D distributions (vector)

    # compute the shannon entropy associated with the 1D
    # and 2D distributions
    H_X = shan_entropy(c_X)
    H_Y = shan_entropy(c_Y)
    H_XY = shan_entropy(c_XY)
   
    # calculate the mutual information
    MI = H_X + H_Y - H_XY

    return MI



# ........................................................................
#
def shan_entropy(c):
    """
    Calculate the shannon entropy of a distribution. Works on a 1D vector 
    or a 2D matrix

    Parameters
    -------------
    c : np.array 
        Array (1D or 2D) or list of numerical values

    Returns
    ------------
    float
        Returns a single value that corresponds to the Shannon entropy
        of the distribution.

    """
    # normalize such that all elements sum up to 1
    c_normalized = c / float(np.sum(c))
    
    # now convert into a single vector of non-zero elements
    c_normalized = c_normalized[np.nonzero(c_normalized)]

    # compute the entropy associated with this vector. The more
    # evenly distributed the greater the entropy
    H = -sum(c_normalized* np.log(c_normalized))  
    return H
