"""
Step 4: train the neural photonic unit and verify it EXHAUSTIVELY (N/N),
then save it in the same .pt form as the neural-aarch64 slice units.

Success criterion (identical in spirit to the aarch64 units): the MLP must
reproduce the golden quantized MZI matrix for ALL 256 phase configurations,
with zero mismatches -- bit-identical to the reference.
"""

import torch
from photonic.unit import (
    NeuralPhotonicMZI, enumerate_domain, verify_exhaustive,
    PHASE_BITS, OUT_BITS,
)

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


def train():
    X, Y = enumerate_domain()
    X = X.to(dev)
    q = 2 ** (OUT_BITS - 1) - 1
    Yf = (Y.float() / q).to(dev)                    # golden as reals in [-1,1]

    unit = NeuralPhotonicMZI().to(dev)
    opt = torch.optim.Adam(unit.parameters(), lr=2e-3)
    sched = torch.optim.lr_scheduler.CosineAnnealingLR(opt, T_max=20000)

    best = -1
    for epoch in range(20000):
        opt.zero_grad()
        pred = unit.forward_real(X)
        loss = (pred - Yf).pow(2).mean()
        loss.backward()
        opt.step()
        sched.step()
        if epoch % 2000 == 0 or epoch == 19999:
            n_ok, n_tot, mae = verify_exhaustive(unit)
            best = max(best, n_ok)
            print(f"  epoch {epoch:5d}  loss {loss.item():.2e}  "
                  f"verified {n_ok}/{n_tot}  max|err| {mae}")
            if n_ok == n_tot:
                print("  -> reached full N/N verification, stopping early")
                break
    return unit


if __name__ == "__main__":
    print("=" * 60)
    print("STEP 4  --  verified neural photonic unit  (MZI2)")
    print(f"  domain: {2**PHASE_BITS} x {2**PHASE_BITS} = {2**(2*PHASE_BITS)} configs, "
          f"output {OUT_BITS}-bit fixed point")
    print("=" * 60)

    unit = train()
    n_ok, n_tot, mae = verify_exhaustive(unit)
    print("-" * 60)
    status = "VERIFIED (N/N)" if n_ok == n_tot else "NOT verified"
    print(f"RESULT: {n_ok}/{n_tot} inputs bit-exact, max|err|={mae}  ->  {status}")

    if n_ok == n_tot:
        path = "MZI2.pt"
        torch.save({
            "state_dict": unit.state_dict(),
            "meta": {
                "unit": "MZI2",
                "role": "phase settings (4b theta, 4b phi) -> quantized 2x2 MZI matrix",
                "phase_bits": PHASE_BITS,
                "out_bits": OUT_BITS,
                "verified": f"{n_ok}/{n_tot} (exhaustive)",
            },
        }, path)
        print(f"  saved verified unit -> {path}  (slots beside ADC8.pt, DECODE_BITMASK.pt ...)")