"""
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.
"""


##
##                                       _ _              _ 
##   ___ __ _ _ __ ___  _ __   __ _ _ __(_) |_ _ __ __ _ (_)
##  / __/ _` | '_ ` _ \| '_ \ / _` | '__| | __| '__/ _` || |
## | (_| (_| | | | | | | |_) | (_| | |  | | |_| | | (_| || |
##  \___\__,_|_| |_| |_| .__/ \__,_|_|  |_|\__|_|  \__,_|/ |
##                     |_|                             |__/ 
##
## Alex Holehouse (Pappu Lab and Holehouse Lab)
## Simulation analysis package
## Copyright 2014 - 2021
##


import mdtraj as md
import numpy as np
import scipy
from .ssexceptions import SSWarning




# ........................................................................
#
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) - requires len(X) == len(Y) == len(weights)

    """
   
    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

    """
    # 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
