"""ShadeNet-2 generator (standalone, no Lightning dependency).

ParallelUNet: dual parallel encoders (vanilla UNet + frozen MobileNetV2),
fused decoder -> 8ch intrinsic maps in [-1, 1]:
  [0:3] albedo | [3:4] relative depth (0=near) | [4:7] normal | [7:8] shading
"""
import torch
import torch.nn as nn
import torchvision
from torchvision.models.feature_extraction import create_feature_extractor

OUT_CH = 8
WIDTH_MULT = 1.45


def _num_groups(ch: int) -> int:
    g = min(32, ch)
    while ch % g != 0:
        g -= 1
    return g


class ChannelLinear(nn.Module):
    """Per-channel learnable affine: y = x * weight + bias."""

    def __init__(self, channels: int, init_scale: float = 0.01):
        super().__init__()
        self.weight = nn.Parameter(
            torch.empty(1, channels, 1, 1).uniform_(-init_scale, init_scale) + 1.0
        )
        self.bias = nn.Parameter(
            torch.empty(1, channels, 1, 1).uniform_(-init_scale, init_scale)
        )

    def forward(self, x):
        return x * self.weight + self.bias


class DoubleConv(nn.Module):
    def __init__(self, in_ch, out_ch):
        super().__init__()
        g = _num_groups(out_ch)
        self.conv = nn.Sequential(
            nn.Conv2d(in_ch, out_ch, (1, 3), padding=(0, 1), bias=False),
            nn.GroupNorm(g, out_ch),
            nn.ELU(inplace=True),
            nn.Conv2d(out_ch, out_ch, (3, 1), padding=(1, 0), bias=False),
            nn.GroupNorm(g, out_ch),
            nn.ELU(inplace=True),
        )

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


class Down(nn.Module):
    def __init__(self, in_ch, out_ch):
        super().__init__()
        self.pool = nn.AvgPool2d(2)
        self.conv = DoubleConv(in_ch, out_ch)

    def forward(self, x):
        return self.conv(self.pool(x))


class Up(nn.Module):
    def __init__(self, in_ch, out_ch, skip_ch=None):
        super().__init__()
        skip_ch = skip_ch or in_ch
        self.up = nn.ConvTranspose2d(in_ch, out_ch, 2, stride=2)
        self.conv = DoubleConv(out_ch + skip_ch, out_ch)

    def forward(self, x, skip):
        x = self.up(x)
        return self.conv(torch.cat([skip, x], dim=1))


