#  ~=~=~=~=~=~=~=~=~=~=~=~=~=~=~=~=~=~=~=~=~=~=~=~=~=~=~=~=~=~=~=~=~=~=~=~=~=~=~
#  MIT License
#
#  Copyright (c) 2021 Nathan Juraj Michlo
#
#  Permission is hereby granted, free of charge, to any person obtaining a copy
#  of this software and associated documentation files (the "Software"), to deal
#  in the Software without restriction, including without limitation the rights
#  to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
#  copies of the Software, and to permit persons to whom the Software is
#  furnished to do so, subject to the following conditions:
#
#  The above copyright notice and this permission notice shall be included in
#  all copies or substantial portions of the Software.
#
#  THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
#  IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
#  FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
#  AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
#  LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
#  OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
#  SOFTWARE.
#  ~=~=~=~=~=~=~=~=~=~=~=~=~=~=~=~=~=~=~=~=~=~=~=~=~=~=~=~=~=~=~=~=~=~=~=~=~=~=~

from dataclasses import dataclass
from typing import Any
from typing import Dict
from typing import Sequence
from typing import Tuple

import torch
from torch.distributions import Distribution
from torch.distributions import Normal
from torch.distributions import kl_divergence

from disent.frameworks.vae._unsupervised__betavae import BetaVae

# ========================================================================= #
# Ada-GVAE                                                                  #
# ========================================================================= #


