"""Implementation of MaxViT module."""

from dataclasses import dataclass
from typing import List, Tuple, Type, Union

import torch
from torch import nn

from metnet.layers.MBConv import MBConv
from metnet.layers.PartitionAttention import BlockAttention, GridAttention


@dataclass
class MaxViTDataClass:
    """
    DataClass for MaxViT.

    Parameters
    ----------
    mb_conv_expansion_rate : int, optional
        MBConv: Expansion rate for the output channels, by default 4
    mb_conv_downscale : bool, optional
        MBConv: Flag to denote downscaling in the conv branch, by default False
    mb_conv_act_layer : Type[nn.Module], optional
        MBConv: activation layer, by default nn.GELU
    mb_conv_norm_layer : Type[nn.Module], optional
        MBConv: norm layer, by default nn.BatchNorm2D
    mb_conv_drop_path : float, optional
        MBConv: Stochastic Depth ratio, by default 0.0
    mb_conv_kernel_size : int, optional
        MBConv: Conv kernel size, by default 3
    mb_conv_se_bottleneck_ratio : float, optional
        MBConv: Squeeze Excite reduction ratio, by default 0.25
    block_attention_num_heads : int, optional
        BlockAttention: Number of attention heads, by default 32
    block_attention_channels : int
        BlockAttention: Number of channels used for attention computations.
        It should be divisible by num_heads, by default 64
    block_attention_attn_grid_window_size : Tuple[int, int], optional
        BlockAttention: Grid/Window size for attention, by default (8, 8)
    block_attention_attn_drop : float, optional
        BlockAttention: Dropout ratio of attention weight, by default 0
    block_attention_proj_drop : float, optional
        BlockAttention: Dropout ratio of output, by default 0
    block_attention_drop_path : float, optional
        BlockAttention: Stochastic depth, by default 0
    block_attention_pre_norm_layer : Type[nn.Module], optional
        BlockAttention: Pre norm layer, by default nn.LayerNorm
    block_attention_post_norm_layer : Type[nn.Module], optional
        BlockAttention: Post norm layer, by default nn.LayerNorm
    block_attention_use_mlp : Type[nn.Module], optional
        BlockAttention: MLP to be used after the attention, by default None
    block_attention_use_normalised_qk : bool, optional
        BlockAttention: Normalise queries and keys as done in Metnet 3, by default True
    grid_attention_num_heads : int, optional
        GridAttention: Number of attention heads, by default 32
    grid_attention_channels : int
        GridAttention: Number of channels used for attention computations.
        It should be divisible by num_heads, by default 64
    grid_attention_attn_grid_window_size : Tuple[int, int], optional
        GridAttention: Grid/Window size for attention, by default (8, 8)
    grid_attention_attn_drop : float, optional
        GridAttention: Dropout ratio of attention weight, by default 0
    grid_attention_proj_drop : float, optional
        GridAttention: Dropout ratio of output, by default 0
    grid_attention_drop_path : float, optional
        GridAttention: Stochastic depth, by default 0
    grid_attention_pre_norm_layer : Type[nn.Module], optional
        GridAttention: Pre norm layer, by default nn.LayerNorm
    grid_attention_post_norm_layer : Type[nn.Module], optional
        GridAttention: Post norm layer, by default nn.LayerNorm
    grid_attention_use_mlp : Type[nn.Module], optional
        GridAttention: MLP to be used after the attention, by default None
    grid_attention_use_normalised_qk : bool, optional
        GridAttention: Normalise queries and keys as done in Metnet 3, by default True
    """

    mb_conv_expansion_rate: int = 4
    mb_conv_downscale: bool = False
    mb_conv_act_layer: Type[nn.Module] = nn.GELU
    mb_conv_norm_layer: Type[nn.Module] = nn.BatchNorm2d
    mb_conv_drop_path: float = 0.0
    mb_conv_kernel_size: int = 3
    mb_conv_se_bottleneck_ratio: float = 0.25
    block_attention_num_heads: int = 32
    block_attention_channels: int = 64
    block_attention_attn_grid_window_size: Tuple[int, int] = (8, 8)
    block_attention_attn_drop: float = 0
    block_attention_proj_drop: float = 0
    block_attention_drop_path: float = 0
    block_attention_pre_norm_layer: Type[nn.Module] = nn.LayerNorm
    block_attention_post_norm_layer: Type[nn.Module] = nn.LayerNorm
    block_attention_use_mlp: Type[nn.Module] = None
    block_attention_use_normalised_qk: bool = True
    grid_attention_num_heads: int = 32
    grid_attention_channels: int = 64
    grid_attention_attn_grid_window_size: Tuple[int, int] = (8, 8)
    grid_attention_attn_drop: float = 0
    grid_attention_proj_drop: float = 0
    grid_attention_drop_path: float = 0
    grid_attention_pre_norm_layer: Type[nn.Module] = nn.LayerNorm
    grid_attention_post_norm_layer: Type[nn.Module] = nn.LayerNorm
    grid_attention_use_mlp: Type[nn.Module] = None
    grid_attention_use_normalised_qk: bool = True


