#######################################################################
#  TARDIS - Transformer And Rapid Dimensionless Instance Segmentation #
#                                                                     #
#  New York Structural Biology Center                                 #
#  Simons Machine Learning Center                                     #
#                                                                     #
#  Robert Kiewisz, Tristan Bepler                                     #
#  MIT License 2021 - 2025                                            #
#######################################################################

from typing import Optional

import torch
import torch.nn as nn

from tardis_em.cnn.model.convolution import (
    DoubleConvolution,
    RecurrentDoubleConvolution,
)
from tardis_em.cnn.model.init_weights import init_weights
from tardis_em.cnn.utils.utils import number_of_features_per_level


class DecoderBlockCNN(nn.Module):
    """
    DecoderBlockCNN implements a CNN-based decoder block architecture suitable
    for use in U-Net-like neural networks. It enables upscaling and optional
    dropout mechanism for regularization.

    This class builds a modular decoder block consisting of upscaling (using
    either bilinear or trilinear interpolation), a double convolution module
    configured for decoding tasks, and optional dropout. The class also
    permits initialization of the weights for its layers.
    """

    def __init__(
        self,
        in_ch: int,
        out_ch: int,
        conv_kernel: int,
        padding: int,
        size: int,
        dropout: Optional[float] = None,
        components="3gcr",
        num_group=8,
    ):
        """
        Initializes a DecoderBlockCNN module composed of multiple layers, including an
        upsampling layer, convolutions, and optional dropout, to process and upsample
        input feature maps according to predefined specifications.

        :param in_ch: Number of input channels for the decoder block.
        :type in_ch: int
        :param out_ch: Number of output channels generated by the decoder block.
        :type out_ch: int
        :param conv_kernel: Size of the convolutional kernel to be used in the
            decoder block.
        :type conv_kernel: int
        :param padding: Amount of padding applied in the convolutional operations.
        :type padding: int
        :param size: Target size for the upsampling layer of the decoder.
        :type size: int
        :param dropout: If specified, represents the dropout probability to apply
            for regularization within the dropout layer.
        :type dropout: Optional[float]
        :param components: String representation of selected features or processing
            options in the decoder block. For example, whether "3gcr" or "2gcr".
        :type components: str
        :param num_group: Number of groups for the group normalization layer in the
            decoder.
        :type num_group: int
        """
        super(DecoderBlockCNN, self).__init__()

        self.dropout = dropout

        """Build decoders"""
        if "3" in components:
            self.upscale = nn.Upsample(size=size, mode="trilinear", align_corners=False)
        elif "2" in components:
            self.upscale = nn.Upsample(size=size, mode="bilinear", align_corners=False)

        self.deconv_module = DoubleConvolution(
            in_ch=in_ch,
            out_ch=out_ch,
            block_type="decoder",
            kernel=conv_kernel,
            padding=padding,
            components=components,
            num_group=num_group,
        )

        """Optional Dropout"""
        if dropout is not None:
            self.dropout_layer = nn.Dropout(dropout)

        """Initialise the blocks"""
        for m in self.children():
            if m.__class__.__name__.find("Conv3d") != -1:
                continue
            if m.__class__.__name__.find("Conv2d") != -1:
                continue
            if m.__class__.__name__.find("GroupNorm") != -1:
                continue

            init_weights(m)

    def forward(self, encoder_features: torch.Tensor, x: torch.Tensor) -> torch.Tensor:
        """
        Combines encoder features and upscale features, applies a deconvolution
        module, and optionally applies a dropout layer to the output during the
        forward pass of the network.

        :param encoder_features: Input tensor representing features extracted
            from the encoder block.
        :param x: Input tensor that has been upsampled by the current block
            of the decoder.
        :return: A tensor obtained by combining the encoder features and
            the upsampled tensor, after applying deconvolution and optional
            dropout.
        """

        x = encoder_features + self.deconv_module(self.upscale(x))

        if self.dropout is not None:
            x = self.dropout_layer(x)

        return x


