# E-OS v5 — FrameGenerator: Glass-Renderer (~1.6M params).
# Input : hint [B,3,HH,HW] (Wireframe-Layout) + state [B,32]
# Output: Glass-Frame [B,3,2H,2W] sigmoid — lernt Apple-Glas + Wasser-Ripples.
# Voll-konvolutional: Training 128x272 -> App 256x544 (Out 512x1088).
# NPU-freundlich: Conv2d/SiLU/FiLM/Concat/NearestUpsample/PixelShuffle/Sigmoid.
import torch
import torch.nn as nn
import torch.nn.functional as F


class FiLM(nn.Module):
    def __init__(self, ch, semb):
        super().__init__()
        self.g = nn.Linear(semb, ch)
        self.b = nn.Linear(semb, ch)

    def forward(self, x, s):
        g = self.g(s).unsqueeze(-1).unsqueeze(-1)
        b = self.b(s).unsqueeze(-1).unsqueeze(-1)
        return x * (1.0 + g) + b


def up2(x):
    return F.interpolate(x, scale_factor=2.0, mode="nearest")


class FrameGeneratorV5(nn.Module):
    def __init__(self, semb=96):
        super().__init__()
        self.stem = nn.Conv2d(3, 40, 3, 2, 1)          # 40 x H/2
        self.e1 = nn.Conv2d(40, 64, 3, 2, 1)           # 64 x H/4
        self.e2 = nn.Conv2d(64, 128, 3, 2, 1)          # 128 x H/8
        self.b0 = nn.Conv2d(128, 128, 3, 1, 1)
        for i in range(7):
            setattr(self, "b%d" % (i + 1), nn.Conv2d(128, 128, 3, 1, 1))
        self.d3 = nn.Conv2d(128 + 64, 96, 3, 1, 1)     # H/4
        self.d2 = nn.Conv2d(96 + 40, 56, 3, 1, 1)      # H/2
        self.d1 = nn.Conv2d(56 + 3, 28, 3, 1, 1)       # H
        self.o1 = nn.Conv2d(28 + 3, 18, 3, 1, 1)
        self.o2 = nn.Conv2d(18, 12, 3, 1, 1)
        self.px = nn.PixelShuffle(2)
        self.semb = nn.Sequential(nn.Linear(32, 64), nn.SiLU(), nn.Linear(64, semb))
        self.film_b = FiLM(128, semb)
        self.film_d = FiLM(96, semb)
        for m in self.modules():
            if isinstance(m, nn.Conv2d):
                nn.init.kaiming_normal_(m.weight, nonlinearity="relu")
                if m.bias is not None:
                    nn.init.zeros_(m.bias)
            if isinstance(m, nn.Linear):
                nn.init.kaiming_normal_(m.weight, nonlinearity="relu")
                nn.init.zeros_(m.bias)

    def forward(self, hint, state):
        s = self.semb(state)
        x0 = F.silu(self.stem(hint))
        x1 = F.silu(self.e1(x0))
        b = F.silu(self.e2(x1))
        b = F.silu(self.b0(b))
        b = self.film_b(b, s)
        for i in range(1, 8):
            b = F.silu(getattr(self, "b%d" % i)(b))
        u3 = F.silu(self.d3(torch.cat([up2(b), x1], 1)))
        u3 = self.film_d(u3, s)
        u2 = F.silu(self.d2(torch.cat([up2(u3), x0], 1)))
        u1 = F.silu(self.d1(torch.cat([up2(u2), hint], 1)))
        y = F.silu(self.o1(torch.cat([u1, hint], 1)))
        y = self.px(self.o2(y))
        return torch.sigmoid(y)


def grad_loss(y, t):
    dyx = y[..., :, 1:] - y[..., :, :-1]
    dtx = t[..., :, 1:] - t[..., :, :-1]
    dyy = y[..., 1:, :] - y[..., :-1, :]
    dty = t[..., 1:, :] - t[..., :-1, :]
    return (dyx - dtx).abs().mean() + (dyy - dty).abs().mean()


if __name__ == "__main__":
    m = FrameGeneratorV5().eval()
    n = sum(p.numel() for p in m.parameters())
    print("params: %.3fM" % (n / 1e6))
    with torch.no_grad():
        y = m(torch.rand(1, 3, 128, 272), torch.rand(1, 32))
    print("train out:", tuple(y.shape))
    with torch.no_grad():
        y = m(torch.rand(1, 3, 256, 544), torch.rand(1, 32))
    print("app out:", tuple(y.shape))
    macs = [0]

    def hook(mod, inp, out):
        macs[0] += mod.weight.numel() * out.shape[-2] * out.shape[-1]
    hs = []
    for mod in m.modules():
        if isinstance(mod, nn.Conv2d):
            hs.append(mod.register_forward_hook(hook))
    with torch.no_grad():
        m(torch.rand(1, 3, 256, 544), torch.rand(1, 32))
    for h in hs:
        h.remove()
    print("app MACs: %.2fG" % (macs[0] / 1e9))