# AUTOGENERATED! DO NOT EDIT! File to edit: 02_joint_entropy_computation.ipynb (unless otherwise specified).

__all__ = ['JointEntropy', 'ExactJointEntropy', 'batch_multi_choices', 'gather_expand', 'SampledJointEntropy',
           'DynamicJointEntropy']

# Cell

import torch
from toma import toma
from tqdm.auto import tqdm

# Cell


class JointEntropy:
    """Random variables (all with the same # of categories $C$) can be added via `JointEntropy.add_variables`.

`JointEntropy.compute` computes the joint entropy.

`JointEntropy.compute_batch` computes the joint entropy of the added variables with each of the variables in the provided batch probabilities in turn."""
    def compute(self) -> torch.Tensor:
        """Computes the entropy of this joint entropy."""
        raise NotImplementedError()

    def add_variables(self, probs_N_K_C: torch.Tensor) -> 'JointEntropy':
        """Expands the joint entropy to include more terms."""
        raise NotImplementedError()

    def compute_batch(self,
                      probs_B_K_C: torch.Tensor,
                      output_entropies_B=None) -> torch.Tensor:
        """Computes the joint entropy of the added variables together with the batch (one by one)."""
        raise NotImplementedError()

# Cell


class ExactJointEntropy(JointEntropy):
    joint_probs_M_K: torch.Tensor

    def __init__(self, joint_probs_M_K: torch.Tensor):
        self.joint_probs_M_K = joint_probs_M_K

    @staticmethod
    def empty(K: int, device=None, dtype=None) -> 'ExactJointEntropy':
        return ExactJointEntropy(torch.ones((1, K), device=device,
                                            dtype=dtype))

    def compute(self) -> torch.Tensor:
        probs_M = torch.mean(self.joint_probs_M_K, dim=1, keepdim=False)
        nats_M = -torch.log(probs_M) * probs_M
        entropy = torch.sum(nats_M)
        return entropy

    def add_variables(self, probs_N_K_C: torch.Tensor) -> 'ExactJointEntropy':
        assert self.joint_probs_M_K.shape[1] == probs_N_K_C.shape[1]

        N, K, C = probs_N_K_C.shape
        joint_probs_K_M_1 = self.joint_probs_M_K.t()[:, :, None]

        # Using lots of memory.
        for i in range(N):
            probs_i__K_1_C = probs_N_K_C[i][:, None, :].to(joint_probs_K_M_1,
                                                           non_blocking=True)
            joint_probs_K_M_C = joint_probs_K_M_1 * probs_i__K_1_C
            joint_probs_K_M_1 = joint_probs_K_M_C.reshape((K, -1, 1))

        self.joint_probs_M_K = joint_probs_K_M_1.squeeze(2).t()
        return self

    def compute_batch(self,
                      probs_B_K_C: torch.Tensor,
                      output_entropies_B=None):
        assert self.joint_probs_M_K.shape[1] == probs_B_K_C.shape[1]

        B, K, C = probs_B_K_C.shape
        M = self.joint_probs_M_K.shape[0]

        if output_entropies_B is None:
            output_entropies_B = torch.empty(B,
                                             dtype=probs_B_K_C.dtype,
                                             device=probs_B_K_C.device)

        pbar = tqdm(total=B, desc="ExactJointEntropy.compute_batch", leave=False)

        @toma.execute.chunked(probs_B_K_C,
                              initial_step=1024,
                              dimension=0)
        def chunked_joint_entropy(chunked_probs_b_K_C: torch.Tensor,
                                  start: int, end: int):
            b = chunked_probs_b_K_C.shape[0]

            probs_b_M_C = torch.empty((b, M, C),
                                      dtype=self.joint_probs_M_K.dtype,
                                      device=self.joint_probs_M_K.device)
            for i in range(b):
                torch.matmul(self.joint_probs_M_K,
                             chunked_probs_b_K_C[i].to(self.joint_probs_M_K,
                                               non_blocking=True),
                             out=probs_b_M_C[i])
            probs_b_M_C /= K

            output_entropies_B[start:end].copy_(torch.sum(
                -torch.log(probs_b_M_C) * probs_b_M_C, dim=(1, 2)),
                                                non_blocking=True)

            pbar.update(end - start)

        pbar.close()

        return output_entropies_B

