"""

Calculate the entropy of the given h5ad

"""

import os

import sys

import argparse

import logging

import numpy as np

from entropy_estimators import continuous as ce

from pyitlib import discrete_random_variable as drv

import scipy.stats

import anndata as ad

SRC_DIR = os.path.join(

    os.path.dirname(os.path.dirname(os.path.abspath(__file__))), "babel",

)

assert os.path.isdir(SRC_DIR)

sys.path.append(SRC_DIR)

import utils

logging.basicConfig(level=logging.INFO)

"""

NOTES

Estimator for continuous variables

https://github.com/paulbrodersen/entropy_estimators

- This estimator is NOT symmetric

>>> x = np.random.randn(3000, 30000)

>>> x

array([[ 1.01757666,  0.14706194,  0.17207894, ..., -0.5776106 ,

         1.27110965, -0.80688082],

       [-0.46566731, -1.65503883,  0.34362236, ..., -0.56790773,

         1.58161324,  0.6875425 ],

       [ 0.21598618,  0.15462247, -0.66670242, ..., -1.28547741,

        -0.1731192 ,  0.19815154],

       ...,

       [ 0.30699781,  0.24104934,  0.30279376, ...,  1.95658979,

         0.78125961,  0.26259683],

       [-1.94023222, -0.79838041, -0.10267371, ..., -0.67825156,

         0.75047044,  0.773398  ],

       [ 0.73951081,  0.3485434 , -0.17277407, ..., -0.32622845,

        -0.59264903,  1.27659335]])

>>> x.shape

(3000, 30000)

>>> h = continuous.get_h(x)

>>> h

69901.37779787864

>>> h = continuous.get_h(x.T)

>>> h

6346.646780095286

(Simple) estimator for discrete variables

https://docs.scipy.org/doc/scipy/reference/generated/scipy.stats.entropy.html

For a binary variable, we can calculate (base e) entropy in several ways

We can specify a torch distribution, and get entory per dimension

We can ask scipy to calculate this from an input of unnormalized probs

Both give us the same results

>>> b = torch.distributions.bernoulli.Bernoulli(torch.tensor([0.1, 0.9, 0.00001, 0.5]))

>>> b.entropy()

tensor([3.2508e-01, 3.2508e-01, 1.2541e-04, 6.9315e-01])

>>> scipy.stats.entropy([0.1, 0.9])

0.3250829733914482

>>> scipy.stats.entropy([1, 1])  # scipy normalizes the input probs

0.6931471805599453

Another estimator for discrete variables

https://github.com/pafoster/pyitlib

- This supports calculation of joint entropy

>>> x

array([[1, 1, 1, 0],

       [0, 0, 0, 1]])

>>> drv.entropy_joint(x)

0.8112781244591328

>>> drv.entropy_joint(x.T)

1.0

"""

def build_parser():

    """Build CLI parser"""

    parser = argparse.ArgumentParser(

        description=__doc__, formatter_class=argparse.ArgumentDefaultsHelpFormatter

    )

    parser.add_argument("h5ad", type=str, help=".h5ad file to evaluate entropy for")

    parser.add_argument(

        "--discrete", action="store_true", help="Use discrete calculation for entropy"

    )

    return parser

def main():

    """Run script"""

    parser = build_parser()

    args = parser.parse_args()

    adata = ad.read_h5ad(args.h5ad)

    logging.info(f"Read {args.h5ad} for adata of {adata.shape}")

    if args.discrete:

        # Use the discrete algorithm from pyitlib

        # https://pafoster.github.io/pyitlib/#discrete_random_variable.entropy_joint

        # https://github.com/pafoster/pyitlib/blob/master/pyitlib/discrete_random_variable.py#L3535

        # Successive realisations of a random variable are indexed by the last axis in the array; multiple random variables may be specified using preceding axes.

        # In other words, different variables are axis 0, samples are axis 1

        # This is contrary to the default ML format which is samples axis 0, variables axes 1

        # Therefore we must transpose

        input_arr = utils.ensure_arr(adata.X).T

        h = drv.entropy_joint(input_arr, base=np.e)

        logging.info(f"Found discrete joint entropy of {h:.6f}")

    else:

        raise NotImplementedError

if __name__ == "__main__":

    main()