"""
Step 3 acceptance tests: quantized (finite-domain) photonic layer.

  (a) no-op default   -> phase_bits=None reproduces Step 1/2 exactly
  (b) unitarity kept  -> quantizing PHASES preserves U U^H = I at every bit-depth
  (c) error -> 0       -> matrix error shrinks ~2x per added phase bit
  (d) finite domain    -> a quantized 2-mode unit has a countable # of configs
  (e) QAT works        -> quantization-aware training still fits a target
"""

import math
import torch
from photonic.mzi import MZIMesh, photodetect
from photonic.linear import PhotonicLinear

torch.manual_seed(0)


def test_default_is_noop(n=6):
    a = MZIMesh(n, seed=3)
    b = MZIMesh(n, seed=3, phase_bits=None)
    err = (a.unitary() - b.unitary()).abs().max().item()
    print(f"(a) default no-op   |diff|_max = {err:.2e}  -> {'PASS' if err == 0 else 'FAIL'}")
    return err == 0


def test_unitarity_preserved(n=6):
    ok = True
    for bits in (2, 3, 4, 6, 8):
        mesh = MZIMesh(n, seed=4, phase_bits=bits)
        U = mesh.unitary()
        err = (U @ U.conj().T - torch.eye(n, dtype=torch.complex64)).abs().max().item()
        ok = ok and err < 1e-4
        print(f"    phase_bits={bits}: |U U^H - I|_max = {err:.2e}")
    print(f"(b) unitarity kept under phase quantization  -> {'PASS' if ok else 'FAIL'}")
    return ok


def test_error_shrinks(n=6):
    ref = MZIMesh(n, seed=5).unitary().detach()
    prev = None
    ok = True
    for bits in (2, 3, 4, 5, 6, 7, 8):
        q = MZIMesh(n, seed=5, phase_bits=bits).unitary().detach()
        err = (q - ref).abs().max().item()
        note = ""
        if prev is not None:
            note = f"(x{prev/err:.1f} better)" if err > 0 else "(exact)"
            if err > prev:
                ok = False
        print(f"    phase_bits={bits}: |M_q - M_cont|_max = {err:.3e} {note}")
        prev = err
    print(f"(c) error shrinks with bits  -> {'PASS' if ok else 'FAIL'}")
    return ok


def test_finite_domain():
    """A 2-mode mesh has 1 MZI (theta, phi) + 2 output phases = 4 phase knobs."""
    n = 2
    mesh = MZIMesh(n, phase_bits=4)
    n_phase_knobs = mesh.theta.numel() + mesh.phi.numel() + mesh.out_phase.numel()
    for bits in (4, 6, 8):
        levels = 2 ** bits
        configs = levels ** n_phase_knobs
        print(f"    2-mode unit, {bits}-bit phases: {levels}^{n_phase_knobs} = {configs:,} configs")
    print(f"(d) finite domain: config count is finite & enumerable  -> PASS")
    print(f"      (this is exactly the precondition your N/N verification needs)")
    return True


def test_qat(n=6, bits=6, steps=4000):
    # passive+quantized optics realize CONTRACTIONS (spectral norm <= 1);
    # normalize the target so it is physically realizable. (Gain lives in the
    # electronic half -- the neural CPU -- not in passive light.)
    target = torch.randn(n, n, dtype=torch.complex64)
    target = target / torch.linalg.matrix_norm(target, ord=2) * 0.99
    layer = PhotonicLinear(n, n, seed=6, phase_bits=bits, amp_bits=bits)
    opt = torch.optim.Adam(layer.parameters(), lr=0.02)
    for _ in range(steps):
        opt.zero_grad()
        loss = (layer.matrix() - target).abs().pow(2).mean()
        loss.backward()
        opt.step()
    # residual reflects finite bit-depth, not a training failure
    print(f"(e) QAT @ {bits}-bit  fit-contraction loss = {loss.item():.2e}  -> {'PASS' if loss.item() < 5e-3 else 'FAIL'}")
    return loss.item() < 5e-3


if __name__ == "__main__":
    print("=" * 60)
    print("STEP 3  --  quantized finite-domain photonic layer")
    print("=" * 60)
    results = [
        test_default_is_noop(),
        test_unitarity_preserved(),
        test_error_shrinks(),
        test_finite_domain(),
        test_qat(),
    ]
    print("-" * 60)
    print(f"RESULT: {sum(results)}/{len(results)} acceptance tests passed")