from typing import Optional
import torch
import torch.nn as nn
import torch.nn.functional as F

_ACTIVATION_MAP = {
    'relu': lambda: nn.ReLU(inplace=True),
    'leaky': lambda: nn.LeakyReLU(negative_slope=0.1, inplace=True),
    'swish': lambda: nn.SiLU(inplace=True),
    'gelu': lambda: nn.GELU(),
    'mish': lambda: nn.Mish(),
    None: lambda: nn.Identity(),
}

class Conv2dBlock(nn.Module):
    """
    A configurable convolutional block: Conv2d -> Norm -> Activation -> Dropout.

    This block is the fundamental building unit used throughout the segmentation
    models. It allows easy switching of normalization, activation, and dropout.
    """

    def __init__(
        self,
        in_channels: int,
        out_channels: int,
        kernel_size: int = 3,
        stride: int = 1,
        padding: int = 1,
        dilation: int = 1,
        groups: int = 1,
        bias: bool = False,
        norm: Optional[str] = 'bn',
        activation: Optional[str] = 'relu',
        dropout: float = 0.0
    ):
        """
        Args:
            in_channels (int): Number of input channels.
            out_channels (int): Number of output channels.
            kernel_size (int, optional): Size of the convolving kernel. Default: 3.
            stride (int, optional): Stride of the convolution. Default: 1.
            padding (int, optional): Padding added to both sides. Default: 1.
            dilation (int, optional): Spacing between kernel elements. Default: 1.
            groups (int, optional): Number of blocked connections from input to output. Default: 1.
            bias (bool, optional): If ``True``, adds a learnable bias to the output. Default: ``False``.
            norm (Optional[str]): Normalization layer to use. Options: 'bn' (BatchNorm),
                'in' (InstanceNorm), or ``None``. Default: 'bn'.
            activation (Optional[str]): Activation function. Options: 'relu', 'leaky',
                'swish', 'gelu', 'mish', or ``None``. Default: 'relu'.
            dropout (float): Dropout probability (2D dropout). If 0, no dropout. Default: 0.0.
        """
        super().__init__()
        self.conv = nn.Conv2d(
            in_channels, out_channels, kernel_size,
            stride=stride, padding=padding, dilation=dilation,
            groups=groups, bias=bias
        )

        # Normalization
        if norm == 'bn':
            self.norm = nn.BatchNorm2d(out_channels)
        elif norm == 'in':
            self.norm = nn.InstanceNorm2d(out_channels)
        elif norm is None:
            self.norm = nn.Identity()
        else:
            raise ValueError(f'Unsupported norm: {norm}')

        # Activation
        if activation not in _ACTIVATION_MAP:
            raise ValueError(f'Unsupported activation: {activation}')
        self.act = _ACTIVATION_MAP[activation]()

        self.dropout = nn.Dropout2d(dropout) if dropout > 0 else nn.Identity()

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

        Args:
            x (torch.Tensor): Input tensor of shape (B, C, H, W).

        Returns:
            torch.Tensor: Output tensor of shape (B, out_channels, H, W) (same spatial size).
        """
        x = self.conv(x)
        x = self.norm(x)
        x = self.act(x)
        x = self.dropout(x)
        return x


class DoubleConv(nn.Module):
    """
    Two consecutive convolutional blocks, often used in U‑Net decoder stages.
    """

    def __init__(
        self,
        in_channels: int,
        out_channels: int,
        mid_channels: Optional[int] = None,
        **kwargs
    ):
        """
        Args:
            in_channels (int): Number of input channels.
            out_channels (int): Number of output channels after both convolutions.
            mid_channels (Optional[int]): Number of channels in the intermediate
                (first) convolution. If ``None``, it is set to ``out_channels``.
                Default: ``None``.
            **kwargs: Additional arguments passed to each ``Conv2dBlock``
                (e.g., norm, activation, dropout, kernel_size, etc.).
        """
        super().__init__()
        if mid_channels is None:
            mid_channels = out_channels
        self.double_conv = nn.Sequential(
            Conv2dBlock(in_channels, mid_channels, **kwargs),
            Conv2dBlock(mid_channels, out_channels, **kwargs)
        )

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

        Args:
            x (torch.Tensor): Input tensor of shape (B, C, H, W).

        Returns:
            torch.Tensor: Output tensor of shape (B, out_channels, H, W).
        """
        return self.double_conv(x)


