"""
Step 1 acceptance tests for the differentiable MZI mesh.

Proves the mesh is a real, coherent, differentiable optical matmul:
  (a) it is unitary                 -> U U^dagger = I (energy conserved)
  (b) it equals a plain matmul      -> mesh(x) == U @ x
  (c) it interferes                 -> one input mode spreads across outputs
  (d) it is trainable               -> fit a target unitary, loss -> 0
"""

import torch
from photonic.mzi import MZIMesh, photodetect

torch.manual_seed(0)


def test_unitary(n=6):
    mesh = MZIMesh(n)
    U = mesh.unitary()
    I = torch.eye(n, dtype=torch.complex64)
    err = (U @ U.conj().T - I).abs().max().item()
    print(f"(a) unitarity   |U U^H - I|_max = {err:.2e}  -> {'PASS' if err < 1e-4 else 'FAIL'}")
    return err < 1e-4


def test_matches_matmul(n=6):
    mesh = MZIMesh(n)
    x = torch.randn(n, dtype=torch.complex64)
    y_forward = mesh(x)
    y_matmul = mesh.unitary() @ x
    err = (y_forward - y_matmul).abs().max().item()
    print(f"(b) is a matmul |mesh(x) - U x|_max = {err:.2e}  -> {'PASS' if err < 1e-5 else 'FAIL'}")
    # energy conservation through the multiply:
    ein = photodetect(x).sum().item()
    eout = photodetect(y_forward).sum().item()
    print(f"      input energy {ein:.4f}  ->  output energy {eout:.4f}  (conserved)")
    return err < 1e-5


def test_interference(n=6):
    """Inject light into ONE mode; a non-mixing device would keep it in one
    output, interference spreads it across many."""
    mesh = MZIMesh(n)
    x = torch.zeros(n, dtype=torch.complex64)
    x[0] = 1.0
    out = photodetect(mesh(x))
    occupied = (out > 1e-3).sum().item()
    print(f"(c) interference: 1 input mode -> {occupied}/{n} output modes lit")
    print("      output intensities:", [f"{v:.3f}" for v in out.tolist()])
    return occupied > 1


def test_trainable(n=6, steps=1500):
    """Fit the mesh to a random TARGET unitary -- proves the optical settings
    are differentiable and the mesh is universal."""
    q, _ = torch.linalg.qr(torch.randn(n, n, dtype=torch.complex64))
    target = q  # a random unitary

    mesh = MZIMesh(n, seed=1)
    opt = torch.optim.Adam(mesh.parameters(), lr=0.05)
    for step in range(steps):
        opt.zero_grad()
        U = mesh.unitary()
        loss = (U - target).abs().pow(2).mean()
        loss.backward()
        opt.step()
    final = loss.item()
    print(f"(d) trainable   fit-a-target-unitary loss = {final:.2e}  -> {'PASS' if final < 1e-3 else 'FAIL'}")
    return final < 1e-3


if __name__ == "__main__":
    print("=" * 60)
    print("STEP 1  --  differentiable MZI mesh (the optical 'compute')")
    print("=" * 60)
    results = [
        test_unitary(),
        test_matches_matmul(),
        test_interference(),
        test_trainable(),
    ]
    print("-" * 60)
    print(f"RESULT: {sum(results)}/{len(results)} acceptance tests passed")