import torch
import torch.nn as nn
import torch.nn.functional as F


def srm_filter(x: torch.Tensor) -> torch.Tensor:
    """Apply fixed high-pass filters to expose noise residuals."""
    kernels = torch.tensor([
        [[ 0,  0,  0], [ 0, -1,  0], [ 0,  1,  0]],
        [[ 0,  0,  0], [ 0, -1,  1], [ 0,  0,  0]],
        [[ 0, -1,  0], [-1,  4, -1], [ 0, -1,  0]],
    ], dtype=torch.float32).unsqueeze(1).to(x.device)  # (3,1,3,3)

    B, C, H, W = x.shape
    residuals  = []
    for c in range(C):
        ch       = x[:, c:c+1, :, :]
        filtered = F.conv2d(ch, kernels, padding=1)
        residuals.append(filtered[:, :1, :, :])
    return torch.cat(residuals, dim=1)   # (B, 3, H, W)


class ConvBlock(nn.Module):
    def __init__(self, in_ch: int, out_ch: int):
        super().__init__()
        self.block = nn.Sequential(
            nn.Conv2d(in_ch, out_ch, 3, padding=1, bias=False),
            nn.BatchNorm2d(out_ch),
            nn.ReLU(inplace=True),
            nn.Conv2d(out_ch, out_ch, 3, padding=1, bias=False),
            nn.BatchNorm2d(out_ch),
            nn.ReLU(inplace=True),
        )

    def forward(self, x):
        return self.block(x)


class TamperNet(nn.Module):
    def __init__(self, base_ch: int = 32):
        super().__init__()
        # Encoder — RGB stream
        self.e1_rgb = ConvBlock(3, base_ch)
        self.e2_rgb = ConvBlock(base_ch,     base_ch * 2)
        self.e3_rgb = ConvBlock(base_ch * 2, base_ch * 4)

        # Encoder — SRM noise stream
        self.e1_srm = ConvBlock(3, base_ch)
        self.e2_srm = ConvBlock(base_ch,     base_ch * 2)
        self.e3_srm = ConvBlock(base_ch * 2, base_ch * 4)

        self.pool = nn.MaxPool2d(2)

        # Bottleneck (both streams fused)
        self.bottleneck = ConvBlock(base_ch * 8, base_ch * 8)

        # Decoder
        self.up3  = nn.ConvTranspose2d(base_ch * 8, base_ch * 4, 2, stride=2)
        self.dec3 = ConvBlock(base_ch * 12, base_ch * 4)
        self.up2  = nn.ConvTranspose2d(base_ch * 4, base_ch * 2, 2, stride=2)
        self.dec2 = ConvBlock(base_ch * 6,  base_ch * 2)
        self.up1  = nn.ConvTranspose2d(base_ch * 2, base_ch,     2, stride=2)
        self.dec1 = ConvBlock(base_ch * 3,  base_ch)

        # Output heads
        self.mask_head = nn.Conv2d(base_ch, 1, 1)
        self.cls_head  = nn.Sequential(
            nn.AdaptiveAvgPool2d(1),
            nn.Flatten(),
            nn.Linear(base_ch * 8, 1),
        )

    def forward(self, x: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
        srm = srm_filter(x)

        # Encode both streams
        e1r = self.e1_rgb(x);              e1s = self.e1_srm(srm)
        e2r = self.e2_rgb(self.pool(e1r)); e2s = self.e2_srm(self.pool(e1s))
        e3r = self.e3_rgb(self.pool(e2r)); e3s = self.e3_srm(self.pool(e2s))

        # Fuse at bottleneck (pool a third time so decoder's 3 upsamplings cancel)
        fused    = self.bottleneck(torch.cat([self.pool(e3r), self.pool(e3s)], dim=1))
        cls_logit = self.cls_head(fused)

        # Decode with skip connections (align spatial size before cat)
        up3_out = self.up3(fused)
        e3r_ = F.interpolate(e3r, size=up3_out.shape[2:])
        e3s_ = F.interpolate(e3s, size=up3_out.shape[2:])
        d3   = self.dec3(torch.cat([up3_out, e3r_, e3s_], dim=1))

        up2_out = self.up2(d3)
        e2r_ = F.interpolate(e2r, size=up2_out.shape[2:])
        e2s_ = F.interpolate(e2s, size=up2_out.shape[2:])
        d2   = self.dec2(torch.cat([up2_out, e2r_, e2s_], dim=1))

        up1_out = self.up1(d2)
        e1r_ = F.interpolate(e1r, size=up1_out.shape[2:])
        e1s_ = F.interpolate(e1s, size=up1_out.shape[2:])
        d1   = self.dec1(torch.cat([up1_out, e1r_, e1s_], dim=1))

        mask = torch.sigmoid(self.mask_head(d1))
        return mask, cls_logit