# E-OS v6 "LIQUID GLASS" — FrameGenerator (~1.0M params).
# Input : hint [B,10,H,W]  (4 shape-class masks + 2x ripple dx/dy/strength)
#         state [B,12]     (theme one-hot/hue, shimmer, day, pressed)
# Output: frame [B,3,2H,2W] in (0,1) via sigmoid  (always 2x hint size)
#
# NPU-friendliness (NNAPI): only Conv2d / SiLU / Mul+Add broadcast (FiLM) /
# Concat / ResizeNearest / Sigmoid / PixelShuffle(=DepthToSpace) — no norm
# layers, no dynamic shapes. Fully convolutional -> trains at half res.
import torch
import torch.nn as nn
import torch.nn.functional as F


class FiLM(nn.Module):
    """scale-shift conditioning: y = x*(1+g) + b  (g,b from state embedding)"""
    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 FrameGenerator(nn.Module):
    def __init__(self, hint_ch=12, state_dim=12, semb=96):
        super().__init__()
        self.stem = nn.Conv2d(hint_ch, 32, 3, 2, 1)    # 32 x H/2
        self.e1   = nn.Conv2d(32, 48, 3, 2, 1)         # 48 x H/4
        self.e2   = nn.Conv2d(48, 96, 3, 2, 1)         # 96 x H/8 (bottleneck)
        self.b0   = nn.Conv2d(96, 96, 3, 1, 1)
        for i in range(7):
            setattr(self, "b%d" % (i + 1), nn.Conv2d(96, 96, 3, 1, 1))
        self.d3   = nn.Conv2d(96 + 48, 96, 3, 1, 1)    # +skip e1 -> H/4
        self.d2   = nn.Conv2d(96 + 32, 48, 3, 1, 1)    # +skip stem -> H/2
        self.d1   = nn.Conv2d(48 + hint_ch, 32, 3, 1, 1)   # +hint @ H/1
        self.o1   = nn.Conv2d(32 + hint_ch, 20, 3, 1, 1)   # +hint @ H/1
        self.o2   = nn.Conv2d(20, 12, 3, 1, 1)
        self.px   = nn.PixelShuffle(2)                 # -> 3 x 2H x 2W
        self.semb = nn.Sequential(nn.Linear(state_dim, 64), nn.SiLU(),
                                  nn.Linear(64, semb))
        self.film_b = FiLM(96, 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))
        y = torch.sigmoid(y)
        return y


if __name__ == "__main__":
    m = FrameGenerator()
    m.eval()
    n = sum(p.numel() for p in m.parameters())
    print("params: %.3fM" % (n / 1e6))
    with torch.no_grad():
        y = m(torch.rand(1, 12, 128, 272), torch.rand(1, 12))
    print("train-res out:", tuple(y.shape))
    with torch.no_grad():
        y2 = m(torch.rand(1, 12, 256, 544), torch.rand(1, 12))
    print("app-res out:", tuple(y2.shape))
    macs = [0]
    hooks = []
    def hook(mod, inp, out):
        macs[0] += mod.weight.numel() * out.shape[-2] * out.shape[-1]
    for mod in m.modules():
        if isinstance(mod, nn.Conv2d):
            hooks.append(mod.register_forward_hook(hook))
    with torch.no_grad():
        m(torch.rand(1, 12, 256, 544), torch.rand(1, 12))
    for hk in hooks:
        hk.remove()
    print("app-res MACs: %.2fG" % (macs[0] / 1e9))