class AdaVae(BetaVae):
    """
    Weakly Supervised Disentanglement Learning Without Compromises: https://arxiv.org/abs/2002.02886
    - pretty much a beta-vae with averaging between decoder outputs to form weak supervision signal.
    - GAdaVAE:   Averaging from https://arxiv.org/abs/1809.02383
    - ML-AdaVAE: Averaging from https://arxiv.org/abs/1705.08841

    MODIFICATION:
    - Symmetric KL Calculation used by default, described in: https://arxiv.org/pdf/2010.14407.pdf
    - adjustable threshold value

    * This class is a little over complicated because it has the added functionality
      listed above, and because we want to re-use features elsewhere. The code can
      be compressed down into about ~20 neat lines for `hook_intercept_ds` if we
      select and chose fixed `cfg` values.
    """

    REQUIRED_OBS = 2

    @dataclass
    class cfg(BetaVae.cfg):
        ada_average_mode: str = "gvae"
        ada_thresh_mode: str = "symmetric_kl"  # kl, symmetric_kl, dist, sampled_dist
        ada_thresh_ratio: float = 0.5

    def hook_intercept_ds(
        self, ds_posterior: Sequence[Distribution], ds_prior: Sequence[Distribution]
    ) -> Tuple[Sequence[Distribution], Sequence[Distribution], Dict[str, Any]]:
        """
        Adaptive VAE Glue Method, putting the various components together
        1. find differences between deltas
        2. estimate a threshold for differences
        3. compute a shared mask from this threshold
        4. average together elements that should be considered shared

        (✓) Visual inspection against reference implementation:
            https://github.com/google-research/disentanglement_lib (aggregate_argmax)
        """
        d0_posterior, d1_posterior = ds_posterior
        # shared elements that need to be averaged, computed per pair in the batch.
        share_mask = self.compute_shared_mask_from_posteriors(
            d0_posterior, d1_posterior, thresh_mode=self.cfg.ada_thresh_mode, ratio=self.cfg.ada_thresh_ratio
        )
        # compute average posteriors
        new_ds_posterior = self.make_shared_posteriors(
            d0_posterior, d1_posterior, share_mask, average_mode=self.cfg.ada_average_mode
        )
        # return new args & generate logs
        return new_ds_posterior, ds_prior, {"shared": share_mask.sum(dim=1).float().mean()}

    # - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - #
    # HELPER - POSTERIORS                                                   #
    # - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - #

    @classmethod
    def compute_deltas_from_posteriors(cls, d0_posterior: Distribution, d1_posterior: Distribution, thresh_mode: str):
        """
        (✓) Visual inspection against reference implementation
        https://github.com/google-research/disentanglement_lib (compute_kl)
        - difference is that they don't multiply by 0.5 to get true kl, but that's not needed

        TODO: this might be numerically unstable with f32 passed to distributions
        """
        # shared elements that need to be averaged, computed per pair in the batch.
        # [𝛿_i ...]
        if thresh_mode == "kl":
            # ORIGINAL
            deltas = kl_divergence(d1_posterior, d0_posterior)
        elif thresh_mode == "symmetric_kl":
            # FROM: https://openreview.net/pdf?id=8VXvj1QNRl1
            kl_deltas_d1_d0 = kl_divergence(d1_posterior, d0_posterior)
            kl_deltas_d0_d1 = kl_divergence(d0_posterior, d1_posterior)
            deltas = (0.5 * kl_deltas_d1_d0) + (0.5 * kl_deltas_d0_d1)
        elif thresh_mode == "dist":
            deltas = cls.compute_deltas_from_zs(d1_posterior.mean, d0_posterior.mean)
        elif thresh_mode == "sampled_dist":
            deltas = cls.compute_deltas_from_zs(d1_posterior.rsample(), d0_posterior.rsample())
        else:
            raise KeyError(f"invalid thresh_mode: {repr(thresh_mode)}")

        # return values
        return deltas

    @classmethod
    def compute_shared_mask_from_posteriors(
        cls, d0_posterior: Distribution, d1_posterior: Distribution, thresh_mode: str, ratio=0.5
    ):
        return cls.estimate_shared_mask(
            z_deltas=cls.compute_deltas_from_posteriors(d0_posterior, d1_posterior, thresh_mode=thresh_mode),
            ratio=ratio,
        )

    @classmethod
    def make_shared_posteriors(
        cls, d0_posterior: Normal, d1_posterior: Normal, share_mask: torch.Tensor, average_mode: str
    ) -> Tuple[Normal, Normal]:
        # compute average posterior
        ave_posterior = AdaVae.compute_average_distribution(
            d0_posterior=d0_posterior, d1_posterior=d1_posterior, average_mode=average_mode
        )
        # select shared elements
        ave_d0_posterior = Normal(
            loc=torch.where(share_mask, ave_posterior.loc, d0_posterior.loc),
            scale=torch.where(share_mask, ave_posterior.scale, d0_posterior.scale),
        )
        ave_d1_posterior = Normal(
            loc=torch.where(share_mask, ave_posterior.loc, d1_posterior.loc),
            scale=torch.where(share_mask, ave_posterior.scale, d1_posterior.scale),
        )
        # return values
        return ave_d0_posterior, ave_d1_posterior

    # - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - #
    # HELPER - MEAN/MU VALUES (same functionality as posterior versions)    #
    # - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - #

    @classmethod
    def compute_deltas_from_zs(cls, z0: torch.Tensor, z1: torch.Tensor) -> torch.Tensor:
        return torch.abs(z0 - z1)

    @classmethod
    def compute_shared_mask_from_zs(cls, z0: torch.Tensor, z1: torch.Tensor, ratio: float = 0.5):
        return cls.estimate_shared_mask(z_deltas=cls.compute_deltas_from_zs(z0, z1), ratio=ratio)

    @classmethod
    def make_shared_zs(
        cls, z0: torch.Tensor, z1: torch.Tensor, share_mask: torch.Tensor
    ) -> Tuple[torch.Tensor, torch.Tensor]:
        # compute average values
        ave = 0.5 * z0 + 0.5 * z1
        # select shared elements
        ave_z0 = torch.where(share_mask, ave, z0)
        ave_z1 = torch.where(share_mask, ave, z1)
        return ave_z0, ave_z1

    # - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - #
    # HELPER - COMMON                                                       #
    # - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - #

    @classmethod
    def estimate_shared_mask(cls, z_deltas: torch.Tensor, ratio: float) -> torch.Tensor:
        """
        Core of the adaptive VAE algorithm, estimating which factors
        have changed (or in this case which are shared and should remained unchanged
        by being be averaged) between pairs of observations.
        - custom ratio is an addition, when ratio==0.5 then
          this is equivalent to the original implementation.

        (✓) Visual inspection against reference implementation:
            https://github.com/google-research/disentanglement_lib (aggregate_argmax)
            - Implementation conversion is non-trivial, items are histogram binned.
              If we are in the second histogram bin, ie. 1, then kl_deltas <= kl_threshs
            - TODO: (aggregate_labels) An alternative mode exists where you can bind the
                    latent variables to any individual label, by one-hot encoding which
                    latent variable should not be shared: "enforce that each dimension
                    of the latent code learns one factor (dimension 1 learns factor 1)
                    and enforce that each factor of variation is encoded in a single
                    dimension."
        """
        assert 0 <= ratio <= 1, f"ratio must be in the range: 0 <= ratio <= 1, got: {repr(ratio)}"
        # threshold τ
        maximums = z_deltas.max(axis=1, keepdim=True).values  # (B, 1)
        minimums = z_deltas.min(axis=1, keepdim=True).values  # (B, 1)
        z_threshs = torch.lerp(minimums, maximums, weight=ratio)  # (B, 1)
        # true if 'unchanged' and should be average
        shared_mask = z_deltas < z_threshs  # broadcast (B, Z) and (B, 1) -> (B, Z)
        # return
        return shared_mask

    # - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - #
    # HELPER - DISTRIBUTIONS                                                #
    # - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - #

    @classmethod
    def compute_average_distribution(cls, d0_posterior: Normal, d1_posterior: Normal, average_mode: str) -> Normal:
        return _COMPUTE_AVE_FNS[average_mode](
            d0_posterior=d0_posterior,
            d1_posterior=d1_posterior,
        )