class DecoderBlockRCNN(nn.Module):
    """
    Implements a Decoder Block for a Recurrent Convolutional Neural Network (RCNN).

    This class defines a decoder block structure for RCNN models. The decoder block
    performs upscaling, residual connections, convolutional operations, and optional
    dropout regularization. It is designed for flexible integration into various neural
    network architectures, particularly those requiring upscaling and deep feature
    convolutions.
    """

    def __init__(
        self,
        in_ch: int,
        out_ch: int,
        conv_kernel: int,
        padding: int,
        size: int,
        dropout: Optional[float] = None,
        components="3gcr",
        num_group=8,
    ):
        """
        Initializes a multi-component decoder block for use in neural networks. The
        DecoderBlockRCNN class supports multiple configurations, including support for
        trilinear or bilinear upsampling, double convolutions, recurrent double
        convolutions, and optional dropout. This class facilitates the creation of
        decoder modules for use in recursive convolutional neural networks with
        flexibility to work with different component configurations and convolution
        parameters.

        :param in_ch: Number of input channels.
        :type in_ch: int
        :param out_ch: Number of output channels.
        :type out_ch: int
        :param conv_kernel: Size of the convolution kernel.
        :type conv_kernel: int
        :param padding: Size of padding to be applied during convolution.
        :type padding: int
        :param size: Target size for the upsampling layer.
        :type size: int
        :param dropout: Dropout probability (if any). Optional.
        :type dropout: Optional[float]
        :param components: String specifying the decoder components to include.
        :type components: str
        :param num_group: Number of groups in group normalization layers.
        :type num_group: int

        :raises ValueError: If unsupported components are specified.
        """
        super(DecoderBlockRCNN, self).__init__()

        self.dropout = dropout

        """Build decoders"""
        if "3" in components:
            self.upscale = nn.Upsample(size=size, mode="trilinear", align_corners=False)
        elif "2" in components:
            self.upscale = nn.Upsample(size=size, mode="bilinear", align_corners=False)

        self.deconv_module = DoubleConvolution(
            in_ch=in_ch,
            out_ch=out_ch,
            block_type="decoder",
            kernel=conv_kernel,
            padding=padding,
            components=components,
            num_group=num_group,
        )

        self.deconv_res_module = RecurrentDoubleConvolution(
            in_ch=out_ch,
            out_ch=out_ch,
            block_type="decoder",
            kernel=conv_kernel,
            padding=padding,
            components=components,
            num_group=num_group,
        )

        """Optional Dropout"""
        if dropout is not None:
            self.dropout_layer = nn.Dropout(dropout)

        """Initialise the blocks"""
        for m in self.children():
            if m.__class__.__name__.find("Conv3d") != -1:
                continue
            if m.__class__.__name__.find("Conv2d") != -1:
                continue
            if m.__class__.__name__.find("GroupNorm") != -1:
                continue

            init_weights(m)

    def forward(self, encoder_features: torch.Tensor, x: torch.Tensor) -> torch.Tensor:
        """
        The forward method implements the forward pass in the neural network layer.
        It modifies input using an upscale operation, applies a deconvolution module, and
        adds the result with encoder features. It further processes the data with a
        deconvolution residual module. If a dropout is specified, it applies the
        dropout layer to the processed output, and finally returns the result.

        :param encoder_features: Features derived from the encoder network
                                 to be combined with the processed tensor.
        :param x: Input tensor passed through this forward method for processing.

        :return: The processed tensor after upscaling, combination with encoder features,
                 residual operations, and optional dropout application.
        """
        x = self.upscale(x)
        x = self.deconv_module(x)

        x = encoder_features + x
        x = self.deconv_res_module(x)

        if self.dropout is not None:
            x = self.dropout_layer(x)

        return x


