# E-OS v4 — FrameGenerator: FULL-FRAME generative renderer (~0.92M params).
# Input : hint [B,3,192,416] (rough layout, /255)  +  state [B,32]
# Output: frame [B,3,384,832] in (0,1) via sigmoid  (always 2x hint size)
#
# NPU-friendliness (NNAPI): only Conv2d / SiLU / Mul+Add broadcast (FiLM) /
# Concat / ResizeNearest / ResizeBilinear / Sigmoid — no norm layers, no dynamic
# shapes. Fully convolutional -> trains at half res, infers at full 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, semb=96):
        super().__init__()
        self.stem = nn.Conv2d(3, 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 res)
        self.b0   = nn.Conv2d(96, 96, 3, 1, 1)
        for i in range(6):
            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 + 3, 24, 3, 1, 1)     # +hint @ H/1
        self.o1   = nn.Conv2d(24 + 3, 16, 3, 1, 1)     # + bilinear hint @ H/1
        self.o2   = nn.Conv2d(16, 12, 3, 1, 1)
        self.px   = nn.PixelShuffle(2)                 # DepthToSpace -> 3 x 2H x 2W
        self.semb = nn.Sequential(nn.Linear(32, 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, 7):
            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


def count_macs(model, hint_hw):
    macs = [0]
    hooks = []
    def hook(mod, inp, out):
        macs[0] += mod.weight.numel() * out.shape[-2] * out.shape[-1]
    for mod in model.modules():
        if isinstance(mod, nn.Conv2d):
            hooks.append(mod.register_forward_hook(hook))
    h = torch.rand(1, 3, hint_hw[0], hint_hw[1])
    s = torch.rand(1, 32)
    with torch.no_grad():
        model(h, s)
    for hk in hooks:
        hk.remove()
    return macs[0]


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, 3, 96, 208), torch.rand(1, 32))
    print("train-res out:", tuple(y.shape))
    with torch.no_grad():
        y2 = m(torch.rand(1, 3, 192, 416), torch.rand(1, 32))
    print("app-res out:", tuple(y2.shape))
    print("app-res MACs: %.2fG" % (count_macs(m, (192, 416)) / 1e9))
    print("train-res MACs: %.2fG" % (count_macs(m, (96, 208)) / 1e9))