# ========================================================================= #
# Averaging Functions                                                       #
# ========================================================================= #


def compute_average_gvae_std(d0_posterior: Normal, d1_posterior: Normal) -> Normal:
    """
    Compute the arithmetic mean of the encoder distributions.
    - This is a custom function based on the Ada-GVAE averaging,
      except over the standard deviation instead of the variance!

    *NB* this is un-official!
    """
    assert isinstance(
        d0_posterior, Normal
    ), f"posterior distributions must be {Normal.__name__} distributions, got: {type(d0_posterior)}"
    assert isinstance(
        d1_posterior, Normal
    ), f"posterior distributions must be {Normal.__name__} distributions, got: {type(d1_posterior)}"
    # averages
    ave_std = 0.5 * (d0_posterior.stddev + d1_posterior.stddev)
    ave_mean = 0.5 * (d0_posterior.mean + d1_posterior.mean)
    # done!
    return Normal(loc=ave_mean, scale=ave_std)


def compute_average_gvae(d0_posterior: Normal, d1_posterior: Normal) -> Normal:
    """
    Compute the arithmetic mean of the encoder distributions.
    - Ada-GVAE Averaging function

    (✓) Visual inspection against reference implementation:
        https://github.com/google-research/disentanglement_lib (GroupVAEBase.model_fn)
    """
    assert isinstance(
        d0_posterior, Normal
    ), f"posterior distributions must be {Normal.__name__} distributions, got: {type(d0_posterior)}"
    assert isinstance(
        d1_posterior, Normal
    ), f"posterior distributions must be {Normal.__name__} distributions, got: {type(d1_posterior)}"
    # averages
    ave_var = 0.5 * (d0_posterior.variance + d1_posterior.variance)
    ave_mean = 0.5 * (d0_posterior.mean + d1_posterior.mean)
    # done!
    return Normal(loc=ave_mean, scale=torch.sqrt(ave_var))


def compute_average_ml_vae(d0_posterior: Normal, d1_posterior: Normal) -> Normal:
    """
    Compute the product of the encoder distributions.
    - Ada-ML-VAE Averaging function

    (✓) Visual inspection against reference implementation:
        https://github.com/google-research/disentanglement_lib (MLVae.model_fn)

    # TODO: recheck
    """
    assert isinstance(
        d0_posterior, Normal
    ), f"posterior distributions must be {Normal.__name__} distributions, got: {type(d0_posterior)}"
    assert isinstance(
        d1_posterior, Normal
    ), f"posterior distributions must be {Normal.__name__} distributions, got: {type(d1_posterior)}"
    # Diagonal matrix inverse: E^-1 = 1 / E
    # https://proofwiki.org/wiki/Inverse_of_Diagonal_Matrix
    z0_invvar, z1_invvar = d0_posterior.variance.reciprocal(), d1_posterior.variance.reciprocal()
    # average var: E^-1 = E1^-1 + E2^-1
    # disentanglement_lib: ave_var = 2 * z0_var * z1_var / (z0_var + z1_var)
    ave_var = 2 * (z0_invvar + z1_invvar).reciprocal()
    # average mean: u^T = (u1^T E1^-1 + u2^T E2^-1) E
    # disentanglement_lib: ave_mean = (z0_mean/z0_var + z1_mean/z1_var) * ave_var * 0.5
    ave_mean = (d0_posterior.mean * z0_invvar + d1_posterior.mean * z1_invvar) * ave_var * 0.5
    # done!
    return Normal(loc=ave_mean, scale=torch.sqrt(ave_var))


