#!/usr/bin/env python3
"""Train a small action-conditioned INTERACTIVE WORLD MODEL and export it as an
INTEGER-exact decoder, so that conformance is a hash. A world model that generates
the next frame from an action through a FLOAT network gives a world that is not
reproducible from one platform to another. Here the world model is integer:

  frame_0 = g_s(z_0);   z_{t+1} = clip( T( z_t ‖ action_t ) );   frame_{t+1} = g_s(z_{t+1})

with BOTH the transition T and the synthesis g_s exported as int_hyper IntModels
(WHM1: per-out-channel int12 weights, fixed-point IGDN/leaky, all i64). The action is
folded in as constant input planes. So given (model, z_0, action trace) every frame is
byte-identical on every machine — a generative interactive world you can put in a
receipt. Paired with wai.world.model's signed/replay-equivalent action trace, both the
trace AND the frames hash-converge.

  python tools/wai_int_world_export.py <out-dir>
"""
import sys, os, json, struct
import numpy as np, torch, torch.nn as nn, torch.nn.functional as F, warnings
warnings.filterwarnings('ignore')
sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
from wai_sota_int_synth_export import (q_weight_pc, q_weight_conv, conv_int, conv_t_raw,
                                       igdn_int, run_int_net, AFRAC, GF, SHIFT)

FRAME, LAT, C, A = 64, 8, 24, 3      # frame px / latent px / channels / action planes
WORLD, SPEED = 96, 4                  # global world image + camera step (px/action)
NROLL, ROLL = 40, 8                   # training rollouts / steps per rollout
STEPS = 2500                          # optimizer steps (CLI-overridable for scale runs)
ONE = 1 << AFRAC
ACTIONS = [(0, 0), (1, 0), (-1, 0), (0, 1), (0, -1)]  # stay, right, left, down, up
SLOPE_Q, SLOPE_FRAC = 410, 12        # leaky 0.1 in fixed point (410/4096)