class MaxViTBlock(nn.Module):
    """MaxViT block."""

    def __init__(self, in_channels: int, maxvit_config: Type[MaxViTDataClass]) -> None:
        """
        Class Constructor Method.

        Parameters
        ----------
        in_channels : int
            Number of input channels
        maxvit_config : Type[MaxViTDataClass]
            MaxVit Config
        """
        super().__init__()
        self.in_channels = in_channels
        self.config = maxvit_config

        self.mb_conv = MBConv(
            in_channels=self.in_channels,
            expansion_rate=self.config.mb_conv_expansion_rate,
            downscale=self.config.mb_conv_downscale,
            act_layer=self.config.mb_conv_act_layer,
            norm_layer=self.config.mb_conv_norm_layer,
            drop_path=self.config.mb_conv_drop_path,
            kernel_size=self.config.mb_conv_kernel_size,
            se_bottleneck_ratio=self.config.mb_conv_se_bottleneck_ratio,
        )

        # Init Block and Grid Attention

        self.block_attention = BlockAttention(
            in_channels=self.in_channels,
            num_heads=self.config.block_attention_num_heads,
            attention_channels=self.config.block_attention_channels,
            attn_grid_window_size=self.config.block_attention_attn_grid_window_size,
            attn_drop=self.config.block_attention_attn_drop,
            proj_drop=self.config.block_attention_proj_drop,
            drop_path=self.config.block_attention_drop_path,
            pre_norm_layer=self.config.block_attention_pre_norm_layer,
            post_norm_layer=self.config.block_attention_post_norm_layer,
            use_mlp=self.config.block_attention_use_mlp,
            use_normalised_qk=self.config.block_attention_use_normalised_qk,
        )

        self.grid_attention = GridAttention(
            in_channels=self.in_channels,
            num_heads=self.config.grid_attention_num_heads,
            attention_channels=self.config.grid_attention_channels,
            attn_grid_window_size=self.config.grid_attention_attn_grid_window_size,
            attn_drop=self.config.grid_attention_attn_drop,
            proj_drop=self.config.grid_attention_proj_drop,
            drop_path=self.config.grid_attention_drop_path,
            pre_norm_layer=self.config.grid_attention_pre_norm_layer,
            post_norm_layer=self.config.grid_attention_post_norm_layer,
            use_mlp=self.config.grid_attention_use_mlp,
            use_normalised_qk=self.config.grid_attention_use_normalised_qk,
        )

    def forward(self, X: torch.Tensor) -> torch.Tensor:
        """
        Forward pass.

        Parameters
        ----------
        X : torch.Tensor
            Input tensor of the shape [N, C, H, W]

        Returns:
        -------
        torch.Tensor
            MaxViT block output tensor of the shape [N, C, H, W]
        """
        output = self.mb_conv(X)
        output = self.block_attention(output)
        output = self.grid_attention(output)

        return output


class MetNetMaxVit(nn.Module):
    """MaxViT block implemented with MetNet 3 modifications."""

    def __init__(
        self,
        in_channels: int = 512,
        out_channels: int = None,
        num_blocks: int = 12,
        maxvit_conf: Union[Type[MaxViTDataClass], List[Type[MaxViTDataClass]]] = MaxViTDataClass(),
        set_linear_stocastic_depth: bool = True,
    ) -> None:
        """
        MetNet3 MaxViT Block.

        Parameters
        ----------
        in_channels : int, optional
            Input Channels, by default 512
        out_channels : int, optional
            Output Channels, by default None
        num_blocks : int, optional
            Number of MaxViT blocks, by default 12
        maxvit_conf : Union[ Type[MaxViTDataClass], List[Type[MaxViTDataClass]] ], optional
            MaxViT config, by default MaxViTDataClass()
        set_linear_stocastic_depth : bool, optional
            Flag to set the stochastic depth linearly in each MaxVit subblock, by default True
        """
        super().__init__()
        self.in_channels = in_channels
        self.out_channels = out_channels if out_channels else in_channels

        self.num_blocks = num_blocks
        self.set_linear_stocastic_depth = set_linear_stocastic_depth

        self.maxvit_blocks = nn.ModuleList()

        if isinstance(maxvit_conf, List):
            assert len(maxvit_conf) == num_blocks
            self.maxvit_conf_list = maxvit_conf
        else:
            self.maxvit_conf_list = [maxvit_conf for _ in range(self.num_blocks)]

        if self.set_linear_stocastic_depth:
            # Linearly sets the stochastic depth a given sub-module
            # (i.e. MBConv, local (block) attention or gridded (grid) attention)
            # from 0 to 0.2, as mentioned in Metnet3 paper
            for conf in self.maxvit_conf_list:
                conf.mb_conv_drop_path = 0
                conf.block_attention_drop_path = 0.1
                conf.grid_attention_drop_path = 0.2

        for conf in self.maxvit_conf_list:
            self.maxvit_blocks.append(MaxViTBlock(in_channels=self.in_channels, maxvit_config=conf))

        self.linear_transform = nn.Conv2d(
            in_channels=self.in_channels, out_channels=self.in_channels, kernel_size=1
        )

    def forward(self, X: torch.Tensor) -> torch.Tensor:
        """
        Forward method.

        Parameters
        ----------
        X : torch.Tensor
            Input tensor

        Returns:
        -------
        torch.Tensor
            Output of the MaxViT block.
        """
        model_output_list = []
        model_output_list.append(self.maxvit_blocks[0](X))
        for i in range(1, self.num_blocks):
            model_output_list.append(self.maxvit_blocks[i](model_output_list[i - 1]))

        output = X + torch.stack(model_output_list).sum(dim=0)

        output = self.linear_transform(output)
        return output