class ResidualConv(nn.Module):
    """
    A residual convolutional block: Conv -> Norm -> Act -> Conv -> Norm + skip connection.
    """

    def __init__(
        self,
        in_channels: int,
        out_channels: int,
        stride: int = 1,
        **kwargs
    ):
        """
        Args:
            in_channels (int): Number of input channels.
            out_channels (int): Number of output channels.
            stride (int, optional): Stride for the first convolution. Default: 1.
            **kwargs: Additional arguments passed to the ``Conv2dBlock``
                (e.g., norm, activation, dropout, kernel_size, etc.).
        """
        super().__init__()
        self.conv1 = Conv2dBlock(in_channels, out_channels, stride=stride, **kwargs)
        self.conv2 = Conv2dBlock(out_channels, out_channels, **kwargs)

        # Skip connection
        if in_channels != out_channels or stride != 1:
            self.skip = Conv2dBlock(
                in_channels, out_channels,
                kernel_size=1, stride=stride, padding=0,
                norm=kwargs.get('norm'), activation=None
            )
        else:
            self.skip = nn.Identity()

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        """
        Forward pass with residual addition.

        Args:
            x (torch.Tensor): Input tensor of shape (B, C, H, W).

        Returns:
            torch.Tensor: Output tensor of shape (B, out_channels, H, W).
        """
        identity = self.skip(x)
        out = self.conv1(x)
        out = self.conv2(out)
        return out + identity