# ---------- procedural world + environment (the "ground truth" dynamics) ----------
def make_world():
    rng = np.random.RandomState(7)
    lo = torch.tensor(rng.rand(1, 3, WORLD // 8, WORLD // 8).astype(np.float32))
    w = F.interpolate(lo, size=(WORLD, WORLD), mode='bilinear', align_corners=True)[0].numpy()
    yy, xx = np.mgrid[0:WORLD, 0:WORLD]
    w[0] = np.clip(w[0] + 0.25 * np.sin(xx / 6.0), 0, 1)          # structure the model must learn
    w[1] = np.clip(w[1] + 0.25 * np.cos(yy / 7.0), 0, 1)
    w[2] = np.clip(w[2] + 0.20 * ((xx // 12 + yy // 12) % 2), 0, 1)
    return w.astype(np.float32)


def crop(world, cx, cy):
    x0 = int(round(cx)) - FRAME // 2
    y0 = int(round(cy)) - FRAME // 2
    x0 = min(max(x0, 0), WORLD - FRAME)
    y0 = min(max(y0, 0), WORLD - FRAME)
    return world[:, y0:y0 + FRAME, x0:x0 + FRAME]


def rollout(world, rng, acts=None):
    cx = cy = WORLD / 2.0
    frames = [crop(world, cx, cy)]
    chosen = []
    for t in range(len(acts) if acts is not None else ROLL):
        a = acts[t] if acts is not None else rng.randint(len(ACTIONS))
        dx, dy = ACTIONS[a]
        cx = min(max(cx + dx * SPEED, FRAME / 2), WORLD - FRAME / 2)
        cy = min(max(cy + dy * SPEED, FRAME / 2), WORLD - FRAME / 2)
        frames.append(crop(world, cx, cy))
        chosen.append(a)
    return np.stack(frames), chosen          # frames[ROLL+1,3,64,64], acts[ROLL]


def action_planes(a):                         # [A,8,8] float at unit scale
    dx, dy = ACTIONS[a]
    p = torch.zeros(A, LAT, LAT)
    p[0] = dx; p[1] = dy; p[2] = 0.0
    return p


# ---------- learned model (float; the encoder/trainer) ----------
class GDN2d(nn.Module):
    def __init__(self, ch):
        super().__init__()
        self.rg = nn.Parameter(torch.randn(ch, ch) * 0.01)
        self.rb = nn.Parameter(torch.zeros(ch))

    def gb(self):
        return F.softplus(self.rg), F.softplus(self.rb + 1.0)

    def forward(self, x):                      # inverse GDN: x * sqrt(beta + gamma·x²)
        g, b = self.gb()
        norm = b[None, :, None, None] + torch.einsum('ij,njhw->nihw', g, x * x)
        return x * torch.sqrt(norm)


class World(nn.Module):
    def __init__(self):
        super().__init__()
        self.E = nn.Sequential(nn.Conv2d(3, C, 4, 2, 1), nn.ReLU(),
                               nn.Conv2d(C, C, 4, 2, 1), nn.ReLU(),
                               nn.Conv2d(C, C, 4, 2, 1))
        self.Tc = nn.ModuleList([nn.Conv2d(C + A, C, 3, 1, 1), nn.Conv2d(C, C, 3, 1, 1)])
        self.G = nn.ModuleList([nn.ConvTranspose2d(C, C, 4, 2, 1), GDN2d(C),
                                nn.ConvTranspose2d(C, C, 4, 2, 1), GDN2d(C),
                                nn.ConvTranspose2d(C, 3, 4, 2, 1)])

    def trans(self, z, a):
        ap = torch.stack([action_planes(ai) for ai in a]).to(z.device)
        x = F.leaky_relu(self.Tc[0](torch.cat([z, ap], 1)), 0.1)
        return torch.clamp(self.Tc[1](x), -1, 1)

    def synth(self, z):
        x = z
        for m in self.G:
            x = m(x)
        return x

    def forward(self, frames, acts):           # frames[B,T+1,3,64,64], acts[B][T]
        z = self.E(frames[:, 0])
        preds = [self.synth(z)]
        for t in range(ROLL):
            z = self.trans(z, [a[t] for a in acts])
            preds.append(self.synth(z))
        return torch.stack(preds, 1)


# ---------- WHM1 writer (matches int_hyper_synth::from_bytes) ----------
def whm1(layers):
    out = bytearray(b'WHM1') + struct.pack('<IIII', AFRAC, GF, SHIFT, len(layers))
    for L in layers:
        if L['type'] == 'convT':
            out += struct.pack('<BIIIIII', 0, L['in_ch'], L['out_ch'], L['k'], L['stride'], L['pad'], L['opad'])
            out += np.array(L['weight'], dtype='<i2').tobytes()
            out += np.array(L['mult'], dtype='<i8').tobytes() + np.array(L['bias'], dtype='<i8').tobytes()
        elif L['type'] == 'conv':
            out += struct.pack('<BIIIII', 1, L['in_ch'], L['out_ch'], L['k'], L['stride'], L['pad'])
            out += np.array(L['weight'], dtype='<i2').tobytes()
            out += np.array(L['mult'], dtype='<i8').tobytes() + np.array(L['bias'], dtype='<i8').tobytes()
        elif L['type'] == 'igdn':
            out += struct.pack('<BI', 2, L['c'])
            out += np.array(L['gamma'], dtype='<i4').tobytes() + np.array(L['beta'], dtype='<i8').tobytes()
        elif L['type'] == 'leaky':
            out += struct.pack('<B', 4) + np.array([L['slope_q']], dtype='<i8').tobytes() + struct.pack('<I', L['slope_frac'])
    return bytes(out)


def ptq_conv(layer):                            # nn.Conv2d → int dict
    wi, sc = q_weight_conv(layer.weight.detach().numpy())
    return {'type': 'conv', 'in_ch': wi.shape[1], 'out_ch': wi.shape[0], 'k': wi.shape[2],
            'stride': layer.stride[0], 'pad': layer.padding[0], 'weight': wi.reshape(-1).tolist(),
            'mult': np.round(sc * (2 ** SHIFT)).astype(np.int64).tolist(),
            'bias': np.round(layer.bias.detach().numpy() * ONE).astype(np.int64).tolist()}


def ptq_convT(layer):
    wi, sc = q_weight_pc(layer.weight.detach().numpy())
    return {'type': 'convT', 'in_ch': wi.shape[0], 'out_ch': wi.shape[1], 'k': wi.shape[2],
            'stride': layer.stride[0], 'pad': layer.padding[0], 'opad': layer.output_padding[0],
            'weight': wi.reshape(-1).tolist(),
            'mult': np.round(sc * (2 ** SHIFT)).astype(np.int64).tolist(),
            'bias': np.round(layer.bias.detach().numpy() * ONE).astype(np.int64).tolist()}


def ptq_igdn(layer):
    with torch.no_grad():
        g, b = layer.gb()
    return {'type': 'igdn', 'c': b.shape[0],
            'gamma': np.round(g.numpy() * (2 ** GF)).astype(np.int64).reshape(-1).tolist(),
            'beta': np.round(b.numpy() * (2 ** (2 * AFRAC))).astype(np.int64).tolist()}


# ---------- integer reference decode (mirrors the Rust int_world loop) ----------
def planes_int(a):
    dx, dy = ACTIONS[a]
    p = np.zeros((A, LAT, LAT), np.int64)
    p[0] = dx * ONE; p[1] = dy * ONE
    return p


def gs_decode(gs, z):                            # z[C,8,8]@2^afrac → frame[64,64,3] u8
    x = run_int_net(gs, z)
    half = 1 << (AFRAC - 1)
    x = np.clip(x, 0, ONE)
    pix = ((x * 255 + half) >> AFRAC).clip(0, 255).astype(np.uint8)   # [3,64,64]
    return np.ascontiguousarray(pix.transpose(1, 2, 0))               # HWC


def fnv1a(buf):
    h = 14695981039346656037
    for b in buf:
        h = ((h ^ b) * 1099511628211) & 0xFFFFFFFFFFFFFFFF
    return h


def decode_world(T, gs, z0, acts):
    z = z0.copy()
    frames = [gs_decode(gs, z)]
    for a in acts:
        inp = np.concatenate([z, planes_int(a)], 0)                  # [C+A,8,8]
        z = run_int_net(T, inp)
        z = np.clip(z, -ONE, ONE)
        frames.append(gs_decode(gs, z))
    return frames


def main(outdir):
    os.makedirs(outdir, exist_ok=True)
    torch.manual_seed(0)
    world = make_world()
    rng = np.random.RandomState(123)
    # Half random walks, half OUT-AND-BACK (go some direction, then reverse) so the
    # model is trained toward long-horizon consistency: reversing a path should
    # return where you came from.
    data = []
    for i in range(NROLL):
        if i % 2 == 0:
            data.append(rollout(world, rng))
        else:
            half = ROLL // 2
            d = rng.randint(1, len(ACTIONS))            # a non-stay direction
            rev = [2, 1, 4, 3][d - 1]                   # its reverse (x↔, y↔)
            data.append(rollout(world, rng, acts=[d] * half + [rev] * (ROLL - half)))
    frames = torch.tensor(np.stack([d[0] for d in data]))
    acts = [d[1] for d in data]

    # GPU when available (torch's `mps` backend) — the free-roll recurrent backprop DEPTH is the
    # training cost, so the GPU is what unblocks a longer horizon (bigger ROLL).
    dev = 'mps' if torch.backends.mps.is_available() else ('cuda' if torch.cuda.is_available() else 'cpu')
    print(f'  training on {dev} (ROLL={ROLL}, C={C}, steps={STEPS})')
    m = World().to(dev)
    frames = frames.to(dev)
    opt = torch.optim.Adam(m.parameters(), 1.5e-3)
    for step in range(STEPS):
        pred = m(frames, acts)
        loss = F.mse_loss(pred, frames)
        opt.zero_grad(); loss.backward(); opt.step()
    m.eval()
    with torch.no_grad():
        snr_f = 10 * np.log10((frames ** 2).sum().item() / ((m(frames, acts) - frames) ** 2).sum().item() + 1e-12)
    m.cpu(); frames = frames.cpu()                     # back to CPU for PTQ + the integer reference decode

    # PTQ both nets → WHM1
    T_layers = [ptq_conv(m.Tc[0]), {'type': 'leaky', 'slope_q': SLOPE_Q, 'slope_frac': SLOPE_FRAC}, ptq_conv(m.Tc[1])]
    gs_layers = [ptq_convT(m.G[0]), ptq_igdn(m.G[1]), ptq_convT(m.G[2]), ptq_igdn(m.G[3]), ptq_convT(m.G[4])]
    T_model = {'afrac': AFRAC, 'gf': GF, 'shift': SHIFT, 'layers': T_layers}
    gs_model = {'afrac': AFRAC, 'gf': GF, 'shift': SHIFT, 'layers': gs_layers}

    # a longer demo walk (showcases the horizon): wander out and partway back
    demo_acts = [1, 1, 1, 3, 3, 1, 4, 2, 2, 2, 3, 4, 1, 0, 2, 4]
    df, _ = rollout(world, None, acts=demo_acts)
    with torch.no_grad():
        z0 = m.E(torch.tensor(df[None, 0]))[0].numpy()
    z0_int = np.round(z0 * ONE).astype(np.int64)

    frames_int = decode_world(T_model, gs_model, z0_int, demo_acts)
    # quality vs ground-truth crops
    gt = (df.transpose(0, 2, 3, 1) * 255).astype(np.float64)
    pr = np.stack(frames_int).astype(np.float64)
    snr_i = 10 * np.log10((gt ** 2).sum() / ((pr - gt) ** 2).sum() + 1e-12)

    # LONG-HORIZON CONSISTENCY probe: walk out and back the same distance; does the
    # integer world return to where it started? (reported honestly — consistency is
    # the hard problem for world models; we measure, not assume.)
    n = max(2, ROLL // 2)
    back = [1] * n + [2] * n                                          # +x n, then −x n
    bf = decode_world(T_model, gs_model, z0_int, back)
    ret_match = bool((bf[0] == bf[-1]).all())
    ret_snr = 10 * np.log10((bf[0].astype(np.float64) ** 2).sum()
                            / ((bf[-1].astype(np.float64) - bf[0].astype(np.float64)) ** 2).sum() + 1e-12)
    print(f'  walk-and-return ({n}+{n}): start==end byte-exact? {ret_match}  (return SNR {ret_snr:.1f} dB)')

    hashes = [format(fnv1a(f.tobytes()), '016x') for f in frames_int]
    clip_hash = format(fnv1a(b''.join(bytes.fromhex(h) for h in hashes)), '016x')

    open(f'{outdir}/world_T.bin', 'wb').write(whm1(T_layers))
    open(f'{outdir}/world_gs.bin', 'wb').write(whm1(gs_layers))
    json.dump({'c': C, 'lat': LAT, 'frame': FRAME, 'a': A, 'afrac': AFRAC,
               'z0': z0_int.reshape(-1).tolist(), 'actions': demo_acts,
               'frame_hashes': hashes, 'clip_hash': clip_hash},
              open(f'{outdir}/world.json', 'w'))
    # raw RGB frames for the demo (optional viewing)
    np.stack(frames_int).astype(np.uint8).tofile(f'{outdir}/world_frames.rgb')
    print(f'int world model: {len(frames_int)} frames {FRAME}x{FRAME}, T {len(whm1(T_layers))}B + g_s '
          f'{len(whm1(gs_layers))}B, float SNR {float(snr_f):.1f} dB / integer {snr_i:.1f} dB')
    print(f'  clip_hash {clip_hash}  frame0 {hashes[0]}')


if __name__ == '__main__':
    out = sys.argv[1] if len(sys.argv) > 1 else '/tmp/int_world'
    # optional full-scale overrides: FRAME LAT C ROLL WORLD (FRAME must = LAT·8)
    if len(sys.argv) > 2:
        FRAME = int(sys.argv[2])
    if len(sys.argv) > 3:
        LAT = int(sys.argv[3])
    if len(sys.argv) > 4:
        C = int(sys.argv[4])
    if len(sys.argv) > 5:
        ROLL = int(sys.argv[5])
    if len(sys.argv) > 6:
        WORLD = int(sys.argv[6])
    if len(sys.argv) > 7:
        STEPS = int(sys.argv[7])
    main(out)