# src/mindglow/seg2d/models/blocks.py
from typing import Optional
import torch
import torch.nn as nn
import torch.nn.functional as F

class Conv2dBlock(nn.Module):
    """(Conv2d -> Norm -> Activation) with optional dropout."""
    def __init__(self,
                 in_channels,
                 out_channels,
                 kernel_size=3,
                 stride=1,
                 padding=1,
                 dilation=1,
                 groups=1,
                 bias=False,
                 norm: Optional['str']='bn',           # 'bn', 'in', or None
                 activation: Optional['str']='relu',   # 'relu', 'leaky', 'swish', None
                 dropout=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 == 'relu':
            self.act = nn.ReLU(inplace=True)
        elif activation == 'leaky':
            self.act = nn.LeakyReLU(negative_slope=0.1, inplace=True)
        elif activation == 'swish':
            self.act = nn.SiLU(inplace=True)
        elif activation is None:
            self.act = nn.Identity()
        else:
            raise ValueError(f'Unsupported activation: {activation}')

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

    def forward(self, x):
        x = self.conv(x)
        x = self.norm(x)
        x = self.act(x)
        x = self.dropout(x)
        return x


class DoubleConv(nn.Module):
    """Two consecutive Conv2dBlocks, used in U-Net decoder stages."""
    def __init__(self,
                 in_channels,
                 out_channels,
                 mid_channels=None,
                 **kwargs):   # passes norm, activation, dropout, 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):
        return self.double_conv(x)


class ResidualConv(nn.Module):
    """Conv -> Norm -> Act -> Conv -> Norm + skip connection."""
    def __init__(self,
                 in_channels,
                 out_channels,
                 stride=1,
                 **kwargs):
        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):
        identity = self.skip(x)
        out = self.conv1(x)
        out = self.conv2(out)
        return out + identity


class Down(nn.Module):
    """Downscaling with maxpool then double conv."""
    def __init__(self,
                 in_channels,
                 out_channels,
                 use_maxpool=True,
                 **kwargs):
        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):
        x = self.pool(x)
        x = self.conv(x)
        return x


class Up(nn.Module):
    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'):
        super().__init__()
        self.bilinear = bilinear

        # Upsampling
        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 after concat = (in_channels // 2) + skip_channels
            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)

        # Pad if necessary
        diffY = x2.size()[2] - x1.size()[2]
        diffX = x2.size()[3] - x1.size()[3]
        x1 = F.pad(x1, [diffX // 2, diffX - diffX // 2,
                        diffY // 2, diffY - diffY // 2])

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

class AttentionGate(nn.Module):
    """Grid attention gating signal for skip connections."""
    def __init__(self, F_g, F_l, F_int):
        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, x):
        g1 = self.W_g(g)
        x1 = self.W_x(x)
        psi = self.relu(g1 + x1)
        psi = self.psi(psi)
        return x * psi


class SEBlock(nn.Module):
    """Channel attention."""
    def __init__(self, channels, reduction=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):
        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)