#!/usr/bin/env python3
"""Train a TWO-AGENT co-control interactive world model and export it INTEGER-exact.
The transition consumes BOTH agents' actions per step (A=4 planes: dxA,dyA,dxB,dyB),
so two participants drive two distinct entities in ONE world and co-render the same
frame — true simultaneous multiplayer, byte-exact. Paired with wai.world.model's
signed action trace (each actor = a slot), the merged log renders the IDENTICAL world
on every peer (convergence), and here both agents move at once rather than serializing.

Same integer discipline as the 1-agent world (wai_int_world_export): transition T +
synthesis g_s are int_hyper IntModels (WHM1), the action folded in as constant input
planes; the whole decode is i64. Reuses that tool's WHM1 writer + PTQ + integer decode.

  python tools/wai_int_world2_export.py <out-dir>
"""
import sys, os, json
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 run_int_net, AFRAC, GF, SHIFT
from wai_int_world_export import whm1, ptq_conv, ptq_convT, ptq_igdn, gs_decode, fnv1a, GDN2d

FRAME, LAT, C, A = 64, 8, 24, 4        # A=4: two agents × (dx,dy)
ONE = 1 << AFRAC
SPEED, R = 5, 4.5                       # agent step (px) / marker radius
ACTIONS = [(0, 0), (1, 0), (-1, 0), (0, 1), (0, -1)]
NROLL, ROLL, STEPS = 48, 8, 2500
COLA, COLB = np.array([1.0, 0.25, 0.25]), np.array([0.25, 1.0, 0.35])   # agent tints


def make_bg():
    rng = np.random.RandomState(7)
    lo = torch.tensor(rng.rand(1, 3, FRAME // 8, FRAME // 8).astype(np.float32))
    bg = F.interpolate(lo, size=(FRAME, FRAME), mode='bilinear', align_corners=True)[0].numpy()
    return (bg * 0.45).astype(np.float32)                # dim background so markers pop


YY, XX = np.mgrid[0:FRAME, 0:FRAME]


def render(bg, pa, pb):
    f = bg.copy()
    for pos, col in ((pa, COLA), (pb, COLB)):
        m = np.exp(-(((XX - pos[0]) ** 2 + (YY - pos[1]) ** 2) / (2 * R * R)))
        for c in range(3):
            f[c] = np.clip(f[c] * (1 - m) + col[c] * m, 0, 1)
    return f.astype(np.float32)


def clampp(p):
    return [min(max(p[0], R), FRAME - R), min(max(p[1], R), FRAME - R)]


def rollout(bg, rng, pairs=None):
    pa, pb = [FRAME * 0.3, FRAME * 0.5], [FRAME * 0.7, FRAME * 0.5]
    frames = [render(bg, pa, pb)]
    chosen = []
    n = len(pairs) if pairs is not None else ROLL
    for t in range(n):
        aa, ab = pairs[t] if pairs is not None else (rng.randint(5), rng.randint(5))
        pa = clampp([pa[0] + ACTIONS[aa][0] * SPEED, pa[1] + ACTIONS[aa][1] * SPEED])
        pb = clampp([pb[0] + ACTIONS[ab][0] * SPEED, pb[1] + ACTIONS[ab][1] * SPEED])
        frames.append(render(bg, pa, pb))
        chosen.append((aa, ab))
    return np.stack(frames), chosen


def planes_t(aa, ab):                                   # [A,LAT,LAT] float
    p = torch.zeros(A, LAT, LAT)
    p[0], p[1] = ACTIONS[aa]
    p[2], p[3] = ACTIONS[ab]
    return p


class World2(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, pairs):
        ap = torch.stack([planes_t(aa, ab) for aa, ab in pairs]).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):
        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)


def planes_int(aa, ab):                                 # [dxA,dyA,dxB,dyB] @ 2^afrac
    return [ACTIONS[aa][0] * ONE, ACTIONS[aa][1] * ONE, ACTIONS[ab][0] * ONE, ACTIONS[ab][1] * ONE]


def decode_world2(T, gs, z0, pairs):
    z = z0.copy()
    frames = [gs_decode(gs, z)]
    for aa, ab in pairs:
        row = np.array(planes_int(aa, ab), dtype=np.int64)[:, None, None] * np.ones((A, LAT, LAT), np.int64)
        z = np.clip(run_int_net(T, np.concatenate([z, row], 0)), -ONE, ONE)
        frames.append(gs_decode(gs, z))
    return frames


def main(outdir):
    os.makedirs(outdir, exist_ok=True)
    torch.manual_seed(0)
    bg = make_bg()
    rng = np.random.RandomState(123)
    data = [rollout(bg, rng) for _ in range(NROLL)]
    frames = torch.tensor(np.stack([d[0] for d in data]))
    acts = [d[1] for d in data]

    m = World2()
    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() / ((m(frames, acts) - frames) ** 2).sum() + 1e-12)

    T_layers = [ptq_conv(m.Tc[0]), {'type': 'leaky', 'slope_q': 410, 'slope_frac': 12}, 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}

    # committed co-control walk: A circles, B zig-zags (both move most ticks)
    demo = [(1, 4), (1, 2), (3, 1), (3, 1), (2, 4), (2, 3), (4, 1), (0, 2)]
    df, _ = rollout(bg, None, pairs=demo)
    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_world2(T_model, gs_model, z0_int, demo)
    gt = (df.transpose(0, 2, 3, 1) * 255).astype(np.float64)
    snr_i = 10 * np.log10((gt ** 2).sum() / ((np.stack(frames_int).astype(np.float64) - gt) ** 2).sum() + 1e-12)
    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}/world2_T.bin', 'wb').write(whm1(T_layers))
    open(f'{outdir}/world2_gs.bin', 'wb').write(whm1(gs_layers))
    json.dump({'c': C, 'lat': LAT, 'frame': FRAME, 'a': A, 'afrac': AFRAC, 'agents': 2,
               'z0': z0_int.reshape(-1).tolist(), 'pairs': demo,
               'frame_hashes': hashes, 'clip_hash': clip_hash},
              open(f'{outdir}/world2.json', 'w'))
    np.stack(frames_int).astype(np.uint8).tofile(f'{outdir}/world2_frames.rgb')
    print(f'2-agent world: {len(frames_int)} frames {FRAME}x{FRAME}, T {len(whm1(T_layers))}B + g_s '
          f'{len(whm1(gs_layers))}B, float {float(snr_f):.1f} / integer {snr_i:.1f} dB')
    print(f'  clip_hash {clip_hash}  frame0 {hashes[0]}')


if __name__ == '__main__':
    main(sys.argv[1] if len(sys.argv) > 1 else '/tmp/int_world2')