# AUTOGENERATED! DO NOT EDIT! File to edit: ../../nbs/layers/04_latent_space.ipynb.

# %% auto 0
__all__ = ['LatentEncoderConfig', 'LatentDecoderConfig', 'GaussianLatentSpaceConfig', 'LatentEncoder3D', 'LatentDecoder3D',
           'GaussianLatentSpace']

# %% ../../nbs/layers/04_latent_space.ipynb 2
from typing import Literal

import torch
from torch import nn
from torch.distributions import Normal, kl_divergence

from ..blocks.cnn import CNNBlock3D, CNNBlockConfig
from ..docstrings import populate_docstring
from ..utils.clamping import symmetric_tanh_clamp
from ..utils.custom_base_model import CustomBaseModel, Field, model_validator
from ..utils.rearrange import rearrange_channels

# %% ../../nbs/layers/04_latent_space.ipynb 4
class LatentEncoderConfig(CNNBlockConfig):
    init_low_var: bool = Field(False, description="Whether to initialize weights such that output variance is low")

    @model_validator(mode="before")
    @classmethod
    def validate_before(cls, data: dict):
        data.setdefault("in_channels", data.get("dim"))
        data.setdefault("out_channels", data.get("latent_dim"))
        return data

    @property
    def dim(self) -> int:
        return self.in_channels

    @property
    def latent_dim(self) -> int:
        return self.out_channels


class LatentDecoderConfig(CNNBlockConfig):
    @model_validator(mode="before")
    @classmethod
    def validate_before(cls, data: dict):
        data.setdefault("in_channels", data.get("latent_dim"))
        data.setdefault("out_channels", data.get("dim"))
        return data

    @property
    def latent_dim(self):
        return self.in_channels

    @property
    def dim(self):
        return self.out_channels


class GaussianLatentSpaceConfig(CustomBaseModel):
    pass

# %% ../../nbs/layers/04_latent_space.ipynb 6
class LatentEncoder3D(nn.Module):
    """Encodes input features into a latent space representation by predicting the mean and standard deviation of the
    latent space."""

    @populate_docstring
    def __init__(self, config: LatentEncoderConfig = {}, checkpointing_level: int = 0, **kwargs):
        """Initializes the LatentEncoder3D.

        Args:
            config: {CONFIG_INSTANCE_DOC}
            checkpointing_level: {CHECKPOINTING_LEVEL_DOC}
            **kwargs: {CONFIG_KWARGS_DOC}
        """
        super().__init__()

        self.config = LatentEncoderConfig.model_validate(config | kwargs)

        latent_dim = self.config.latent_dim

        self.dim_mapper = CNNBlock3D(self.config, checkpointing_level)
        self.quant_conv_mu = nn.Conv3d(latent_dim, latent_dim, 1)
        self.quant_conv_log_var = nn.Conv3d(latent_dim, latent_dim, 1)

        if self.config.init_low_var:
            self.init_low_var()

    def init_low_var(self, bias_constant: float = -1.0):
        nn.init.normal_(self.quant_conv_log_var.weight, std=0.001)
        nn.init.constant_(self.quant_conv_log_var.bias, bias_constant)

    @populate_docstring
    def forward(
        self,
        x: torch.Tensor,
        prior_mu: torch.Tensor | None = None,
        prior_log_var: torch.Tensor | None = None,
        return_log_var: bool = False,
        max_mu: float = 100.0,
        max_log_var: float = 10.0,
        channels_first: bool = True,
    ) -> tuple[torch.Tensor, torch.Tensor]:
        """Get latent space representation of the input by mapping it to the latent dimension and then extracting the
        mean and standard deviation of the latent space. If a prior distribution is provided, the input is expected to
        predict the deviation from the prior. If it is not provided, one can think of it as the deviation from a
        standard normal distribution is being predicted. The output is the mean and standard deviation of the latent
        space.

        Args:
            x: {INPUT_3D_DOC}
            prior_mu: The mean of the prior distribution. If None, it is assumed to be the mean of a standard normal
                distribution. {INPUT_3D_DOC}
            prior_log_var: The log-variance of the prior distribution. If None, it is assumed to be log-variance of a
                standard normal distribution. {INPUT_3D_DOC}
            return_log_var: Whether to return the log-variance too.
            max_mu: Clamps sigma to the minimum and maximum values allowed i.e. to the range `[-max_mu,
                max_mu]`.
            max_log_var: Clamps log-variance to the minimum and maximum values allowed i.e. to the range `[-max_log_var,
                max_log_var]`. Defaults to 10.0, which corresponds to a variance from 0.000045 (std=0.006737) to
                22026.465 (std=148.413).
            channels_first: {CHANNELS_FIRST_DOC}

        Returns:
            z_mu: The mean of the latent space.
            z_sigma: The standard deviation of the latent space.
        """
        # x: (b, [dim], z, y, x, [dim])
        # z_mu_prior: (b, [latent_dim], z, y, x, [latent_dim])
        # z_log_var_prior: (b, [latent_dim], z, y, x, [latent_dim])

        assert (
            prior_mu is None and prior_log_var is None or (prior_mu is not None and prior_log_var is not None)
        ), "If prior_mu or prior_log_var are provided, both must be provided."

        x = rearrange_channels(x, channels_first, True)
        # (b, dim, z, y, x)
        if prior_mu is not None:
            prior_mu = rearrange_channels(prior_mu, channels_first, True)
            prior_log_var = rearrange_channels(prior_log_var, channels_first, True)
            # (b, latent_dim, z, y, x)

        x = self.dim_mapper(x, channels_first=True)
        # (b, latent_dim, z, y, x)

        z_mu = self.quant_conv_mu(x)
        z_log_var = self.quant_conv_log_var(x)
        # (b, latent_dim, z, y, x)

        if prior_mu is not None:
            z_mu = prior_mu + z_mu
            z_log_var = prior_log_var + z_log_var
            # (b, latent_dim, z, y, x)

        z_mu = symmetric_tanh_clamp(z_mu, max_mu)
        z_log_var = symmetric_tanh_clamp(z_log_var, max_log_var)
        z_sigma = torch.exp(z_log_var / 2)
        # (b, latent_dim, z, y, x)

        z_mu = rearrange_channels(z_mu, True, channels_first)
        z_sigma = rearrange_channels(z_sigma, True, channels_first)
        if return_log_var:
            z_log_var = rearrange_channels(z_log_var, True, channels_first)
        # (b, [latent_dim], z, y, x, [latent_dim])

        if return_log_var:
            return z_mu, z_sigma, z_log_var

        return z_mu, z_sigma