# Cell
def batch_multi_choices(probs_b_C, M: int):
    """
    probs_b_C: Ni... x C

    Returns:
        choices: Ni... x M
    """
    probs_B_C = probs_b_C.reshape((-1, probs_b_C.shape[-1]))

    # samples: Ni... x draw_per_xx
    choices = torch.multinomial(probs_B_C, num_samples=M, replacement=True)

    choices_b_M = choices.reshape(list(probs_b_C.shape[:-1]) + [M])
    return choices_b_M


def gather_expand(data, dim, index):
    if gather_expand.DEBUG_CHECKS:
        assert len(data.shape) == len(index.shape)
        assert all(dr == ir or 1 in (dr, ir) for dr, ir in zip(data.shape, index.shape))

    max_shape = [max(dr, ir) for dr, ir in zip(data.shape, index.shape)]
    new_data_shape = list(max_shape)
    new_data_shape[dim] = data.shape[dim]

    new_index_shape = list(max_shape)
    new_index_shape[dim] = index.shape[dim]

    data = data.expand(new_data_shape)
    index = index.expand(new_index_shape)

    return torch.gather(data, dim, index)


gather_expand.DEBUG_CHECKS = False

# Cell


class SampledJointEntropy(JointEntropy):
    """Random variables (all with the same # of categories $C$) can be added via `SampledJointEntropy.add_variables`.

`SampledJointEntropy.compute` computes the joint entropy.

`SampledJointEntropy.compute_batch` computes the joint entropy of the added variables with each of the variables in the provided batch probabilities in turn."""
    sampled_joint_probs_M_K: torch.Tensor

    def __init__(self, sampled_joint_probs_M_K: torch.Tensor):
        self.sampled_joint_probs_M_K = sampled_joint_probs_M_K

    @staticmethod
    def empty(K: int, device=None, dtype=None) -> 'SampledJointEntropy':
        return SampledJointEntropy(
            torch.ones((1, K), device=device, dtype=dtype))

    @staticmethod
    def sample(probs_N_K_C: torch.Tensor, M: int) -> 'SampledJointEntropy':
        K = probs_N_K_C.shape[1]

        # S: num of samples per w
        S = M // K

        choices_N_K_S = batch_multi_choices(probs_N_K_C, S).long()

        expanded_choices_N_1_K_S = choices_N_K_S[:, None, :, :]
        expanded_probs_N_K_1_C = probs_N_K_C[:, :, None, :]

        probs_N_K_K_S = gather_expand(expanded_probs_N_K_1_C,
                                      dim=-1,
                                      index=expanded_choices_N_1_K_S)
        # exp sum log seems necessary to avoid 0s?
        probs_K_K_S = torch.exp(
            torch.sum(torch.log(probs_N_K_K_S), dim=0, keepdim=False))
        samples_K_M = probs_K_K_S.reshape((K, -1))

        samples_M_K = samples_K_M.t()
        return SampledJointEntropy(samples_M_K)

    def compute(self) -> torch.Tensor:
        sampled_joint_probs_M = torch.mean(self.sampled_joint_probs_M_K,
                                           dim=1,
                                           keepdim=False)
        nats_M = -torch.log(sampled_joint_probs_M)
        entropy = torch.mean(nats_M)
        return entropy

    def add_variables(self, probs_N_K_C: torch.Tensor,
                      M2: int) -> 'SampledJointEntropy':
        assert self.sampled_joint_probs_M_K.shape[1] == probs_N_K_C.shape[1]

        sample_K_M1_1 = self.sampled_joint_probs_M_K.t()[:, :, None]

        new_sample_M2_K = self.sample(probs_N_K_C, M2).sampled_joint_probs_M_K
        new_sample_K_1_M2 = new_sample_M2_K.t()[:, None, :]

        merged_sample_K_M1_M2 = sample_K_M1_1 * new_sample_K_1_M2
        merged_sample_K_M = merged_sample_K_M1_M2.reshape((K, -1))

        self.sampled_joint_probs_M_K = merged_sample_K_M.t()

        return self

    def compute_batch(self,
                      probs_B_K_C: torch.Tensor,
                      output_entropies_B=None):
        assert self.sampled_joint_probs_M_K.shape[1] == probs_B_K_C.shape[1]

        B, K, C = probs_B_K_C.shape
        M = self.sampled_joint_probs_M_K.shape[0]

        if output_entropies_B is None:
            output_entropies_B = torch.empty(B,
                                             dtype=probs_B_K_C.dtype,
                                             device=probs_B_K_C.device)

        pbar = tqdm(total=B, desc="SampledJointEntropy.compute_batch", leave=False)

        @toma.execute.chunked(probs_B_K_C,
                              initial_step=1024,
                              dimension=0)
        def chunked_joint_entropy(chunked_probs_b_K_C: torch.Tensor,
                                  start: int, end: int):
            b = chunked_probs_b_K_C.shape[0]

            probs_b_M_C = torch.empty(
                (b, M, C),
                dtype=self.sampled_joint_probs_M_K.dtype,
                device=self.sampled_joint_probs_M_K.device)
            for i in range(b):
                torch.matmul(self.sampled_joint_probs_M_K,
                             probs_B_K_C[i].to(self.sampled_joint_probs_M_K,
                                               non_blocking=True),
                             out=probs_b_M_C[i])
            probs_b_M_C /= K

            q_1_M_1 = self.sampled_joint_probs_M_K.mean(dim=1,
                                                        keepdim=True)[None]

            output_entropies_B[start:end].copy_(
                torch.sum(-torch.log(probs_b_M_C) * probs_b_M_C / q_1_M_1,
                          dim=(1, 2)) / M,
                non_blocking=True)

            pbar.update(end - start)

        pbar.close()

        return output_entropies_B