class DecoderBlockUnet3Plus(nn.Module):
    """
    DecoderBlockUnet3Plus is a class for decoding layers in a UNet3+ architecture. It is
    designed for multi-scale feature fusion with optional attention mechanism, skip
    connections, and dropout capabilities. This block helps upscale encoded spatial
    representations and recombine them with corresponding encoder features.
    """

    def __init__(
        self,
        in_ch: int,
        out_ch: int,
        conv_kernel: int,
        padding: int,
        size: int,
        components: str,
        num_group: int,
        num_layer: int,
        encoder_feature_ch: list,
        attn_features=False,
        dropout: Optional[float] = None,
    ):
        """
        DecoderBlockUnet3Plus is a class for decoding layers in a UNet3+ architecture. It is
        designed for multi-scale feature fusion with optional attention mechanism, skip
        connections, and dropout capabilities. This block helps upscale encoded spatial
        representations and recombine them with corresponding encoder features.

        :param in_ch: Number of input channels for the up-convolution.
        :param out_ch: Number of output channels for the up-convolution.
        :param conv_kernel: Size of the convolution kernel used in decoding layers.
        :param padding: Padding applied in the decoding convolution operations.
        :param size: Target spatial size for upscaling operation.
        :param components: Components indicator (e.g., "2" for 2D, "3" for 3D operations).
        :param num_group: Number of groups for group normalization in the convolution layers.
        :param num_layer: Number of encoder layers whose features will be extracted and
            concatenated with decoder features.
        :param encoder_feature_ch: List of input channel sizes for the corresponding
            encoder layers. Each channel corresponds to an encoder.
        :param attn_features: Flag to enable or disable the attention mechanism. Defaults to False.
        :param dropout: Optional value for specifying the dropout probability. If set to
            None, dropout is disabled.
        """
        super(DecoderBlockUnet3Plus, self).__init__()

        self.attn_features = attn_features
        self.dropout = dropout

        """Main Block Up-Convolution"""
        if "3" in components:
            self.upscale = nn.Upsample(size=size, mode="trilinear", align_corners=False)
        elif "2" in components:
            self.upscale = nn.Upsample(size=size, mode="bilinear", align_corners=False)

        self.deconv = DoubleConvolution(
            in_ch=in_ch,
            out_ch=out_ch,
            block_type="decoder",
            kernel=conv_kernel,
            padding=padding,
            components=components,
            num_group=num_group,
        )

        """Skip-Connection Encoders"""
        num_layer = num_layer - 1
        pool_kernels = [None, 2, 4, 8, 16, 32]
        pool_kernels = pool_kernels[:num_layer]

        self.encoder_max_pool = nn.ModuleList([])
        self.encoder_feature_conv = nn.ModuleList([])

        for i, en_in_channel in enumerate(encoder_feature_ch):
            pool_kernel = pool_kernels[i]

            if pool_kernel is not None and "3" in components:
                max_pool = nn.MaxPool3d(
                    kernel_size=pool_kernel, stride=pool_kernel, ceil_mode=True
                )
            elif pool_kernel is not None and "2" in components:
                max_pool = nn.MaxPool2d(
                    kernel_size=pool_kernel, stride=pool_kernel, ceil_mode=True
                )
            else:
                max_pool = None

            conv = DoubleConvolution(
                in_ch=en_in_channel,
                out_ch=out_ch,
                block_type="decoder",
                kernel=conv_kernel,
                padding=padding,
                components=components,
                num_group=num_group,
            )
            self.encoder_max_pool.append(max_pool)
            self.encoder_feature_conv.append(conv)

        if attn_features is not None:
            self.attn_conv = DoubleConvolution(
                in_ch=out_ch + in_ch,
                out_ch=out_ch,
                block_type="decoder",
                kernel=conv_kernel,
                padding=padding,
                components=components,
                num_group=num_group,
            )

        """Optional Dropout"""
        if dropout is not None:
            self.dropout_layer = nn.Dropout(dropout)

        """Initialise the blocks"""
        for m in self.children():
            if m.__class__.__name__.find("Conv3d") != -1:
                continue
            if m.__class__.__name__.find("Conv2d") != -1:
                continue
            if m.__class__.__name__.find("GroupNorm") != -1:
                continue

            init_weights(m)

    def forward(self, x: torch.Tensor, encoder_features: list) -> torch.Tensor:
        """
        Processes the input tensor through the forward pass of a neural network block.

        This function handles different operations, including upscaling, deconvolution,
        skip-connections with encoder features, optional attention mechanisms, and dropout
        application. It integrates different components such as upscaling the input tensor,
        merging encoder features, applying optional attention layers, and optionally applying
        dropout operation at the end.

        :param x: Input tensor to be processed through the block.
        :type x: torch.Tensor
        :param encoder_features: List of tensors representing the outputs from encoder layers.
        :type encoder_features: list[torch.Tensor]

        :return: Processed tensor obtained after applying all operations in the block.
        :rtype: torch.Tensor
        """

        """Main Block"""
        if self.attn_features:
            x_attn = self.upscale(x)
            x = self.deconv(x_attn)
        else:
            x = self.deconv(self.upscale(x))

        """Skip-Connections Encoder"""
        for i, encoder in enumerate(encoder_features):
            if self.encoder_max_pool[i] is not None:
                encoder = self.encoder_feature_conv[i](
                    self.encoder_max_pool[i](encoder)
                )
            else:
                encoder = self.encoder_feature_conv[i](encoder)
            x = x + encoder

        if self.attn_features:
            x = self.attn_conv(torch.cat((x, x_attn), dim=1))

        # Additional Dropout
        if self.dropout is not None:
            x = self.dropout_layer(x)

        return x