class ParallelUNet(nn.Module):
    def __init__(self, in_ch=3, out_ch=OUT_CH, dropout=0.0,
                 width_mult=WIDTH_MULT):
        super().__init__()
        w = float(width_mult)

        def C(n):
            return max(8, int(round(n * w / 8.0) * 8))

        self.u_inc = DoubleConv(in_ch, C(64))
        self.u_down1 = Down(C(64), C(128))
        self.u_down2 = Down(C(128), C(256))
        self.u_down3 = Down(C(256), C(512))
        self.u_down4 = Down(C(512), C(256))
        self.u_down5 = Down(C(256), C(512))

        mean = torch.tensor([0.485, 0.456, 0.406]).view(1, 3, 1, 1)
        std = torch.tensor([0.229, 0.224, 0.225]).view(1, 3, 1, 1)
        self.register_buffer("m_shift", 1.0 - 2.0 * mean)
        self.register_buffer("m_scale", 1.0 / (2.0 * std))

        mbnet = torchvision.models.mobilenet_v2(weights="IMAGENET1K_V1")
        self.mobile = create_feature_extractor(
            mbnet,
            return_nodes={
                "features.0": "m0",
                "features.2": "m1",
                "features.6": "m2",
                "features.13": "m3",
                "features.17": "m4",
            },
        )
        for p in self.mobile.parameters():
            p.requires_grad = False

        self.bottleneck_fusion = DoubleConv(C(512) + 320, C(512))
        self.dropout = nn.Dropout2d(p=dropout) if dropout > 0 else nn.Identity()

        self.mem_u1 = ChannelLinear(C(64))
        self.mem_u2 = ChannelLinear(C(128))
        self.mem_u3 = ChannelLinear(C(256))
        self.mem_u4 = ChannelLinear(C(512))
        self.mem_u5 = ChannelLinear(C(256))
        self.mem_u6 = ChannelLinear(C(512))
        self.mem_b = ChannelLinear(C(512))
        self.mem_d0 = ChannelLinear(C(256))
        self.mem_d1 = ChannelLinear(C(256))
        self.mem_d2 = ChannelLinear(C(256))
        self.mem_d3 = ChannelLinear(C(128))
        self.mem_d4 = ChannelLinear(C(64))
        self.up0 = Up(C(512), C(256), skip_ch=C(256) + 96)
        self.up1 = Up(C(256), C(256), skip_ch=C(512) + 32)
        self.up2 = Up(C(256), C(256), skip_ch=C(256) + 24)
        self.up3 = Up(C(256), C(128), skip_ch=C(128) + 32)
        self.up4 = Up(C(128), C(64), skip_ch=C(64))

        self.head = DoubleConv(C(64) + 3, C(32))
        self.head_out = nn.Conv2d(C(32), out_ch, 3, padding=1)

    def forward(self, x):
        u1 = self.mem_u1(self.u_inc(x))
        u2 = self.mem_u2(self.u_down1(u1))
        u3 = self.mem_u3(self.u_down2(u2))
        u4 = self.mem_u4(self.u_down3(u3))
        u5 = self.mem_u5(self.u_down4(u4))
        u6 = self.mem_u6(self.u_down5(u5))

        x_imagenet = (x + self.m_shift) * self.m_scale
        mf = self.mobile(x_imagenet)
        m0, m1, m2, m3, m4 = mf["m0"], mf["m1"], mf["m2"], mf["m3"], mf["m4"]

        b = self.mem_b(self.dropout(
            self.bottleneck_fusion(torch.cat([u6, m4], dim=1))))
        d0 = self.mem_d0(self.dropout(self.up0(b, torch.cat([u5, m3], dim=1))))
        d1 = self.mem_d1(self.dropout(self.up1(d0, torch.cat([u4, m2], dim=1))))
        d2 = self.mem_d2(self.dropout(self.up2(d1, torch.cat([u3, m1], dim=1))))
        d3 = self.mem_d3(self.dropout(self.up3(d2, torch.cat([u2, m0], dim=1))))
        d4 = self.mem_d4(self.dropout(self.up4(d3, u1)))

        h = self.head(torch.cat([d4, x], dim=1))
        return torch.tanh(self.head_out(h))


def load_shadenet2(checkpoint_path, device="cpu", use_ema=True,
                   width_mult=WIDTH_MULT) -> ParallelUNet:
    """Build the generator and load shadenet2.ckpt (Lightning or raw format).

    Strips the `generator.` prefix from Lightning checkpoints and applies the
    EMA shadow (validated best) unless use_ema=False.
    """
    model = ParallelUNet(out_ch=OUT_CH, width_mult=width_mult)
    ckpt = torch.load(checkpoint_path, map_location=device, weights_only=False)
    sd = ckpt.get("state_dict", ckpt)
    if any(k.startswith("generator.") for k in sd):
        sd = {k[len("generator."):]: v for k, v in sd.items()
              if k.startswith("generator.")}
    model.load_state_dict(sd, strict=False)
    if use_ema:
        ema = ckpt.get("ema_generator") or {}
        if ema:
            model.load_state_dict(ema, strict=False)
            print(f"Using EMA weights ({len(ema)} tensors).")
        else:
            print("No EMA shadow in checkpoint, using raw weights.")
    model.eval().to(device)
    return model