# Cell

class DynamicJointEntropy(JointEntropy):
    inner: JointEntropy
    probs_max_N_K_C: torch.Tensor
    N: int
    M: int

    def __init__(self, M:int, max_N: int, K: int, C: int, dtype=None, device=None):
        self.M = M
        self.N = 0
        self.max_N = max_N

        self.inner = ExactJointEntropy.empty(K, dtype=dtype, device=device)
        self.probs_max_N_K_C = torch.empty((max_N, K, C), dtype=dtype, device=device)


    def add_variables(self, probs_N_K_C: torch.Tensor) -> 'DynamicJointEntropy':
        C = self.probs_max_N_K_C.shape[2]
        add_N = probs_N_K_C.shape[0]

        assert self.probs_max_N_K_C.shape[0] >= self.N + add_N
        assert self.probs_max_N_K_C.shape[2] == C

        self.probs_max_N_K_C[self.N:self.N+add_N] = probs_N_K_C
        self.N += add_N

        num_exact_samples = C**self.N
        if num_exact_samples > self.M:
            self.inner = SampledJointEntropy.sample(self.probs_max_N_K_C[:self.N], self.M)
        else:
            self.inner.add_variables(probs_N_K_C)

        return self

    def compute(self) -> torch.Tensor:
        return self.inner.compute()

    def compute_batch(self,
                      probs_B_K_C: torch.Tensor,
                      output_entropies_B=None) -> torch.Tensor:
        """Computes the joint entropy of the added variables together with the batch (one by one)."""
        return self.inner.compute_batch(probs_B_K_C, output_entropies_B)