class Down(nn.Module):
    """
    Downscaling block: optionally max‑pool then double convolution.
    Can also use strided convolution for learnable downsampling.
    """

    def __init__(
        self,
        in_channels: int,
        out_channels: int,
        use_maxpool: bool = True,
        **kwargs
    ):
        """
        Args:
            in_channels (int): Number of input channels.
            out_channels (int): Number of output channels after downsampling.
            use_maxpool (bool, optional): If ``True``, use max‑pool (2x2) followed by
                double convolution. If ``False``, use a strided convolution for
                downsampling. Default: ``True``.
            **kwargs: Additional arguments passed to the underlying ``DoubleConv``
                (e.g., norm, activation, dropout, etc.).
        """
        super().__init__()
        if use_maxpool:
            self.pool = nn.MaxPool2d(2)
            self.conv = DoubleConv(in_channels, out_channels, **kwargs)
        else:
            # Strided convolution for learnable downsampling
            self.conv = DoubleConv(in_channels, out_channels, stride=2, **kwargs)
            self.pool = nn.Identity()

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

        Args:
            x (torch.Tensor): Input tensor of shape (B, C, H, W).

        Returns:
            torch.Tensor: Output tensor of shape (B, out_channels, H/2, W/2)
                if downsampling was applied, otherwise spatial size unchanged.
        """
        x = self.pool(x)
        x = self.conv(x)
        return x


class Up(nn.Module):
    """
    Upscaling block for U‑Net: upsample (bilinear or transpose conv),
    reduce channels, concatenate with skip connection, then double convolution.
    """

    def __init__(
        self,
        in_channels: int,
        skip_channels: int,
        out_channels: int,
        bilinear: bool = False,
        dropout: float = 0.0,
        norm: str = 'bn',
        activation: str = 'relu'
    ):
        """
        Args:
            in_channels (int): Number of channels from the decoder (lower level).
            skip_channels (int): Number of channels from the encoder skip connection.
            out_channels (int): Desired number of output channels after the block.
            bilinear (bool, optional): If ``True``, use bilinear upsampling followed by
                a 1x1 convolution to reduce channels. If ``False``, use a transposed
                convolution. Default: ``False``.
            dropout (float): Dropout probability for the convolutional blocks.
            norm (str): Normalization type for internal convolutions.
            activation (str): Activation function for internal convolutions.
        """
        super().__init__()
        self.bilinear = bilinear

        if bilinear:
            self.up = nn.Upsample(scale_factor=2, mode='bilinear', align_corners=True)
            self.reduce = Conv2dBlock(
                in_channels, in_channels // 2,
                kernel_size=1, padding=0,
                norm=norm, activation=activation, dropout=dropout
            )
            total_channels = (in_channels // 2) + skip_channels
        else:
            self.up = nn.ConvTranspose2d(in_channels, in_channels // 2, kernel_size=2, stride=2)
            self.reduce = nn.Identity()
            total_channels = (in_channels // 2) + skip_channels

        self.conv = DoubleConv(
            total_channels, out_channels,
            norm=norm, activation=activation, dropout=dropout
        )
    def forward(self, x1: torch.Tensor, x2: torch.Tensor) -> torch.Tensor:
        x1 = self.up(x1)
        if not isinstance(self.reduce, nn.Identity):
            x1 = self.reduce(x1)

        # Resize x1 to exactly match x2's spatial dimensions
        x1 = F.interpolate(x1, size=x2.shape[-2:], mode='bilinear', align_corners=False)

        x = torch.cat([x2, x1], dim=1)
        return self.conv(x)


class AttentionGate(nn.Module):
    """
    Attention gate for U‑Net (grid attention). It weights the skip connection
    features using a gating signal from the decoder.
    """

    def __init__(self, F_g: int, F_l: int, F_int: int):
        """
        Args:
            F_g (int): Number of channels in the gating signal (from decoder).
            F_l (int): Number of channels in the skip connection (from encoder).
            F_int (int): Number of intermediate channels (reduction dimension).
        """
        super().__init__()
        self.W_g = nn.Sequential(
            nn.Conv2d(F_g, F_int, kernel_size=1, stride=1, padding=0, bias=True),
            nn.BatchNorm2d(F_int)
        )
        self.W_x = nn.Sequential(
            nn.Conv2d(F_l, F_int, kernel_size=1, stride=1, padding=0, bias=True),
            nn.BatchNorm2d(F_int)
        )
        self.psi = nn.Sequential(
            nn.Conv2d(F_int, 1, kernel_size=1, stride=1, padding=0, bias=True),
            nn.BatchNorm2d(1),
            nn.Sigmoid()
        )
        self.relu = nn.ReLU(inplace=True)

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

        Args:
            g (torch.Tensor): Gating signal from decoder, shape (B, F_g, H, W).
            x (torch.Tensor): Skip connection features, shape (B, F_l, H, W).

        Returns:
            torch.Tensor: Attention-weighted skip features, same shape as `x`.
        """
        g1 = self.W_g(g)
        x1 = self.W_x(x)
        psi = self.relu(g1 + x1)
        psi = self.psi(psi)
        return x * psi