# %% ../../nbs/layers/04_latent_space.ipynb 9
class LatentDecoder3D(nn.Module):
    """Decodes latent space representations back to the original feature space."""

    @populate_docstring
    def __init__(self, config: LatentDecoderConfig = {}, checkpointing_level: int = 0, **kwargs):
        """Initializes the LatentDecoder3D.

        Args:
            config: {CONFIG_INSTANCE_DOC}
            checkpointing_level: {CHECKPOINTING_LEVEL_DOC}
            **kwargs: {CONFIG_KWARGS_DOC}
        """
        super().__init__()

        self.config = LatentDecoderConfig.model_validate(config | kwargs)

        latent_dim = self.config.latent_dim

        self.post_quant_conv = nn.Conv3d(latent_dim, latent_dim, 1)
        self.dim_mapper = CNNBlock3D(self.config, checkpointing_level)

    @populate_docstring
    def forward(self, z: torch.Tensor, channels_first: bool = True) -> torch.Tensor:
        """Decode latent space representation back to the original feature space.

        Args:
            z: Sampled vector in the latent space. {INPUT_3D_DOC}
            channels_first: {CHANNELS_FIRST_DOC}

        Returns:
            x: The decoded feature tensor. {OUTPUT_3D_DOC}
        """
        # z: (b, [latent_dim], z, y, x, [latent_dim])

        z = rearrange_channels(z, channels_first, True)
        # (b, latent_dim, z, y, x)

        z = self.post_quant_conv(z)
        # (b, latent_dim, z, y, x)

        x = self.dim_mapper(z, channels_first=True)
        # (b, dim, z, y, x)

        x = rearrange_channels(x, True, channels_first)
        # (b, [dim], z, y, x, [dim])

        return x

