"""CPU check of rapid_moss.model against the release MOSSImageTransformer.

    python tests/test_equiv.py MODEL_DIR LODESTONES_DIR
Runs everything in fp32 (rapid_moss.model.BF16 patched) so the logic must match to ~1e-5:
plain forward, negative-branch skip 22, the lodestones generator path (fp32 time path +
unmerged LoRA hooks, as its model card describes), discriminator taps and a gradient check.
"""
import sys
import time

import torch
from safetensors.torch import load_file

import rapid_moss.model as M
from rapid_moss.nets import Discriminator, FakeScore, Generator

torch.manual_seed(0)
model_dir, lode_dir = sys.argv[1], sys.argv[2]
M.BF16 = torch.float32
release = M.load_release_module(model_dir)
t0 = time.time()
tr = release.MOSSImageTransformer.from_pretrained(f'{model_dir}/transformer', torch_dtype=torch.float32,
                                                  local_files_only=True)
model = M.Moss(tr, release)
print(f'load {time.time() - t0:.1f}s')

h, w, L = 6, 10, 12
latent = torch.randn(1, 128, h, w)
emb = torch.randn(L, 4096) * 2
ctx = model.context(emb)


def rel(a, b):
    return ((a - b).norm() / b.norm()).item()


with torch.no_grad():
    for t in (999.0, 500.0, 37.0):
        ref = tr(latent, t, emb[None]).sample.float()
        got = model(latent, t, ctx)
        print(f'forward t={t}: rel {rel(got, ref):.2e}')
        assert rel(got, ref) < 1e-4
    ref = tr(latent, 600.0, emb[None], skip_layers=[22]).sample.float()
    got = model(latent, 600.0, ctx, skip=(22,))
    print(f'skip22: rel {rel(got, ref):.2e} (vs no-skip {rel(got, model(latent, 600.0, ctx)):.2e})')
    assert rel(got, ref) < 1e-4

    # lodestones generator: release model + replaced time path + LoRA hooks (model card recipe)
    gen = Generator(model, 128, 128, 4)
    gen.load_lodestones(lode_dir)
    te = load_file(f'{lode_dir}/time_embedding.safetensors')
    lora = load_file(f'{lode_dir}/lora.safetensors')
    import copy
    tr2 = copy.deepcopy(tr)
    tr2.load_state_dict(te, strict=False)
    for i, layer in enumerate(tr2.layers):
        for name, lin in layer.proj.items():
            p = f'diffusion_model.layers.{i}.proj.{name}'
            down, up, alpha = lora[p + '.lora_down.weight'], lora[p + '.lora_up.weight'], lora[p + '.alpha']
            scale = float(alpha) / down.shape[0]
            lin.register_forward_hook(lambda m, inp, out, d=down, u=up, s=scale: out + s * (inp[0] @ d.t()) @ u.t())
    Q = load_file(f'{lode_dir}/dataspace.safetensors')['latent_mix.Q']
    codes = load_file(f'{lode_dir}/dataspace.safetensors')['step_codes']
    z = torch.randn(1, 128, h, w)
    x0_ref = torch.zeros_like(z)
    x0_got = torch.zeros_like(z)
    for r in range(3):
        x = z + torch.einsum('oc,bchw->bohw', Q, x0_ref)
        x0_ref = z - tr2(x, codes[r], emb[None]).sample.float()
        x0_got = gen.step(model, z, x0_got, r, ctx)
        print(f'generator r={r}: rel {rel(x0_got, x0_ref):.2e}')
        assert rel(x0_got, x0_ref) < 1e-4
    del tr2

    disc = Discriminator(model.dim, (5, 11), 256)
    feats = model(latent, 300.0, ctx, taps=disc.taps)
    print('taps', {k: tuple(v.shape) for k, v in feats.items()}, 'logits', tuple(disc(feats, h, w).shape))

# gradients: generator step through checkpointed blocks, fake score, D input-gradient
fake = FakeScore(model, 16, 16, 4)
x0 = gen.step(model, z, torch.randn_like(z), 1, ctx, ckpt=True)
x0.backward(torch.randn_like(x0))
nz = {n: p.grad.abs().sum().item() > 0 for n, p in gen.named_parameters()}
print('gen grads nonzero:', all(nz.values()), [n for n, v in nz.items() if not v][:5])
assert all(nz.values())
v = fake(model, latent, 400.0, ctx, 2, ckpt=True)
v.square().mean().backward()
print('fake r_emb grad rows nonzero:', (fake.r_emb.grad.abs().sum(1) > 0).tolist())
leaf = latent.clone().requires_grad_(True)
g, = torch.autograd.grad(-disc(model(leaf, 300.0, ctx, taps=disc.taps, ckpt=True), h, w).mean(), leaf)
print('D input grad norm', g.norm().item())
assert g.norm() > 0
print('OK')