"""Train the neural photodetector and verify it exhaustively (N/N), then save PD.pt."""

import torch
from photonic.detector import (
    NeuralPhotodetector, enumerate_domain, verify_exhaustive, RE_BITS, IM_BITS,
)

torch.manual_seed(0)
dev = "cuda" if torch.cuda.is_available() else "cpu"

X, Y = enumerate_domain()
X, Yf = X.to(dev), (Y / 255.0).to(dev)
unit = NeuralPhotodetector().to(dev)
opt = torch.optim.Adam(unit.parameters(), lr=2e-3)
sched = torch.optim.lr_scheduler.CosineAnnealingLR(opt, T_max=45000)

print("=" * 56)
print(f"NEURAL PHOTODETECTOR  domain=2^{RE_BITS+IM_BITS}={2**(RE_BITS+IM_BITS)} configs, 8-bit output")
print("=" * 56)
for epoch in range(45000):
    opt.zero_grad()
    loss = (unit.forward_frac(X) - Yf).pow(2).mean()
    loss.backward(); opt.step(); sched.step()
    if epoch % 2000 == 0 or epoch == 44999:
        n_ok, n_tot, mae = verify_exhaustive(unit.cpu()); unit.to(dev)
        print(f"  epoch {epoch:5d}  loss {loss.item():.2e}  verified {n_ok}/{n_tot}  max|err| {mae}")
        if n_ok == n_tot:
            print("  -> full N/N verification reached")
            break

unit = unit.cpu()
n_ok, n_tot, mae = verify_exhaustive(unit)
print("-" * 56)
status = "VERIFIED (N/N)" if n_ok == n_tot else "NOT verified"
print(f"RESULT: {n_ok}/{n_tot} bit-exact, max|err|={mae}  ->  {status}")
if n_ok == n_tot:
    torch.save({"state_dict": unit.state_dict(),
                "meta": {"unit": "PD", "role": "complex field (5b re, 5b im) -> intensity byte",
                         "verified": f"{n_ok}/{n_tot} (exhaustive)"}}, "PD.pt")
    print("  saved -> PD.pt")