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

## 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
from .ssexceptions import SSWarning, SSException


# ........................................................................
#
def calc_MI(X, Y, bins, weights=False,normalize=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. Optionally,
    one can normalize the MI by the joint shannon entropy to obtain the NMI. 

    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

    normalize : boolean (Default=False)
        Normalize a given mutual information matrix derived from two observables, X and Y.
        The normalized mutual information is given by:

        NMI = I / H_XY

        where I is the mutual information function and H is the shannon entropy 
        of the joint probability distribution between X and Y.

    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]
              
    H_X = shan_entropy(c_X)
    H_Y = shan_entropy(c_Y)
    H_XY = shan_entropy(c_XY)

    MI = H_X + H_Y - H_XY

    if normalize:
        # Uncomment to *locally* suppress numpy divide by zero runtime warning.
        # This will be cleaned up before returning nmi.
        # with np.errstate(invalid='ignore'):
        nmi = MI / H_XY
        # aforementioned cleanup
        if np.isnan(nmi):
            # print(f"Normalization led to {np.nan} - replacing with 0.")
            nmi = 0
        return nmi
    else:
        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