"""CPU check of the DMAD critic wiring (fp32, tiny latent).

    python tests/test_dmad.py MODEL_DIR LODESTONES_DIR
Generator phase (critic frozen): gradients reach the generator and nothing reaches the critic.
Critic phase: gradients reach the critic head and its LoRA B on blocks >= start only.
"""
import sys

import torch
from torch.nn import functional as F

import rapid_moss.model as M
from rapid_moss.nets import Critic, 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)
tr = release.MOSSImageTransformer.from_pretrained(f'{model_dir}/transformer', torch_dtype=torch.float32,
                                                  local_files_only=True)
model = M.Moss(tr, release)
ctx = model.context(torch.randn(10, 4096) * 2)
z = torch.randn(1, 128, 4, 6)

gen = Generator(model, 128, 128, 4)
gen.load_lodestones(lode_dir)
critic = Critic(model, 16, 16, start=30, hidden=64)
assert critic.lora.A['q'].shape[0] == 12

# generator phase
critic.requires_grad_(False)
x0 = gen.step(model, z, gen.rollout(model, z, ctx, 1), 1, ctx, ckpt=True)
logit = critic(model, 0.6 * x0 + 0.4 * torch.randn_like(z), 400.0, ctx, ckpt=True)
assert logit.shape == (1,)
(-logit.mean()).backward()
assert all(p.grad is not None and p.grad.abs().sum() > 0 for p in gen.parameters()), 'generator grads missing'
assert all(p.grad is None for p in critic.parameters()), 'critic got generator-phase grads'
print('generator phase OK, logit', logit.item())

# critic phase
critic.requires_grad_(True)
for p in gen.parameters():
    p.grad = None
d_q = critic(model, 0.5 * torch.randn_like(z) + 0.5 * z, 300.0, ctx, ckpt=True)
d_g = critic(model, 0.5 * x0.detach() + 0.5 * torch.randn_like(z), 300.0, ctx, ckpt=True)
(F.softplus(-d_q) + F.softplus(d_g)).mean().backward()
assert all(p.grad is None for p in gen.parameters()), 'generator got critic-phase grads'
assert all(p.grad.abs().sum() > 0 for p in critic.head.parameters()) and critic.out.weight.grad.abs().sum() > 0
b_rows = (critic.lora.B['q'].grad.flatten(1).abs().sum(1) > 0).tolist()
print('critic LoRA B grad per block (30..41):', b_rows)
assert all(b_rows)
print('OK')