"""
Ladder VAE (LVAE) Model.

The current implementation is based on "Interpretable Unsupervised Diversity Denoising
and Artefact Removal, Prakash et al."
"""

from collections.abc import Iterable
from typing import Union

import numpy as np
import torch
import torch.nn as nn

from careamics.config.architectures import register_model

from ..activation import get_activation
from .layers import (
    BottomUpDeterministicResBlock,
    BottomUpLayer,
    GateLayer,
    TopDownDeterministicResBlock,
    TopDownLayer,
)
from .utils import Interpolate, ModelType, crop_img_tensor


@register_model("LVAE")
class LadderVAE(nn.Module):
    """
    Constructor.

    Parameters
    ----------
    input_shape : int
        The size of the input image.
    output_channels : int
        The number of output channels.
    multiscale_count : int
        The number of scales for multiscale processing.
    z_dims : list[int]
        The dimensions of the latent space for each layer.
    encoder_n_filters : int
        The number of filters in the encoder.
    decoder_n_filters : int
        The number of filters in the decoder.
    encoder_conv_strides : list[int]
        The strides for the conv layers encoder.
    decoder_conv_strides : list[int]
        The strides for the conv layers decoder.
    encoder_dropout : float
        The dropout rate for the encoder.
    decoder_dropout : float
        The dropout rate for the decoder.
    nonlinearity : str
        The nonlinearity function to use.
    predict_logvar : bool
        Whether to predict the log variance.
    analytical_kl : bool
        Whether to use analytical KL divergence.

    Raises
    ------
    NotImplementedError
        If only 2D convolutions are supported.
    """

    def __init__(
        self,
        input_shape: int,
        output_channels: int,
        multiscale_count: int,
        z_dims: list[int],
        encoder_n_filters: int,
        decoder_n_filters: int,
        encoder_conv_strides: list[int],
        decoder_conv_strides: list[int],
        encoder_dropout: float,
        decoder_dropout: float,
        nonlinearity: str,
        predict_logvar: bool,
        analytical_kl: bool,
    ):
        super().__init__()

        # -------------------------------------------------------
        # Customizable attributes
        self.image_size = input_shape
        """Input image size. (Z, Y, X) or (Y, X) if the data is 2D."""
        # TODO: we need to be careful with this since used to be an int.
        # the tuple of shapes used to be `self.input_shape`.
        self.target_ch = output_channels
        self.encoder_conv_strides = encoder_conv_strides
        self.decoder_conv_strides = decoder_conv_strides
        self._multiscale_count = multiscale_count
        self.z_dims = z_dims
        self.encoder_n_filters = encoder_n_filters
        self.decoder_n_filters = decoder_n_filters
        self.encoder_dropout = encoder_dropout
        self.decoder_dropout = decoder_dropout
        self.nonlin = nonlinearity
        self.predict_logvar = predict_logvar
        self.analytical_kl = analytical_kl
        # -------------------------------------------------------

        # -------------------------------------------------------
        # Model attributes -> Hardcoded
        self.model_type = ModelType.LadderVae  # TODO remove !
        self.encoder_blocks_per_layer = 1
        self.decoder_blocks_per_layer = 1
        self.bottomup_batchnorm = True
        self.topdown_batchnorm = True
        self.topdown_conv2d_bias = True
        self.gated = True
        self.encoder_res_block_kernel = 3
        self.decoder_res_block_kernel = 3
        self.encoder_res_block_skip_padding = False
        self.decoder_res_block_skip_padding = False
        self.merge_type = "residual"
        self.no_initial_downscaling = True
        self.skip_bottomk_buvalues = 0
        self.stochastic_skip = True
        self.learn_top_prior = True
        self.res_block_type = "bacdbacd"  # TODO remove !
        self.mode_pred = False
        self.logvar_lowerbound = -5
        self._var_clip_max = 20
        self._stochastic_use_naive_exponential = False
        self._enable_topdown_normalize_factor = True

        # Attributes that handle LC -> Hardcoded
        self.enable_multiscale = self._multiscale_count > 1
        self.multiscale_retain_spatial_dims = True
        self.multiscale_lowres_separate_branch = False
        self.multiscale_decoder_retain_spatial_dims = (
            self.multiscale_retain_spatial_dims and self.enable_multiscale
        )

        # Derived attributes
        self.n_layers = len(self.z_dims)

        # Others...
        self._tethered_to_input = False
        self._tethered_ch1_scalar = self._tethered_ch2_scalar = None
        if self._tethered_to_input:
            target_ch = 1
            requires_grad = False
            self._tethered_ch1_scalar = nn.Parameter(
                torch.ones(1) * 0.5, requires_grad=requires_grad
            )
            self._tethered_ch2_scalar = nn.Parameter(
                torch.ones(1) * 2.0, requires_grad=requires_grad
            )
        # -------------------------------------------------------

        # -------------------------------------------------------
        # Data attributes
        self.color_ch = 1  # TODO for now we only support 1 channel
        self.normalized_input = True
        # -------------------------------------------------------

        # -------------------------------------------------------
        # Loss attributes
        # enabling reconstruction loss on mixed input
        self.mixed_rec_w = 0
        self.nbr_consistency_w = 0

        # -------------------------------------------------------
        # 3D related stuff
        self._mode_3D = len(self.image_size) == 3  # TODO refac
        self._model_3D_depth = self.image_size[0] if self._mode_3D else 1
        self._decoder_mode_3D = len(self.decoder_conv_strides) == 3
        if self._mode_3D and not self._decoder_mode_3D:
            assert self._model_3D_depth % 2 == 1, "3D model depth should be odd"
        assert (
            self._mode_3D is True or self._decoder_mode_3D is False
        ), "Decoder cannot be 3D when encoder is 2D"
        self._squish3d = self._mode_3D and not self._decoder_mode_3D
        self._3D_squisher = (
            None
            if not self._squish3d
            else nn.ModuleList(
                [
                    GateLayer(
                        channels=self.encoder_n_filters,
                        conv_strides=self.encoder_conv_strides,
                    )
                    for k in range(len(self.z_dims))
                ]
            )
        )
        # TODO: this bit is in the Ashesh's confusing-hacky style... Can we do better?

        # -------------------------------------------------------
        # # Training attributes
        # # can be used to tile the validation predictions
        # self._val_idx_manager = val_idx_manager
        # self._val_frame_creator = None
        # # initialize the learning rate scheduler params.
        # self.lr_scheduler_monitor = self.lr_scheduler_mode = None
        # self._init_lr_scheduler_params(config)
        # self._global_step = 0
        # -------------------------------------------------------

        # -------------------------------------------------------

        # Calculate the downsampling happening in the network
        self.downsample = [1] * self.n_layers
        self.overall_downscale_factor = np.power(2, sum(self.downsample))
        if not self.no_initial_downscaling:  # by default do another downscaling
            self.overall_downscale_factor *= 2

        assert max(self.downsample) <= self.encoder_blocks_per_layer
        assert len(self.downsample) == self.n_layers
        # -------------------------------------------------------

        # -------------------------------------------------------
        ### CREATE MODEL BLOCKS
        # First bottom-up layer: change num channels + downsample by factor 2
        # unless we want to prevent this
        self.encoder_conv_op = getattr(nn, f"Conv{len(self.encoder_conv_strides)}d")
        # TODO these should be defined for all layers here ?
        self.decoder_conv_op = getattr(nn, f"Conv{len(self.decoder_conv_strides)}d")
        # TODO: would be more readable to have a derived parameters to use like
        # `conv_dims = len(self.encoder_conv_strides)` and then use `Conv{conv_dims}d`
        stride = 1 if self.no_initial_downscaling else 2
        self.first_bottom_up = self.create_first_bottom_up(stride)

        # Input Branches for Lateral Contextualization
        self.lowres_first_bottom_ups = None
        self._init_multires()

        # Other bottom-up layers
        self.bottom_up_layers = self.create_bottom_up_layers(
            self.multiscale_lowres_separate_branch
        )

        # Top-down layers
        self.top_down_layers = self.create_top_down_layers()
        self.final_top_down = self.create_final_topdown_layer(
            not self.no_initial_downscaling
        )

        # Likelihood module
        # self.likelihood = self.create_likelihood_module()

        # Output layer --> Project to target_ch many channels
        logvar_ch_needed = self.predict_logvar is not None
        self.output_layer = self.parameter_net = self.decoder_conv_op(
            self.decoder_n_filters,
            self.target_ch * (1 + logvar_ch_needed),
            kernel_size=3,
            padding=1,
            bias=self.topdown_conv2d_bias,
        )

        # # gradient norms. updated while training. this is also logged.
        # self.grad_norm_bottom_up = 0.0
        # self.grad_norm_top_down = 0.0
        # PSNR computation on validation.
        # self.label1_psnr = RunningPSNR()
        # self.label2_psnr = RunningPSNR()
        # TODO: did you add this?

        # msg =f'[{self.__class__.__name__}] Stoc:{not self.non_stochastic_version} RecMode:{self.reconstruction_mode} TethInput:{self._tethered_to_input}'
        # msg += f' TargetCh: {self.target_ch}'
        # print(msg)

    ### SET OF METHODS TO CREATE MODEL BLOCKS
    def create_first_bottom_up(
        self,
        init_stride: int,
        num_res_blocks: int = 1,
    ) -> nn.Sequential:
        """
        Method creates the first bottom-up block of the Encoder.

        Its role is to perform a first image compression step.
        It is composed by a sequence of nn.Conv2d + non-linearity +
        BottomUpDeterministicResBlock (1 or more, default is 1).

        Parameters
        ----------
        init_stride: int
            The stride used by the intial Conv2d block.
        num_res_blocks: int, optional
            The number of BottomUpDeterministicResBlocks, default is 1.
        """
        # From what I got from Ashesh, Z should not be touched in any case.
        nonlin = get_activation(self.nonlin)
        conv_block = self.encoder_conv_op(
            in_channels=self.color_ch,
            out_channels=self.encoder_n_filters,
            kernel_size=self.encoder_res_block_kernel,
            padding=(
                0
                if self.encoder_res_block_skip_padding
                else self.encoder_res_block_kernel // 2
            ),
            stride=init_stride,
        )

        modules = [conv_block, nonlin]

        for _ in range(num_res_blocks):
            modules.append(
                BottomUpDeterministicResBlock(
                    conv_strides=self.encoder_conv_strides,
                    c_in=self.encoder_n_filters,
                    c_out=self.encoder_n_filters,
                    nonlin=nonlin,
                    downsample=False,
                    batchnorm=self.bottomup_batchnorm,
                    dropout=self.encoder_dropout,
                    res_block_type=self.res_block_type,
                    res_block_kernel=self.encoder_res_block_kernel,
                )
            )

        return nn.Sequential(*modules)

    def create_bottom_up_layers(self, lowres_separate_branch: bool) -> nn.ModuleList:
        """
        Method creates the stack of bottom-up layers of the Encoder.

        that are used to generate the so-called `bu_values`.

        NOTE:
            If `self._multiscale_count < self.n_layers`, then LC is done only in the first
            `self._multiscale_count` bottom-up layers (starting from the bottom).

        Parameters
        ----------
        lowres_separate_branch: bool
            Whether the residual block(s) used for encoding the low-res input are shared
            (`False`) or not (`True`) with the "same-size" residual block(s) in the
            `BottomUpLayer`'s primary flow.
        """
        multiscale_lowres_size_factor = 1
        nonlin = get_activation(self.nonlin)

        bottom_up_layers = nn.ModuleList([])
        for i in range(self.n_layers):
            # Whether this is the top layer
            is_top = i == self.n_layers - 1

            # LC is applied only to the first (_multiscale_count - 1) bottom-up layers
            layer_enable_multiscale = (
                self.enable_multiscale and self._multiscale_count > i + 1
            )

            # This factor determines the factor by which the low-resolution tensor is larger
            # N.B. Only used if layer_enable_multiscale == True, so we updated it only in that case
            multiscale_lowres_size_factor *= 1 + int(layer_enable_multiscale)

            # TODO: check correctness of this
            if self._multiscale_count > 1:
                output_expected_shape = (dim // 2 ** (i + 1) for dim in self.image_size)
            else:
                output_expected_shape = None

            # Add bottom-up deterministic layer at level i.
            # It's a sequence of residual blocks (BottomUpDeterministicResBlock), possibly with downsampling between them.
            bottom_up_layers.append(
                BottomUpLayer(
                    n_res_blocks=self.encoder_blocks_per_layer,
                    n_filters=self.encoder_n_filters,
                    downsampling_steps=self.downsample[i],
                    nonlin=nonlin,
                    conv_strides=self.encoder_conv_strides,
                    batchnorm=self.bottomup_batchnorm,
                    dropout=self.encoder_dropout,
                    res_block_type=self.res_block_type,
                    res_block_kernel=self.encoder_res_block_kernel,
                    gated=self.gated,
                    lowres_separate_branch=lowres_separate_branch,
                    enable_multiscale=self.enable_multiscale,  # TODO: shouldn't the arg be `layer_enable_multiscale` here?
                    multiscale_retain_spatial_dims=self.multiscale_retain_spatial_dims,
                    multiscale_lowres_size_factor=multiscale_lowres_size_factor,
                    decoder_retain_spatial_dims=self.multiscale_decoder_retain_spatial_dims,
                    output_expected_shape=output_expected_shape,
                )
            )

        return bottom_up_layers

    def create_top_down_layers(self) -> nn.ModuleList:
        """
        Method creates the stack of top-down layers of the Decoder.

        In these layer the `bu`_values` from the Encoder are merged with the `p_params` from the previous layer
        of the Decoder to get `q_params`. Then, a stochastic layer generates a sample from the latent distribution
        with parameters `q_params`. Finally, this sample is fed through a TopDownDeterministicResBlock to
        compute the `p_params` for the layer below.

        NOTE 1:
            The algorithm for generative inference approximately works as follows:
                - p_params = output of top-down layer above
                - bu = inferred bottom-up value at this layer
                - q_params = merge(bu, p_params)
                - z = stochastic_layer(q_params)
                - (optional) get and merge skip connection from prev top-down layer
                - top-down deterministic ResNet

        NOTE 2:
            When doing unconditional generation, bu_value is not available. Hence the
            merge layer is not used, and z is sampled directly from p_params.

        """
        top_down_layers = nn.ModuleList([])
        nonlin = get_activation(self.nonlin)
        # NOTE: top-down layers are created starting from the bottom-most
        for i in range(self.n_layers):
            # Check if this is the top layer
            is_top = i == self.n_layers - 1

            if self._enable_topdown_normalize_factor:  # TODO: What is this?
                normalize_latent_factor = (
                    1 / np.sqrt(2 * (1 + i)) if len(self.z_dims) > 4 else 1.0
                )
            else:
                normalize_latent_factor = 1.0

            top_down_layers.append(
                TopDownLayer(
                    z_dim=self.z_dims[i],
                    n_res_blocks=self.decoder_blocks_per_layer,
                    n_filters=self.decoder_n_filters,
                    is_top_layer=is_top,
                    conv_strides=self.decoder_conv_strides,
                    upsampling_steps=self.downsample[i],
                    nonlin=nonlin,
                    merge_type=self.merge_type,
                    batchnorm=self.topdown_batchnorm,
                    dropout=self.decoder_dropout,
                    stochastic_skip=self.stochastic_skip,
                    learn_top_prior=self.learn_top_prior,
                    top_prior_param_shape=self.get_top_prior_param_shape(),
                    res_block_type=self.res_block_type,
                    res_block_kernel=self.decoder_res_block_kernel,
                    gated=self.gated,
                    analytical_kl=self.analytical_kl,
                    vanilla_latent_hw=self.get_latent_spatial_size(i),
                    retain_spatial_dims=self.multiscale_decoder_retain_spatial_dims,
                    input_image_shape=self.image_size,
                    normalize_latent_factor=normalize_latent_factor,
                    conv2d_bias=self.topdown_conv2d_bias,
                    stochastic_use_naive_exponential=self._stochastic_use_naive_exponential,
                )
            )
        return top_down_layers

    def create_final_topdown_layer(self, upsample: bool) -> nn.Sequential:
        """Create the final top-down layer of the Decoder.

        NOTE: In this layer, (optional) upsampling is performed by bilinear interpolation
        instead of transposed convolution (like in other TD layers).

        Parameters
        ----------
        upsample: bool
            Whether to upsample the input of the final top-down layer
            by bilinear interpolation with `scale_factor=2`.
        """
        # Final top-down layer
        modules = list()

        if upsample:
            modules.append(Interpolate(scale=2))

        for i in range(self.decoder_blocks_per_layer):
            modules.append(
                TopDownDeterministicResBlock(
                    c_in=self.decoder_n_filters,
                    c_out=self.decoder_n_filters,
                    nonlin=get_activation(self.nonlin),
                    conv_strides=self.decoder_conv_strides,
                    batchnorm=self.topdown_batchnorm,
                    dropout=self.decoder_dropout,
                    res_block_type=self.res_block_type,
                    res_block_kernel=self.decoder_res_block_kernel,
                    gated=self.gated,
                    conv2d_bias=self.topdown_conv2d_bias,
                )
            )
        return nn.Sequential(*modules)

    def _init_multires(self, config=None) -> nn.ModuleList:
        """
        Method defines the input block/branch to encode/compress low-res lateral inputs.

        at different hierarchical levels
        in the multiresolution approach (LC). The role of the input branches is similar
        to the one of the first bottom-up layer in the primary flow of the Encoder,
        namely to compress the lateral input image to a degree that is compatible with
        the one of the primary flow.

        NOTE 1: Each input branch consists of a sequence of Conv2d + non-linearity
        + BottomUpDeterministicResBlock. It is meaningful to observe that the
        `BottomUpDeterministicResBlock` shares the same model attributes with the blocks
        in the primary flow of the Encoder (e.g., c_in, c_out, dropout, etc. etc.).
        Moreover, it does not perform downsampling.

        NOTE 2: `_multiscale_count` attribute defines the total number of inputs to the
        bottom-up pass. In other terms if we have the input patch and n_LC additional
        lateral inputs, we will have a total of (n_LC + 1) inputs.
        """
        stride = 1 if self.no_initial_downscaling else 2
        nonlin = get_activation(self.nonlin)
        if self._multiscale_count is None:
            self._multiscale_count = 1

        msg = (
            f"Multiscale count ({self._multiscale_count}) should not exceed the number"
            f"of bottom up layers ({self.n_layers}) by more than 1.\n"
        )
        assert (
            self._multiscale_count <= 1 or self._multiscale_count <= 1 + self.n_layers
        ), msg  # TODO how ?

        msg = (
            "Multiscale approach only supports monocrome images. "
            f"Found instead color_ch={self.color_ch}."
        )
        # assert self._multiscale_count == 1 or self.color_ch == 1, msg

        lowres_first_bottom_ups = []
        for _ in range(1, self._multiscale_count):
            first_bottom_up = nn.Sequential(
                self.encoder_conv_op(
                    in_channels=self.color_ch,
                    out_channels=self.encoder_n_filters,
                    kernel_size=5,
                    padding="same",
                    stride=stride,
                ),
                nonlin,
                BottomUpDeterministicResBlock(
                    c_in=self.encoder_n_filters,
                    c_out=self.encoder_n_filters,
                    conv_strides=self.encoder_conv_strides,
                    nonlin=nonlin,
                    downsample=False,
                    batchnorm=self.bottomup_batchnorm,
                    dropout=self.encoder_dropout,
                    res_block_type=self.res_block_type,
                ),
            )
            lowres_first_bottom_ups.append(first_bottom_up)

        self.lowres_first_bottom_ups = (
            nn.ModuleList(lowres_first_bottom_ups)
            if len(lowres_first_bottom_ups)
            else None
        )

    ### SET OF FORWARD-LIKE METHODS
    def bottomup_pass(self, inp: torch.Tensor) -> list[torch.Tensor]:
        """Wrapper of _bottomup_pass()."""
        # TODO Remove wrapper
        return self._bottomup_pass(
            inp,
            self.first_bottom_up,
            self.lowres_first_bottom_ups,
            self.bottom_up_layers,
        )

    def _bottomup_pass(
        self,
        inp: torch.Tensor,
        first_bottom_up: nn.Sequential,
        lowres_first_bottom_ups: nn.ModuleList,
        bottom_up_layers: nn.ModuleList,
    ) -> list[torch.Tensor]:
        """
        Method defines the forward pass through the LVAE Encoder, the so-called.

        Bottom-Up pass.

        Parameters
        ----------
        inp: torch.Tensor
            The input tensor to the bottom-up pass of shape (B, 1+n_LC, H, W), where n_LC
            is the number of lateral low-res inputs used in the LC approach.
            In particular, the first channel corresponds to the input patch, while the
            remaining ones are associated to the lateral low-res inputs.
        first_bottom_up: nn.Sequential
            The module defining the first bottom-up layer of the Encoder.
        lowres_first_bottom_ups: nn.ModuleList
            The list of modules defining Lateral Contextualization.
        bottom_up_layers: nn.ModuleList
            The list of modules defining the stack of bottom-up layers of the Encoder.
        """
        if self._multiscale_count > 1:
            x = first_bottom_up(inp[:, :1])
        else:
            x = first_bottom_up(inp)

        # Loop from bottom to top layer, store all deterministic nodes we
        # need for the top-down pass in bu_values list
        bu_values = []
        for i in range(self.n_layers):
            lowres_x = None
            if self._multiscale_count > 1 and i + 1 < inp.shape[1]:
                lowres_x = lowres_first_bottom_ups[i](inp[:, i + 1 : i + 2])
            x, bu_value = bottom_up_layers[i](x, lowres_x=lowres_x)
            bu_values.append(bu_value)

        return bu_values

    def topdown_pass(
        self,
        bu_values: Union[torch.Tensor, None] = None,
        n_img_prior: Union[torch.Tensor, None] = None,
        constant_layers: Union[Iterable[int], None] = None,
        forced_latent: Union[list[torch.Tensor], None] = None,
        top_down_layers: Union[nn.ModuleList, None] = None,
        final_top_down_layer: Union[nn.Sequential, None] = None,
    ) -> tuple[torch.Tensor, dict[str, torch.Tensor]]:
        """
        Method defines the forward pass through the LVAE Decoder, the so-called.

        Top-Down pass.

        Parameters
        ----------
        bu_values: torch.Tensor, optional
            Output of the bottom-up pass. It will have values from multiple layers of
            the ladder.
        n_img_prior: optional
            When `bu_values` is `None`, `n_img_prior` indicates the number of images to
            generate
            from the prior (so bottom-up pass is not used at all here).
        constant_layers: Iterable[int], optional
            A sequence of indexes associated to the layers in which a single instance's
            z is copied over the entire batch (bottom-up path is not used, so only prior
            is used here). Set to `None` to avoid this behaviour.
        forced_latent: list[torch.Tensor], optional
            A list of tensors that are used as fixed latent variables (hence, sampling
            doesn't take place in this case).
        top_down_layers: nn.ModuleList, optional
            A list of top-down layers to use in the top-down pass. If `None`, the method
            uses the default layers defined in the constructor.
        final_top_down_layer: nn.Sequential, optional
            The last top-down layer of the top-down pass. If `None`, the method uses the
            default layers defined in the constructor.
        """
        if top_down_layers is None:
            top_down_layers = self.top_down_layers
        if final_top_down_layer is None:
            final_top_down_layer = self.final_top_down

        # Default: no layer is sampled from the distribution's mode
        if constant_layers is None:
            constant_layers = []
        prior_experiment = len(constant_layers) > 0

        # If the bottom-up inference values are not given, don't do
        # inference, sample from prior instead
        inference_mode = bu_values is not None

        # Check consistency of arguments
        if inference_mode != (n_img_prior is None):
            msg = (
                "Number of images for top-down generation has to be given "
                "if and only if we're not doing inference"
            )
            raise RuntimeError(msg)
        if inference_mode and prior_experiment:
            msg = (
                "Prior experiments (e.g. sampling from mode) are not"
                " compatible with inference mode"
            )
            raise RuntimeError(msg)

        # Sampled latent variables at each layer
        z = [None] * self.n_layers
        # KL divergence of each layer
        kl = [None] * self.n_layers
        # Kl divergence restricted, only for the LC enabled setup denoiSplit.
        kl_restricted = [None] * self.n_layers
        # mean from which z is sampled.
        q_mu = [None] * self.n_layers
        # log(var) from which z is sampled.
        q_lv = [None] * self.n_layers
        # Spatial map of KL divergence for each layer
        kl_spatial = [None] * self.n_layers
        debug_qvar_max = [None] * self.n_layers
        kl_channelwise = [None] * self.n_layers
        if forced_latent is None:
            forced_latent = [None] * self.n_layers

        # Top-down inference/generation loop
        out = None
        for i in reversed(range(self.n_layers)):
            # If available, get deterministic node from bottom-up inference
            try:
                bu_value = bu_values[i]
            except TypeError:
                bu_value = None

            # Whether the current layer should be sampled from the mode
            constant_out = i in constant_layers

            # Input for skip connection
            skip_input = out

            # Full top-down layer, including sampling and deterministic part
            out, aux = top_down_layers[i](
                input_=out,
                skip_connection_input=skip_input,
                inference_mode=inference_mode,
                bu_value=bu_value,
                n_img_prior=n_img_prior,
                force_constant_output=constant_out,
                forced_latent=forced_latent[i],
                mode_pred=self.mode_pred,
                var_clip_max=self._var_clip_max,
            )
            # Save useful variables
            z[i] = aux["z"]  # sampled variable at this layer (batch, ch, h, w)
            kl[i] = aux["kl_samplewise"]  # (batch, )
            kl_restricted[i] = aux["kl_samplewise_restricted"]
            kl_spatial[i] = aux["kl_spatial"]  # (batch, h, w)
            q_mu[i] = aux["q_mu"]
            q_lv[i] = aux["q_lv"]

            kl_channelwise[i] = aux["kl_channelwise"]
            debug_qvar_max[i] = aux["qvar_max"]
            # if self.mode_pred is False:
            #     logprob_p += aux['logprob_p'].mean()  # mean over batch
            # else:
            #     logprob_p = None

        # Final top-down layer
        out = final_top_down_layer(out)

        # Store useful variables in a dict to return them
        data = {
            "z": z,  # list of tensors with shape (batch, ch[i], h[i], w[i])
            "kl": kl,  # list of tensors with shape (batch, )
            "kl_restricted": kl_restricted,  # list of tensors with shape (batch, )
            "kl_spatial": kl_spatial,  # list of tensors w shape (batch, h[i], w[i])
            "kl_channelwise": kl_channelwise,  # list of tensors with shape (batch, ch[i])
            # 'logprob_p': logprob_p,  # scalar, mean over batch
            "q_mu": q_mu,
            "q_lv": q_lv,
            "debug_qvar_max": debug_qvar_max,
        }
        return out, data

    def forward(self, x: torch.Tensor) -> tuple[torch.Tensor, dict[str, torch.Tensor]]:
        """
        Forward pass through the LVAE model.

        Parameters
        ----------
        x: torch.Tensor
            The input tensor of shape (B, C, H, W).
        """
        img_size = x.size()[2:]

        # Bottom-up inference: return list of length n_layers (bottom to top)
        bu_values = self.bottomup_pass(x)
        for i in range(0, self.skip_bottomk_buvalues):
            bu_values[i] = None

        if self._squish3d:
            bu_values = [
                torch.mean(self._3D_squisher[k](bu_value), dim=2)
                for k, bu_value in enumerate(bu_values)
            ]

        # Top-down inference/generation
        out, td_data = self.topdown_pass(bu_values)

        if out.shape[-1] > img_size[-1]:
            # Restore original image size
            out = crop_img_tensor(out, img_size)

        out = self.output_layer(out)

        return out, td_data

    ### SET OF GETTERS
    def get_padded_size(self, size):
        """
        Returns the smallest size (H, W) of the image with actual size given
        as input, such that H and W are powers of 2.
        :param size: input size, tuple either (N, C, H, W) or (H, W)
        :return: 2-tuple (H, W)
        """
        # Make size argument into (heigth, width)
        # assert len(size) in [2, 4, 5] # TODO commented out cuz it's weird
        # We're only interested in the Y,X dimensions
        size = size[-2:]

        if self.multiscale_decoder_retain_spatial_dims is True:
            # In this case, we can go much more deeper and so this is not required
            # (in the way it is. ;). More work would be needed if this was to be correctly implemented )
            return list(size)

        # Overall downscale factor from input to top layer (power of 2)
        dwnsc = self.overall_downscale_factor

        # Output smallest powers of 2 that are larger than current sizes
        padded_size = [((s - 1) // dwnsc + 1) * dwnsc for s in size]
        # TODO Needed for pad/crop odd sizes. Move to dataset?
        return padded_size

    def get_latent_spatial_size(self, level_idx: int):
        """Level_idx: 0 is the bottommost layer, the highest resolution one."""
        actual_downsampling = level_idx + 1
        dwnsc = 2**actual_downsampling
        sz = self.get_padded_size(self.image_size)
        h = sz[0] // dwnsc
        w = sz[1] // dwnsc
        assert h == w
        return h

    def get_top_prior_param_shape(self, n_imgs: int = 1):

        # Compute the total downscaling performed in the Encoder
        if self.multiscale_decoder_retain_spatial_dims is False:
            dwnsc = self.overall_downscale_factor
        else:
            # LC allow the encoder latents to keep the same (H, W) size at different levels
            actual_downsampling = self.n_layers + 1 - self._multiscale_count
            dwnsc = 2**actual_downsampling

        h = self.image_size[-2] // dwnsc
        w = self.image_size[-1] // dwnsc
        mu_logvar = self.z_dims[-1] * 2  # mu and logvar
        top_layer_shape = (n_imgs, mu_logvar, h, w)
        # TODO refactor!
        if self._model_3D_depth > 1 and self._decoder_mode_3D is True:
            # TODO check if model_3D_depth is needed ?
            top_layer_shape = (n_imgs, mu_logvar, self._model_3D_depth, h, w)
        return top_layer_shape