# %% ../../nbs/layers/04_latent_space.ipynb 11
class GaussianLatentSpace(nn.Module):
    """Implements a Gaussian latent space with sampling and KL divergence computation."""

    @populate_docstring
    def __init__(self, config: GaussianLatentSpaceConfig = {}, **kwargs):
        """Initializes the GaussianLatentSpace.

        Args:
            config: {CONFIG_INSTANCE_DOC}
            **kwargs: {CONFIG_KWARGS_DOC}
        """
        super().__init__()

        self.config = GaussianLatentSpaceConfig.model_validate(config | kwargs)

    @populate_docstring
    def sample(self, z_mu: torch.Tensor, z_sigma: torch.Tensor, force_sampling: bool = False) -> torch.Tensor:
        """Samples from the latent space using the reparameterization trick.

        Args:
            z_mu: The mean of the latent space. {INPUT_3D_DOC}
            z_sigma: The standard deviation of the latent space. {INPUT_3D_DOC}
            force_sampling: Whether to force sampling even in evaluation mode.

        Returns:
            z_vae: The sampled latent vector. {OUTPUT_3D_DOC}
        """
        if not self.training and not force_sampling:
            return z_mu
        eps = torch.randn_like(z_sigma)
        z_vae = z_mu + eps * z_sigma
        return z_vae

    @staticmethod
    @populate_docstring
    def kl_divergence(
        z_mu: torch.Tensor,
        z_sigma: torch.Tensor,
        prior_mu: torch.Tensor | None = None,
        prior_sigma: torch.Tensor | None = None,
        reduction: Literal["allsum", "channelssum"] | None = "allsum",
        channels_first: bool = True,
    ) -> torch.Tensor:
        """Computes the KL divergence between the latent space distribution and a prior distribution.

        Args:
            z_mu: The mean of the latent space. {INPUT_3D_DOC}
            z_sigma: The standard deviation of the latent space. {INPUT_3D_DOC}
            prior_mu: The mean of the prior distribution. If None, it is assumed to be the mean of a standard normal
                distribution. {INPUT_3D_DOC}
            prior_sigma: The standard deviation of the prior distribution. If None, it is assumed to be the standard
                deviation of a standard normal distribution. {INPUT_3D_DOC}
            reduction: The reduction method to apply to the KL divergence. If "allsum", sums over all dimensions except
                batch and averages over the batch. If "channelssum", sums over the channel dimension and averages over
                the batch. If None, no reduction is applied.
            channels_first: {CHANNELS_FIRST_DOC}

        Returns:
            kl_div: The KL divergence between the latent space distribution and the prior distribution.
        """
        if prior_mu is None:
            prior_mu = torch.zeros_like(z_mu)
        if prior_sigma is None:
            prior_sigma = torch.ones_like(z_sigma)

        prior = Normal(prior_mu, prior_sigma)
        posterior = Normal(z_mu, z_sigma)
        kl_div = kl_divergence(posterior, prior)

        if reduction is None:
            pass
        elif reduction == "allsum":
            kl_div = kl_div.sum(dim=list(range(1, kl_div.ndim))).mean()
        elif reduction == "channelssum":
            if channels_first:
                channels_dim = 1
            else:
                channels_dim = -1
            kl_div = kl_div.sum(dim=channels_dim).mean()
        else:
            raise NotImplementedError(f"Reduction {reduction} not implemented")

        return kl_div

    @populate_docstring
    def forward(
        self,
        z_mu: torch.Tensor,
        z_sigma: torch.Tensor,
        prior_mu: torch.Tensor | None = None,
        prior_sigma: torch.Tensor | None = None,
        kl_divergence_reduction: Literal["allsum", "channelssum"] | None = "allsum",
        force_sampling: bool = False,
        channels_first: bool = True,
    ) -> tuple[torch.Tensor, torch.Tensor]:
        """Samples from the latent space and computes the KL divergence between the latent space distribution and a
        prior distribution.

        Args:
            z_mu: The mean of the latent space. {INPUT_3D_DOC}
            z_sigma: The standard deviation of the latent space. {INPUT_3D_DOC}
            prior_mu: The mean of the prior distribution. If None, it is assumed to be the mean of a standard normal
                distribution. {INPUT_3D_DOC}
            prior_sigma: The standard deviation of the prior distribution. If None, it is assumed to be the standard
                deviation of a standard normal distribution. {INPUT_3D_DOC}
            kl_divergence_reduction: The reduction method to apply to the KL divergence. If "allsum", sums over all
                dimensions except batch and averages over the batch. If "channelssum", sums over the channel dimension
                and averages over the batch. If None, no reduction is applied.
            force_sampling: Whether to force sampling even in evaluation mode.
            channels_first: {CHANNELS_FIRST_DOC}

        Returns:
            z_vae: The sampled latent vector. {OUTPUT_3D_DOC}
            kl_div: The KL divergence between the latent space distribution and the prior distribution.
        """
        return self.sample(z_mu, z_sigma, force_sampling), self.kl_divergence(
            z_mu, z_sigma, prior_mu, prior_sigma, kl_divergence_reduction, channels_first
        )
