"""PyTorch module for Contrastive VI for single cell expression data."""

import numpy as np
import torch
import torch.nn.functional as F
from torch.distributions import Normal
from torch.distributions import kl_divergence as kl

from scvi import REGISTRY_KEYS
from scvi.distributions import ZeroInflatedNegativeBinomial
from scvi.module.base import BaseModuleClass, LossOutput, auto_move_data
from scvi.nn import DecoderSCVI, Encoder

torch.backends.cudnn.benchmark = True


class ContrastiveVAE(BaseModuleClass):
    """Variational inference for contrastive analysis of RNA-seq data.

    Implements the contrastiveVI model of :cite:p:`Weinberger23`.

    Parameters
    ----------
    n_input
        Number of input genes.
    n_batch
        Number of batches. If 0, no batch effect correction is performed.
    n_hidden
        Number of nodes per hidden layer.
    n_background_latent
        Dimensionality of the background latent space.
    n_salient_latent
        Dimensionality of the salient latent space.
    n_layers
        Number of hidden layers used for encoder and decoder NNs.
    dropout_rate
        Dropout rate for neural networks.
    use_observed_lib_size
        Use observed library size for RNA as scaling factor in mean of conditional
        distribution.
    library_log_means
        1 x n_batch array of means of the log library sizes. Parameterize prior on
        library size if not using observed library size.
    library_log_vars
        1 x n_batch array of variances of the log library sizes. Parameterize prior
        on library size if not using observed library size.
    wasserstein_penalty
        Weight of the Wasserstein distance loss that further discourages shared
        variations from leaking into the salient latent space.
    """

    def __init__(
        self,
        n_input: int,
        n_batch: int = 0,
        n_hidden: int = 128,
        n_background_latent: int = 10,
        n_salient_latent: int = 10,
        n_layers: int = 1,
        dropout_rate: float = 0.1,
        use_observed_lib_size: bool = True,
        library_log_means: np.ndarray | None = None,
        library_log_vars: np.ndarray | None = None,
        wasserstein_penalty: float = 0,
    ) -> None:
        super().__init__()
        self.n_input = n_input
        self.n_batch = n_batch
        self.n_hidden = n_hidden
        self.n_background_latent = n_background_latent
        self.n_salient_latent = n_salient_latent
        self.n_layers = n_layers
        self.dropout_rate = dropout_rate
        self.latent_distribution = "normal"
        self.dispersion = "gene"
        self.px_r = torch.nn.Parameter(torch.randn(n_input))
        self.use_observed_lib_size = use_observed_lib_size
        self.wasserstein_penalty = wasserstein_penalty

        if not self.use_observed_lib_size:
            if library_log_means is None or library_log_vars is None:
                raise ValueError(
                    "If not using observed_lib_size, "
                    "must provide library_log_means and library_log_vars."
                )
            self.register_buffer("library_log_means", torch.from_numpy(library_log_means).float())
            self.register_buffer("library_log_vars", torch.from_numpy(library_log_vars).float())

        cat_list = [n_batch]
        # Background encoder.
        self.z_encoder = Encoder(
            n_input,
            n_background_latent,
            n_cat_list=cat_list,
            n_layers=n_layers,
            n_hidden=n_hidden,
            dropout_rate=dropout_rate,
            distribution=self.latent_distribution,
            inject_covariates=True,
            use_batch_norm=True,
            use_layer_norm=False,
            var_activation=None,
        )
        # Salient encoder.
        self.s_encoder = Encoder(
            n_input,
            n_salient_latent,
            n_cat_list=cat_list,
            n_layers=n_layers,
            n_hidden=n_hidden,
            dropout_rate=dropout_rate,
            distribution=self.latent_distribution,
            inject_covariates=True,
            use_batch_norm=True,
            use_layer_norm=False,
            var_activation=None,
        )
        # Library size encoder.
        self.l_encoder = Encoder(
            n_input,
            1,
            n_layers=1,
            n_cat_list=cat_list,
            n_hidden=n_hidden,
            dropout_rate=dropout_rate,
            inject_covariates=True,
            use_batch_norm=True,
            use_layer_norm=False,
            var_activation=None,
        )
        # Decoder from latent variable to distribution parameters in data space.
        n_total_latent = n_background_latent + n_salient_latent
        self.decoder = DecoderSCVI(
            n_total_latent,
            n_input,
            n_cat_list=cat_list,
            n_layers=n_layers,
            n_hidden=n_hidden,
            inject_covariates=True,
            use_batch_norm=True,
            use_layer_norm=False,
        )

    @auto_move_data
    def _compute_local_library_params(
        self, batch_index: torch.Tensor
    ) -> tuple[torch.Tensor, torch.Tensor]:
        """Computes local library parameters.

        Compute two tensors of shape (batch_index.shape[0], 1) where each
        element corresponds to the mean and variances, respectively, of the
        log library sizes in the batch the cell corresponds to.
        """
        n_batch = self.library_log_means.shape[1]
        local_library_log_means = F.linear(
            F.one_hot(batch_index.squeeze(-1), n_batch).float(), self.library_log_means
        )
        local_library_log_vars = F.linear(
            F.one_hot(batch_index.squeeze(-1), n_batch).float(), self.library_log_vars
        )
        return local_library_log_means, local_library_log_vars

    @staticmethod
    def _get_min_batch_size(concat_tensors: dict[str, dict[str, torch.Tensor]]) -> int:
        return min(
            concat_tensors["background"][REGISTRY_KEYS.X_KEY].shape[0],
            concat_tensors["target"][REGISTRY_KEYS.X_KEY].shape[0],
        )

    @staticmethod
    def _reduce_tensors_to_min_batch_size(
        tensors: dict[str, torch.Tensor], min_batch_size: int
    ) -> None:
        for name, tensor in tensors.items():
            tensors[name] = tensor[:min_batch_size, :]

    @staticmethod
    def _get_inference_input_from_concat_tensors(
        concat_tensors: dict[str, dict[str, torch.Tensor]], index: str
    ) -> dict[str, torch.Tensor]:
        tensors = concat_tensors[index]
        x = tensors[REGISTRY_KEYS.X_KEY]
        batch_index = tensors[REGISTRY_KEYS.BATCH_KEY]
        input_dict = {"x": x, "batch_index": batch_index}
        return input_dict

    def _get_inference_input(
        self, concat_tensors: dict[str, dict[str, torch.Tensor]]
    ) -> dict[str, dict[str, torch.Tensor]]:
        background = self._get_inference_input_from_concat_tensors(concat_tensors, "background")
        target = self._get_inference_input_from_concat_tensors(concat_tensors, "target")
        # Ensure batch sizes are the same.
        min_batch_size = self._get_min_batch_size(concat_tensors)
        self._reduce_tensors_to_min_batch_size(background, min_batch_size)
        self._reduce_tensors_to_min_batch_size(target, min_batch_size)
        return {"background": background, "target": target}

    @staticmethod
    def _get_generative_input_from_concat_tensors(
        concat_tensors: dict[str, dict[str, torch.Tensor]], index: str
    ) -> dict[str, torch.Tensor]:
        tensors = concat_tensors[index]
        batch_index = tensors[REGISTRY_KEYS.BATCH_KEY]
        input_dict = {"batch_index": batch_index}
        return input_dict

    @staticmethod
    def _get_generative_input_from_inference_outputs(
        inference_outputs: dict[str, dict[str, torch.Tensor]], data_source: str
    ) -> dict[str, torch.Tensor]:
        z = inference_outputs[data_source]["z"]
        s = inference_outputs[data_source]["s"]
        library = inference_outputs[data_source]["library"]
        return {"z": z, "s": s, "library": library}

    def _get_generative_input(
        self,
        concat_tensors: dict[str, dict[str, torch.Tensor]],
        inference_outputs: dict[str, dict[str, torch.Tensor]],
    ) -> dict[str, dict[str, torch.Tensor]]:
        background_tensor_input = self._get_generative_input_from_concat_tensors(
            concat_tensors, "background"
        )
        target_tensor_input = self._get_generative_input_from_concat_tensors(
            concat_tensors, "target"
        )
        # Ensure batch sizes are the same.
        min_batch_size = self._get_min_batch_size(concat_tensors)
        self._reduce_tensors_to_min_batch_size(background_tensor_input, min_batch_size)
        self._reduce_tensors_to_min_batch_size(target_tensor_input, min_batch_size)

        background_inference_outputs = self._get_generative_input_from_inference_outputs(
            inference_outputs, "background"
        )
        target_inference_outputs = self._get_generative_input_from_inference_outputs(
            inference_outputs, "target"
        )
        background = {**background_tensor_input, **background_inference_outputs}
        target = {**target_tensor_input, **target_inference_outputs}
        return {"background": background, "target": target}

    @staticmethod
    def _reshape_tensor_for_samples(tensor: torch.Tensor, n_samples: int):
        return tensor.unsqueeze(0).expand((n_samples, tensor.size(0), tensor.size(1)))

    @auto_move_data
    def _generic_inference(
        self,
        x: torch.Tensor,
        batch_index: torch.Tensor,
        n_samples: int = 1,
    ) -> dict[str, torch.Tensor]:
        x_ = x
        if self.use_observed_lib_size:
            library = torch.log(x.sum(1)).unsqueeze(1)
        x_ = torch.log(1 + x_)

        qz_m, qz_v, z = self.z_encoder(x_, batch_index)
        qs_m, qs_v, s = self.s_encoder(x_, batch_index)

        ql_m, ql_v = None, None
        if not self.use_observed_lib_size:
            ql_m, ql_v, library_encoded = self.l_encoder(x_, batch_index)
            library = library_encoded

        if n_samples > 1:
            qz_m = self._reshape_tensor_for_samples(qz_m, n_samples)
            qz_v = self._reshape_tensor_for_samples(qz_v, n_samples)
            z = self._reshape_tensor_for_samples(z, n_samples)
            qs_m = self._reshape_tensor_for_samples(qs_m, n_samples)
            qs_v = self._reshape_tensor_for_samples(qs_v, n_samples)
            s = self._reshape_tensor_for_samples(s, n_samples)

            if self.use_observed_lib_size:
                library = self._reshape_tensor_for_samples(library, n_samples)
            else:
                ql_m = self._reshape_tensor_for_samples(ql_m, n_samples)
                ql_v = self._reshape_tensor_for_samples(ql_v, n_samples)
                library = Normal(ql_m, ql_v.sqrt()).sample()

        outputs = {
            "z": z,
            "qz_m": qz_m,
            "qz_v": qz_v,
            "s": s,
            "qs_m": qs_m,
            "qs_v": qs_v,
            "library": library,
            "ql_m": ql_m,
            "ql_v": ql_v,
        }
        return outputs

    @auto_move_data
    def inference(
        self,
        background: dict[str, torch.Tensor],
        target: dict[str, torch.Tensor],
        n_samples: int = 1,
    ) -> dict[str, dict[str, torch.Tensor]]:
        background_batch_size = background["x"].shape[0]
        target_batch_size = target["x"].shape[0]
        inference_input = {}
        for key in background.keys():
            inference_input[key] = torch.cat([background[key], target[key]], dim=0)
        outputs = self._generic_inference(**inference_input, n_samples=n_samples)
        batch_size_dim = 0 if n_samples == 1 else 1
        background_outputs, target_outputs = {}, {}
        for key in outputs.keys():
            if outputs[key] is not None:
                background_tensor, target_tensor = torch.split(
                    outputs[key],
                    [background_batch_size, target_batch_size],
                    dim=batch_size_dim,
                )
            else:
                background_tensor, target_tensor = None, None
            background_outputs[key] = background_tensor
            target_outputs[key] = target_tensor
        background_outputs["s"] = torch.zeros_like(background_outputs["s"])
        return {"background": background_outputs, "target": target_outputs}

    @auto_move_data
    def _generic_generative(
        self,
        z: torch.Tensor,
        s: torch.Tensor,
        library: torch.Tensor,
        batch_index: torch.Tensor,
    ) -> dict[str, torch.Tensor]:
        latent = torch.cat([z, s], dim=-1)
        px_scale, px_r, px_rate, px_dropout = self.decoder(
            self.dispersion,
            latent,
            library,
            batch_index,
        )
        px_r = torch.exp(self.px_r)
        return {
            "px_scale": px_scale,
            "px_r": px_r,
            "px_rate": px_rate,
            "px_dropout": px_dropout,
        }

    @auto_move_data
    def generative(
        self,
        background: dict[str, torch.Tensor],
        target: dict[str, torch.Tensor],
    ) -> dict[str, dict[str, torch.Tensor]]:
        latent_z_shape = background["z"].shape
        batch_size_dim = 0 if len(latent_z_shape) == 2 else 1
        background_batch_size = background["z"].shape[batch_size_dim]
        target_batch_size = target["z"].shape[batch_size_dim]
        generative_input = {}
        for key in ["z", "s", "library"]:
            generative_input[key] = torch.cat([background[key], target[key]], dim=batch_size_dim)
        generative_input["batch_index"] = torch.cat(
            [background["batch_index"], target["batch_index"]], dim=0
        )
        outputs = self._generic_generative(**generative_input)
        background_outputs, target_outputs = {}, {}
        for key in ["px_scale", "px_rate", "px_dropout"]:
            if outputs[key] is not None:
                background_tensor, target_tensor = torch.split(
                    outputs[key],
                    [background_batch_size, target_batch_size],
                    dim=batch_size_dim,
                )
            else:
                background_tensor, target_tensor = None, None
            background_outputs[key] = background_tensor
            target_outputs[key] = target_tensor
        background_outputs["px_r"] = outputs["px_r"]
        target_outputs["px_r"] = outputs["px_r"]
        return {"background": background_outputs, "target": target_outputs}

    @staticmethod
    def reconstruction_loss(
        x: torch.Tensor,
        px_rate: torch.Tensor,
        px_r: torch.Tensor,
        px_dropout: torch.Tensor,
    ) -> torch.Tensor:
        """Computes likelihood loss for zero-inflated negative binomial distribution.

        Parameters
        ----------
        x
            Input data.
        px_rate
            Mean of distribution.
        px_r
            Inverse dispersion.
        px_dropout
            Logits scale of zero inflation probability.

        Returns
        -------
        Negative log likelihood (reconstruction loss) for each data point. If number of
        latent samples == 1, the tensor has shape (batch_size, ). If number of latent
        samples > 1, the tensor has shape (n_samples, batch_size).
        """
        recon_loss = (
            -ZeroInflatedNegativeBinomial(mu=px_rate, theta=px_r, zi_logits=px_dropout)
            .log_prob(x)
            .sum(dim=-1)
        )
        return recon_loss

    @staticmethod
    def latent_kl_divergence(
        variational_mean: torch.Tensor,
        variational_var: torch.Tensor,
        prior_mean: torch.Tensor,
        prior_var: torch.Tensor,
    ) -> torch.Tensor:
        """Computes KL divergence between a variational posterior and prior Gaussian.

        Parameters
        ----------
        variational_mean
            Mean of the variational posterior Gaussian.
        variational_var
            Variance of the variational posterior Gaussian.
        prior_mean
            Mean of the prior Gaussian.
        prior_var
            Variance of the prior Gaussian.

        Returns
        -------
        KL divergence for each data point. If number of latent samples == 1, the tensor
        has shape (batch_size, ). If number of latent samples > 1, the tensor has shape
        (n_samples, batch_size).
        """
        return kl(
            Normal(variational_mean, variational_var.sqrt()),
            Normal(prior_mean, prior_var.sqrt()),
        ).sum(dim=-1)

    def library_kl_divergence(
        self,
        batch_index: torch.Tensor,
        variational_library_mean: torch.Tensor,
        variational_library_var: torch.Tensor,
        library: torch.Tensor,
    ) -> torch.Tensor:
        """Computes KL divergence between library size variational posterior and prior.

        Both the variational posterior and prior are Log-Normal.

        Parameters
        ----------
        batch_index
            Batch indices for batch-specific library size mean and variance.
        variational_library_mean
            Mean of variational Log-Normal.
        variational_library_var
            Variance of variational Log-Normal.
        library
            Sampled library size.

        Returns
        -------
        KL divergence for each data point. If number of latent samples == 1, the tensor
        has shape (batch_size, ). If number of latent samples > 1, the tensor has shape
        (n_samples, batch_size).
        """
        if not self.use_observed_lib_size:
            (
                local_library_log_means,
                local_library_log_vars,
            ) = self._compute_local_library_params(batch_index)

            kl_library = kl(
                Normal(variational_library_mean, variational_library_var.sqrt()),
                Normal(local_library_log_means, local_library_log_vars.sqrt()),
            )
        else:
            kl_library = torch.zeros_like(library)
        return kl_library.sum(dim=-1)

    def _generic_loss(
        self,
        tensors: dict[str, torch.Tensor],
        inference_outputs: dict[str, torch.Tensor],
        generative_outputs: dict[str, torch.Tensor],
    ) -> dict[str, torch.Tensor]:
        x = tensors[REGISTRY_KEYS.X_KEY]
        batch_index = tensors[REGISTRY_KEYS.BATCH_KEY]

        qz_m = inference_outputs["qz_m"]
        qz_v = inference_outputs["qz_v"]
        qs_m = inference_outputs["qs_m"]
        qs_v = inference_outputs["qs_v"]
        library = inference_outputs["library"]
        ql_m = inference_outputs["ql_m"]
        ql_v = inference_outputs["ql_v"]
        px_rate = generative_outputs["px_rate"]
        px_r = generative_outputs["px_r"]
        px_dropout = generative_outputs["px_dropout"]

        prior_z_m = torch.zeros_like(qz_m)
        prior_z_v = torch.ones_like(qz_v)
        prior_s_m = torch.zeros_like(qs_m)
        prior_s_v = torch.ones_like(qs_v)

        recon_loss = self.reconstruction_loss(x, px_rate, px_r, px_dropout)
        kl_z = self.latent_kl_divergence(qz_m, qz_v, prior_z_m, prior_z_v)
        kl_s = self.latent_kl_divergence(qs_m, qs_v, prior_s_m, prior_s_v)
        kl_library = self.library_kl_divergence(batch_index, ql_m, ql_v, library)
        return {
            "recon_loss": recon_loss,
            "kl_z": kl_z,
            "kl_s": kl_s,
            "kl_library": kl_library,
        }

    def loss(
        self,
        concat_tensors: dict[str, dict[str, torch.Tensor]],
        inference_outputs: dict[str, dict[str, torch.Tensor]],
        generative_outputs: dict[str, dict[str, torch.Tensor]],
        kl_weight: float = 1.0,
    ) -> LossOutput:
        """Computes loss terms for contrastiveVI.

        Parameters
        ----------
        concat_tensors
            Tuple of data mini-batch. The first element contains background data
            mini-batch. The second element contains target data mini-batch.
        inference_outputs
            Dictionary of inference step outputs. The keys are "background" and
            "target" for the corresponding outputs.
        generative_outputs
            Dictionary of generative step outputs. The keys are "background" and
            "target" for the corresponding outputs.
        kl_weight
            Importance weight for KL divergence of background and salient latent
            variables, relative to KL divergence of library size.

        Returns
        -------
        An scvi.module.base.LossOutput instance that records the following:
        loss
            One-dimensional tensor for overall loss used for optimization.
        reconstruction_loss
            Reconstruction loss with shape (n_samples, batch_size) if number of latent
            samples > 1, or (batch_size, ) if number of latent samples == 1.
        kl_local
            KL divergence term with shape (n_samples, batch_size) if number of latent
            samples > 1, or (batch_size, ) if number of latent samples == 1.
        """
        background_tensors = concat_tensors["background"]
        target_tensors = concat_tensors["target"]
        # Ensure batch sizes are the same.
        min_batch_size = self._get_min_batch_size(concat_tensors)
        self._reduce_tensors_to_min_batch_size(background_tensors, min_batch_size)
        self._reduce_tensors_to_min_batch_size(target_tensors, min_batch_size)

        background_losses = self._generic_loss(
            background_tensors,
            inference_outputs["background"],
            generative_outputs["background"],
        )
        target_losses = self._generic_loss(
            target_tensors,
            inference_outputs["target"],
            generative_outputs["target"],
        )
        reconst_loss = background_losses["recon_loss"] + target_losses["recon_loss"]
        kl_divergence_z = background_losses["kl_z"] + target_losses["kl_z"]
        kl_divergence_s = target_losses["kl_s"]
        kl_divergence_l = background_losses["kl_library"] + target_losses["kl_library"]

        wasserstein_loss = torch.norm(
            inference_outputs["background"]["qs_m"], dim=-1
        ) ** 2 + torch.sum(inference_outputs["background"]["qs_v"], dim=-1)

        kl_local_for_warmup = kl_divergence_z + kl_divergence_s
        kl_local_no_warmup = kl_divergence_l

        weighted_kl_local = (
            kl_weight * (self.wasserstein_penalty * wasserstein_loss + kl_local_for_warmup)
            + kl_local_no_warmup
        )

        loss = torch.mean(reconst_loss + weighted_kl_local)

        kl_local = {
            "kl_divergence_l": kl_divergence_l,
            "kl_divergence_z": kl_divergence_z,
            "kl_divergence_s": kl_divergence_s,
        }

        # LossOutput internally sums the reconst_loss, kl_local, and kl_global
        # terms before logging, so we do the same for the `wasserstein_loss` term.
        return LossOutput(
            loss=loss,
            reconstruction_loss=reconst_loss,
            kl_local=kl_local,
            extra_metrics={"wasserstein_loss_sum": torch.sum(wasserstein_loss)},
        )

    def sample(self):
        raise NotImplementedError("No sampling method for contrastiveVI")
