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

ParallelUNet v3: dual parallel encoders (vanilla UNet + frozen MobileNetV2),
depthwise-separable factorized convs, full H/32 bottleneck, and a patch
dictionary output tail -> 8ch intrinsic maps in [-1, 1]:
  [0:3] albedo | [3:4] relative depth (0=near) | [4:7] normal | [7:8] shading
"""
import math

import torch
import torch.nn as nn
import torchvision
from torchvision.models.feature_extraction import create_feature_extractor

OUT_CH = 8
WIDTH_MULT = 0.9


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


def _scale(width_mult):
    """Channel scaler: width_mult rounded to multiples of 8 (GroupNorm-safe)."""
    w = float(width_mult)

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

    return C


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):
    """Factorized (1,3)+(3,1) conv pair with GN+ELU, reflect padding."""

    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),
                      padding_mode="reflect", bias=False),
            nn.GroupNorm(g, out_ch),
            nn.ELU(inplace=True),
            nn.Conv2d(out_ch, out_ch, (3, 1), padding=(1, 0),
                      padding_mode="reflect", bias=False),
            nn.GroupNorm(g, out_ch),
            nn.ELU(inplace=True),
        )

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


class DSDoubleConv(nn.Module):
    """Depthwise-separable DoubleConv: same geometry, ~3x fewer params."""

    def __init__(self, in_ch, out_ch):
        super().__init__()
        g = _num_groups(out_ch)
        self.conv = nn.Sequential(
            nn.Conv2d(in_ch, in_ch, (1, 3), padding=(0, 1),
                      padding_mode="reflect", groups=in_ch, bias=False),
            nn.Conv2d(in_ch, out_ch, 1, bias=False),
            nn.GroupNorm(g, out_ch),
            nn.ELU(inplace=True),
            nn.Conv2d(out_ch, out_ch, (3, 1), padding=(1, 0),
                      padding_mode="reflect", groups=out_ch, bias=False),
            nn.Conv2d(out_ch, out_ch, 1, 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, block=DSDoubleConv):
        super().__init__()
        self.pool = nn.AvgPool2d(2)
        self.conv = block(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, block=DSDoubleConv):
        super().__init__()
        skip_ch = skip_ch or in_ch
        self.up = nn.ConvTranspose2d(in_ch, out_ch, 2, stride=2)
        self.conv = block(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 PatchDictionaryBias(nn.Module):
    """Soft patch-dictionary prior: tile -> softmax over atoms -> blend."""

    def __init__(self, channels=OUT_CH, patch=16, n_atoms=1024, hidden=128,
                 deblock=True, blur_radius=2):
        super().__init__()
        self.patch = int(patch)
        self.n_atoms = int(n_atoms)
        self.atoms = nn.Parameter(
            torch.zeros(n_atoms, channels, patch, patch))
        self.desc = nn.Conv2d(channels, channels, 4, stride=4,
                              groups=channels, bias=False)
        d = channels * (patch // 4) * (patch // 4)
        self.addr = nn.Sequential(
            nn.Linear(d, hidden),
            nn.ELU(inplace=True),
            nn.Linear(hidden, n_atoms),
        )
        self.blend = nn.Conv2d(2 * channels, channels, 1, bias=True)
        self.blur_recon = bool(deblock)
        self.blur_radius = int(blur_radius)
        r = self.blur_radius
        row = torch.tensor([math.comb(2 * r, i) for i in range(2 * r + 1)],
                           dtype=torch.float32)
        row = row / row.sum()
        k2 = (row[:, None] * row[None, :])[None, None].expand(
            channels, 1, -1, -1).contiguous()
        self.register_buffer("blur_k", k2)

    def _gauss_blur(self, t):
        r = self.blur_radius
        t = torch.nn.functional.pad(t, (r, r, r, r), mode="reflect")
        return torch.nn.functional.conv2d(t, self.blur_k, groups=t.shape[1])

    def _deblock(self, recon, b, c, nh, nw, p):
        img = (recon.reshape(b, nh, nw, c, p, p)
               .permute(0, 3, 1, 4, 2, 5)
               .reshape(b, c, nh * p, nw * p))
        r = self.blur_radius
        blurred = self._gauss_blur(img)
        H, W = img.shape[-2:]
        ih = torch.arange(H, device=img.device)
        iw = torch.arange(W, device=img.device)
        mh = ((ih % p) < r) | ((ih % p) >= p - r)
        mw = ((iw % p) < r) | ((iw % p) >= p - r)
        mask = (mh[:, None] | mw[None, :]).to(img.dtype)[None, None]
        img = img + mask * (blurred - img)
        return (img.reshape(b, c, nh, p, nw, p)
                .permute(0, 2, 4, 1, 3, 5)
                .reshape(b * nh * nw, c, p, p))

    def forward(self, x):
        b, c, h, w = x.shape
        p = self.patch
        nh, nw = h // p, w // p
        tiles = (x.reshape(b, c, nh, p, nw, p)
                 .permute(0, 2, 4, 1, 3, 5)
                 .reshape(b * nh * nw, c, p, p))
        feat = self.desc(tiles).flatten(1)
        wts = torch.softmax(self.addr(feat), dim=1)
        recon = torch.einsum("na,achw->nchw", wts, self.atoms)
        if self.blur_recon:
            recon = self._deblock(recon, b, c, nh, nw, p)
        both = torch.cat([tiles, recon], dim=1)
        out = self.blend(both).reshape(b, nh, nw, c, p, p)
        return out.permute(0, 3, 1, 4, 2, 5).reshape(b, c, nh * p, nw * p)


class VanillaEncoder(nn.Module):
    def __init__(self, in_ch=3, width_mult=1.0):
        super().__init__()
        C = _scale(width_mult)
        self.inc = DoubleConv(in_ch, C(64))
        self.down1 = Down(C(64), C(128))
        self.down2 = Down(C(128), C(256))
        self.down3 = Down(C(256), C(512))
        self.down4 = Down(C(512), C(256))
        self.down5 = Down(C(256), C(512))
        self.mem1 = ChannelLinear(C(64))
        self.mem2 = ChannelLinear(C(128))
        self.mem3 = ChannelLinear(C(256))
        self.mem4 = ChannelLinear(C(512))
        self.mem5 = ChannelLinear(C(256))
        self.mem6 = ChannelLinear(C(512))

    def forward(self, x):
        u1 = self.mem1(self.inc(x))
        u2 = self.mem2(self.down1(u1))
        u3 = self.mem3(self.down2(u2))
        u4 = self.mem4(self.down3(u3))
        u5 = self.mem5(self.down4(u4))
        u6 = self.mem6(self.down5(u5))
        return u1, u2, u3, u4, u5, u6


class MobileEncoder(nn.Module):
    def __init__(self):
        super().__init__()
        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("shift", 1.0 - 2.0 * mean)
        self.register_buffer("scale", 1.0 / (2.0 * std))
        mbnet = torchvision.models.mobilenet_v2(weights="IMAGENET1K_V1")
        self.trunk = 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.trunk.parameters():
            p.requires_grad = False

    def forward(self, x):
        feats = self.trunk((x + self.shift) * self.scale)
        return feats["m0"], feats["m1"], feats["m2"], feats["m3"], feats["m4"]


class FusionBottleneck(nn.Module):
    def __init__(self, width_mult=1.0):
        super().__init__()
        C = _scale(width_mult)
        self.m_proj = nn.Conv2d(320, C(96), 1, bias=False)
        self.fusion = DSDoubleConv(C(512) + C(96), C(512))
        self.mem = ChannelLinear(C(512))

    def forward(self, u6, m4):
        return self.mem(self.fusion(torch.cat([u6, self.m_proj(m4)], dim=1)))


class FusedDecoder(nn.Module):
    def __init__(self, width_mult=1.0):
        super().__init__()
        C = _scale(width_mult)
        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.mem0 = ChannelLinear(C(256))
        self.mem1 = ChannelLinear(C(256))
        self.mem2 = ChannelLinear(C(256))
        self.mem3 = ChannelLinear(C(128))
        self.mem4 = ChannelLinear(C(64))

    def forward(self, b, uni, mob):
        u1, u2, u3, u4, u5 = uni
        m0, m1, m2, m3 = mob
        d0 = self.mem0(self.up0(b, torch.cat([u5, m3], dim=1)))
        d1 = self.mem1(self.up1(d0, torch.cat([u4, m2], dim=1)))
        d2 = self.mem2(self.up2(d1, torch.cat([u3, m1], dim=1)))
        d3 = self.mem3(self.up3(d2, torch.cat([u2, m0], dim=1)))
        return self.mem4(self.up4(d3, u1))


class IntrinsicHead(nn.Module):
    def __init__(self, out_ch=OUT_CH, width_mult=1.0):
        super().__init__()
        C = _scale(width_mult)
        self.head = DoubleConv(C(64) + 3, C(32))
        self.head_out = nn.Conv2d(C(32), out_ch, 3, padding=1,
                                  padding_mode="reflect")

    def forward(self, d4, x):
        return self.head_out(self.head(torch.cat([d4, x], dim=1)))


class DictTail(nn.Module):
    """Patch-dictionary prior (blends internally, deblocked) -> tanh."""

    def __init__(self, out_ch=OUT_CH, deblock=True, dict_atoms=1024):
        super().__init__()
        self.dict = PatchDictionaryBias(channels=out_ch, patch=16,
                                        n_atoms=dict_atoms, deblock=deblock)

    def forward(self, logits):
        return torch.tanh(self.dict(logits))


class ParallelUNet(nn.Module):
    def __init__(self, in_ch=3, out_ch=OUT_CH, width_mult=WIDTH_MULT,
                 bias_res=384, deblock=True, dict_atoms=1024):
        super().__init__()
        self.encoder = VanillaEncoder(in_ch, width_mult)
        self.mobile = MobileEncoder()
        self.bottleneck = FusionBottleneck(width_mult)
        self.decoder = FusedDecoder(width_mult)
        self.head = IntrinsicHead(out_ch, width_mult)
        self.tail = DictTail(out_ch, deblock=deblock, dict_atoms=dict_atoms)

    def forward(self, x):
        u1, u2, u3, u4, u5, u6 = self.encoder(x)
        m0, m1, m2, m3, m4 = self.mobile(x)
        b = self.bottleneck(u6, m4)
        d4 = self.decoder(b, (u1, u2, u3, u4, u5), (m0, m1, m2, m3))
        return self.tail(self.head(d4, x))


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

    Strips the `generator.` prefix from Lightning checkpoints and applies the
    EMA shadow 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