def build_decoder(
    conv_layers: int,
    conv_layer_scaler: int,
    components: str,
    num_group: int,
    conv_kernel: int,
    padding: int,
    sizes: list,
    dropout: Optional[float] = None,
    deconv_module="CNN",
    attn_features=False,
):
    """
    Builds a decoder module consisting of multiple decoder blocks. The choice of decoder
    block type depends on the `deconv_module` parameter, which supports CNN, RCNN, and
    Unet3Plus decoder types. The decoder is constructed based on the feature map channels
    derived from the number of convolution layers and a scaling factor. Each decoder block
    is appended to a list, which is eventually returned as a ModuleList.

    :param conv_layers: Number of convolution layers to base the feature computation.
    :param conv_layer_scaler: Scaling factor to adjust the number of features for each layer.
    :param components: Type or characteristic of components used in the decoder blocks.
    :param num_group: Number of groups for grouped convolutions, if applicable.
    :param conv_kernel: Size of the convolution kernel used within the decoder.
    :param padding: Padding to be applied in the convolution layers.
    :param sizes: List of spatial sizes for each decoder block output.
    :param dropout: Fraction for dropout regularization to prevent overfitting (optional).
    :param deconv_module: Type of the decoder module, options are "CNN", "RCNN", or
        "unet3plus".
    :param attn_features: Flag to include attention features for Unet3Plus decoder
        (specifically when `deconv_module="unet3plus"`).

    :return: A torch.nn.ModuleList containing the decoder blocks with the configured
        parameters.
    """
    decoders = []
    feature_map = number_of_features_per_level(
        channel_scaler=conv_layer_scaler, num_levels=conv_layers
    )
    feature_map = list(reversed(feature_map))

    """Forward decoder"""
    if deconv_module == "CNN":
        # Unet decoder
        for i in range(len(feature_map) - 1):
            in_ch = feature_map[i]
            out_ch = feature_map[i + 1]
            size = sizes[i]

            decoder = DecoderBlockCNN(
                in_ch=in_ch,
                out_ch=out_ch,
                conv_kernel=conv_kernel,
                padding=padding,
                size=size,
                dropout=dropout,
                components=components,
                num_group=num_group,
            )
            decoders.append(decoder)
    elif deconv_module == "RCNN":
        # ResNet decoder
        for i in range(len(feature_map) - 1):
            in_ch = feature_map[i]
            out_ch = feature_map[i + 1]
            size = sizes[i]

            decoder = DecoderBlockRCNN(
                in_ch=in_ch,
                out_ch=out_ch,
                conv_kernel=conv_kernel,
                padding=padding,
                size=size,
                dropout=dropout,
                components=components,
                num_group=num_group,
            )

            decoders.append(decoder)
    elif deconv_module == "unet3plus":
        # Unet3Plus decoder
        idx_en = 1

        for i in range(len(feature_map) - 1):
            # Main Module features
            in_ch = feature_map[i]
            out_ch = feature_map[i + 1]
            size = sizes[i]

            # Encoder De Convolution
            encoder_feature_ch = feature_map[idx_en:]
            idx_en += 1

            decoder = DecoderBlockUnet3Plus(
                in_ch=in_ch,
                out_ch=out_ch,
                conv_kernel=conv_kernel,
                padding=padding,
                size=size,
                components=components,
                num_group=num_group,
                num_layer=conv_layers,
                encoder_feature_ch=encoder_feature_ch,
                attn_features=attn_features,
                dropout=dropout,
            )
            decoders.append(decoder)

    return nn.ModuleList(decoders)
