"""
Generators for the baseline / ablation study. All variants output tanh in [-1, 1], like the
released S2F generator, so they share one training and evaluation path.

  arch='s2f'   S2FGenerator (models/s2f_model.py) with optional switches:
               use_multi_scale_input  (False -> single 3x3 input conv)
               bridge                 'cbam' (paper) | 'cnn' (no CBAM) | 'plain' (dilated conv block only:
                                      no CBAM, no Global Context Module, no hierarchical attention)
               attention_gates        False -> skip connections pass encoder features unchanged
  arch='unet'  plain U-Net (Ronneberger et al., 2015): 4 levels, 64-1024 channels, double 3x3 conv + BN +
               ReLU, transposed-conv upsampling, concatenation skips; no attention modules.
"""
import torch
import torch.nn as nn

from models.blocks import ResidualBlock
from models.s2f_model import PatchGANDiscriminator, S2FGenerator


class _PassSkip(nn.Module):
    """Replaces an attention gate: returns the encoder skip feature unchanged."""
    def forward(self, g, x):
        return x


def _double_conv(i, o):
    return nn.Sequential(nn.Conv2d(i, o, 3, padding=1, bias=False), nn.BatchNorm2d(o), nn.ReLU(inplace=True),
                         nn.Conv2d(o, o, 3, padding=1, bias=False), nn.BatchNorm2d(o), nn.ReLU(inplace=True))


class UNet(nn.Module):
    def __init__(self, in_channels=1, out_channels=1, base=64):
        super().__init__()
        c = [base, base * 2, base * 4, base * 8, base * 16]
        self.downs = nn.ModuleList([_double_conv(in_channels, c[0])] + [_double_conv(c[i], c[i + 1]) for i in range(3)])
        self.pool = nn.MaxPool2d(2)
        self.bottleneck = _double_conv(c[3], c[4])
        self.ups = nn.ModuleList([nn.ConvTranspose2d(c[i + 1], c[i], 2, stride=2) for i in reversed(range(4))])
        self.decs = nn.ModuleList([_double_conv(c[i] * 2, c[i]) for i in reversed(range(4))])
        self.out = nn.Sequential(nn.Conv2d(c[0], out_channels, 1), nn.Tanh())

    def forward(self, x):
        skips = []
        for d in self.downs:
            x = d(x); skips.append(x); x = self.pool(x)
        x = self.bottleneck(x)
        for up, dec, s in zip(self.ups, self.decs, reversed(skips)):
            x = dec(torch.cat([up(x), s], dim=1))
        return self.out(x)


def build_generator(arch="s2f", in_channels=3, use_multi_scale_input=True, bridge="cbam",
                    attention_gates=True, **_):
    if arch == "unet":
        return UNet(in_channels=in_channels)
    if arch != "s2f":
        raise ValueError(f"unknown arch {arch}")
    g = S2FGenerator(in_channels=in_channels, out_channels=1, img_size=1024,
                     bridge_type="cnn" if bridge in ("cnn", "plain") else "cbam",
                     use_multi_scale_input=use_multi_scale_input)
    if bridge == "plain":
        g.bridge = nn.Sequential(nn.Conv2d(512, 1024, 3, padding=2, dilation=2), nn.BatchNorm2d(1024),
                                 nn.ReLU(inplace=True), ResidualBlock(1024, 1024))
    if not attention_gates:
        g.att4, g.att3, g.att2, g.att1 = _PassSkip(), _PassSkip(), _PassSkip(), _PassSkip()
    return g


def build_discriminator(in_channels=3):
    return PatchGANDiscriminator(in_channels=in_channels + 1, ndf=64, n_layers=3)


def n_params(module):
    return sum(p.numel() for p in module.parameters())