#!/usr/bin/env python
# Written by Greg Ver Steeg
# See readme.pdf for documentation
# Or go to http://www.isi.edu/~gregv/npeet.html

import warnings

import numpy as np
import numpy.linalg as la
from numpy import log
from scipy.special import digamma
from sklearn.neighbors import BallTree, KDTree

import torch
from torch.nn.functional import normalize
from sklearn.preprocessing import normalize as norm


class MI(object):
    """docstring for MI"""
    def __init__(self):
        super(MI, self).__init__()        


    # CONTINUOUS ESTIMATORS

    def entropy(self, x, k=3, base=2):
        """ The classic K-L k-nearest neighbor continuous entropy estimator
            x should be a list of vectors, e.g. x = [[1.3], [3.7], [5.1], [2.4]]
            if x is a one-dimensional scalar and we have four samples
        """
        assert k <= len(x) - 1, "Set k smaller than num. samples - 1"
        x = np.asarray(x)
        n_elements, n_features = x.shape
        x = self.add_noise(x)
        tree = self.build_tree(x)
        nn = self.query_neighbors(tree, x, k)
        const = digamma(n_elements) - digamma(k) + n_features * log(2)
        #return (const + n_features * np.log(nn).mean()) / log(base)
        return (const + n_features * np.log(nn/np.median(nn)).mean()) / log(base)

    def tc(self, xs, k=3, base=2):
        xs_columns = np.expand_dims(xs, axis=0).T
        entropy_features = [self.entropy(col, k=k, base=base) for col in xs_columns]
        return np.sum(entropy_features) - self.entropy(xs, k, base)


    # UTILITY FUNCTIONS

    def add_noise(self, x, intens=1e-10):
        # small noise to break degeneracy, see doc.
        return x + intens * np.random.random_sample(x.shape)

    def build_tree(self, points):
        if points.shape[1] >= 20:
            return BallTree(points, metric='chebyshev')
        return KDTree(points, metric='chebyshev')

    def query_neighbors(self, tree, x, k):
        return tree.query(x, k=k + 1)[0][:, k]


    #GRAM Matrix

    def get_gram_matrix(self,input_tensor): # input_tensor should be of dims (bs, channels, *)
        bs, channels = input_tensor.shape[:2]
        input_tensor  = input_tensor.view(bs, channels, -1)
        input_tensor = normalize(input_tensor, p=2, dim=2)
        gram = torch.matmul(input_tensor, input_tensor.permute(0,2,1))
        return gram