class AttentionUp(nn.Module):
    """
    Upsampling block with an attention gate for U‑Net decoders.

    Args:
        in_channels (int): Number of channels in the input decoder feature map (x1).
        skip_channels (int): Number of channels in the encoder skip connection (x2).
        out_channels (int): Desired number of output channels after the block.
        bilinear (bool, optional): If True, use bilinear upsampling (non‑learnable)
            followed by a 1×1 convolution to reduce channels. If False, use a
            transposed convolution (learnable) that both upsamples and reduces channels.
            Default: False.
        dropout (float, optional): Dropout probability applied after each convolutional
            block inside the `DoubleConv`. Default: 0.0.
        norm (str, optional): Normalization type for internal convolutions.
            Options: 'bn' (BatchNorm), 'in' (InstanceNorm), or None. Default: 'bn'.
        activation (str, optional): Activation function for internal convolutions.
            Options: 'relu', 'leaky', 'swish', 'gelu', 'mish', or None. Default: 'relu'.
        reduction (int, optional): Reduction factor for the internal bottleneck dimension
            of the attention gate. The internal channels are set to
            `min(up_out_channels, skip_channels) // reduction`. Default: 2.
            Larger reduction (e.g., 4, 8) → smaller internal dimension → fewer parameters, lighter attention module, but may lose some capacity.
            Smaller reduction (e.g., 1, 2) → larger internal dimension → more expressive attention, but adds more parameters.

    Attributes:
        up (nn.Module): Upsampling layer (either `nn.Upsample` or `nn.ConvTranspose2d`).
        reduce (nn.Module): Channel reduction layer (either `Conv2dBlock` or `nn.Identity`).
        attention (AttentionGate): The attention gate module.
        conv (DoubleConv): Final double convolution after concatenation.

    Forward:
        Args:
            x1 (torch.Tensor): Decoder feature map from the lower (deeper) level.
                Shape: (B, in_channels, H, W).
            x2 (torch.Tensor): Skip connection feature map from the encoder.
                Shape: (B, skip_channels, H_skip, W_skip), where H_skip and W_skip
                are expected to be twice the size of H and W (after upsampling).

        Returns:
            torch.Tensor: Output feature map of shape (B, out_channels, H_skip, W_skip).
                Spatial size matches the skip connection (i.e., twice the input size).
    """

    def __init__(
        self,
        in_channels: int,
        skip_channels: int,
        out_channels: int,
        bilinear: bool = False,
        dropout: float = 0.0,
        norm: str = 'bn',
        activation: str = 'relu',
        reduction: int = 2
    ):
        super().__init__()
        self.bilinear = bilinear

        if bilinear:
            # Non‑learnable bilinear upsampling (doubles spatial size, channels unchanged)
            self.up = nn.Upsample(scale_factor=2, mode='bilinear', align_corners=True)
            self.reduce = Conv2dBlock(
                in_channels, in_channels // 2,
                kernel_size=1, padding=0,
                norm=norm, activation=activation, dropout=dropout
            )
            up_out_channels = in_channels // 2
        else:
            self.up = nn.ConvTranspose2d(in_channels, in_channels // 2, kernel_size=2, stride=2)
            self.reduce = nn.Identity()  # no further reduction needed
            up_out_channels = in_channels // 2

        # Determine internal dimension for attention gate (avoid zero)
        F_int = min(up_out_channels, skip_channels) // reduction
        if F_int == 0:
            F_int = 1

        # Attention gate: gating signal = upsampled feature, input = skip connection
        self.attention = AttentionGate(F_g=up_out_channels, F_l=skip_channels, F_int=F_int)

        # Final convolution after concatenation of attended skip and upsampled feature
        total_channels = up_out_channels + skip_channels
        self.conv = DoubleConv(
            total_channels, out_channels,
            norm=norm, activation=activation, dropout=dropout
        )

    def forward(self, x1: torch.Tensor, x2: torch.Tensor) -> torch.Tensor:
        """
        Forward pass through the attention upsampling block.

        Args:
            x1 (torch.Tensor): Decoder feature from lower level. Shape (B, in_channels, H, W).
            x2 (torch.Tensor): Skip connection from encoder. Shape (B, skip_channels, H_skip, W_skip),
                where H_skip ≈ 2*H, W_skip ≈ 2*W.

        Returns:
            torch.Tensor: Output tensor of shape (B, out_channels, H_skip, W_skip).
        """
        x1 = self.up(x1)
        if not isinstance(self.reduce, nn.Identity):
            x1 = self.reduce(x1)

        x1 = F.interpolate(x1, size=x2.shape[-2:], mode='bilinear', align_corners=False)

        attended_skip = self.attention(g=x1, x=x2)

        x = torch.cat([attended_skip, x1], dim=1)
        return self.conv(x)


class SEBlock(nn.Module):
    """
    Squeeze‑and‑Excitation block for channel attention.
    """

    def __init__(self, channels: int, reduction: int = 16):
        """
        Args:
            channels (int): Number of input channels.
            reduction (int, optional): Reduction ratio for the bottleneck.
                Default: 16.
        """
        super().__init__()
        self.global_avgpool = nn.AdaptiveAvgPool2d(1)
        self.fc = nn.Sequential(
            nn.Linear(channels, channels // reduction, bias=False),
            nn.ReLU(inplace=True),
            nn.Linear(channels // reduction, channels, bias=False),
            nn.Sigmoid()
        )

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

        Args:
            x (torch.Tensor): Input tensor of shape (B, C, H, W).

        Returns:
            torch.Tensor: Channel‑recalibrated tensor, same shape as input.
        """
        b, c, _, _ = x.size()
        y = self.global_avgpool(x).view(b, c)
        y = self.fc(y).view(b, c, 1, 1)
        return x * y.expand_as(x)