_COMPUTE_AVE_FNS = {
    "gvae": compute_average_gvae,
    "ml-vae": compute_average_ml_vae,
    "gvae_std": compute_average_gvae_std,  # this is un-official!
}


# ========================================================================= #
# Ada-GVAE                                                                  #
# ========================================================================= #


class AdaGVaeMinimal(BetaVae):
    """
    This is a direct implementation of the Ada-GVAE,
    which should be equivalent to the AdaVae with config values:

    >>> AdaVae.cfg(
    >>>    ada_average_mode='gvae',
    >>>    ada_thresh_mode='symmetric_kl',
    >>>    ada_thresh_ratio=0.5,
    >>> )
    """

    REQUIRED_OBS = 2

    def hook_intercept_ds(
        self, ds_posterior: Sequence[Distribution], ds_prior: Sequence[Distribution]
    ) -> Tuple[Sequence[Distribution], Sequence[Distribution], Dict[str, Any]]:
        """
        Adaptive VAE Method, putting the various components together
            1. compute differences between representations
            2. estimate a threshold for differences
            3. compute a shared mask from this threshold
            4. average together elements that are marked as shared

        (x) Visual inspection against reference implementation:
            https://github.com/google-research/disentanglement_lib (aggregate_argmax)
        """
        d0_posterior, d1_posterior = ds_posterior
        assert isinstance(
            d0_posterior, Normal
        ), f"posterior distributions must be {Normal.__name__} distributions, got: {type(d0_posterior)}"
        assert isinstance(
            d1_posterior, Normal
        ), f"posterior distributions must be {Normal.__name__} distributions, got: {type(d1_posterior)}"

        # [1] symmetric KL Divergence FROM: https://openreview.net/pdf?id=8VXvj1QNRl1
        z_deltas = 0.5 * kl_divergence(d1_posterior, d0_posterior) + 0.5 * kl_divergence(d0_posterior, d1_posterior)

        # [2] estimate threshold from deltas
        z_deltas_min = z_deltas.min(axis=1, keepdim=True).values  # (B, 1)
        z_deltas_max = z_deltas.max(axis=1, keepdim=True).values  # (B, 1)
        z_thresh = 0.5 * z_deltas_min + 0.5 * z_deltas_max  # (B, 1)

        # [3] shared elements that need to be averaged, computed per pair in the batch
        share_mask = z_deltas < z_thresh  # broadcast (B, Z) and (B, 1) to get (B, Z)

        # [4.a] compute average representations
        # - this is the only difference between the Ada-ML-VAE
        ave_mean = 0.5 * d0_posterior.mean + 0.5 * d1_posterior.mean
        ave_std = (0.5 * d0_posterior.variance + 0.5 * d1_posterior.variance) ** 0.5

        # [4.b] select shared or original values based on mask
        z0_mean = torch.where(share_mask, ave_mean, d0_posterior.loc)
        z1_mean = torch.where(share_mask, ave_mean, d1_posterior.loc)
        z0_std = torch.where(share_mask, ave_std, d0_posterior.scale)
        z1_std = torch.where(share_mask, ave_std, d1_posterior.scale)

        # construct distributions
        ave_d0_posterior = Normal(loc=z0_mean, scale=z0_std)
        ave_d1_posterior = Normal(loc=z1_mean, scale=z1_std)
        new_ds_posterior = (ave_d0_posterior, ave_d1_posterior)

        # [done] return new args & generate logs
        return new_ds_posterior, ds_prior, {"shared": share_mask.sum(dim=1).float().mean()}


# ========================================================================= #
# END                                                                       #
# ========================================================================= #
