"""
Step 2 acceptance tests: arbitrary photonic linear layer (M = U Sigma V^H).

  (a) staged == matrix   -> physical rotate/attenuate/rotate equals M @ x
  (b) universal          -> fit an arbitrary complex target matrix, loss -> 0
  (c) rectangular        -> works for out != in
  (d) passive is passive -> passive=True never amplifies (||y|| <= ||x||)
"""

import torch
from photonic.linear import PhotonicLinear
from photonic.mzi import photodetect

torch.manual_seed(0)


def test_staged_equals_matrix(n=6):
    layer = PhotonicLinear(n, n)
    x = torch.randn(n, dtype=torch.complex64)
    err = (layer(x, staged=True) - layer(x, staged=False)).abs().max().item()
    print(f"(a) staged==matrix  |diff|_max = {err:.2e}  -> {'PASS' if err < 1e-4 else 'FAIL'}")
    return err < 1e-4


def test_universal(n=6, steps=3000):
    target = torch.randn(n, n, dtype=torch.complex64)   # arbitrary matrix
    layer = PhotonicLinear(n, n, seed=2)
    opt = torch.optim.Adam(layer.parameters(), lr=0.03)
    for _ in range(steps):
        opt.zero_grad()
        loss = (layer.matrix() - target).abs().pow(2).mean()
        loss.backward()
        opt.step()
    print(f"(b) universal       fit-arbitrary-matrix loss = {loss.item():.2e}  -> {'PASS' if loss.item() < 1e-3 else 'FAIL'}")
    return loss.item() < 1e-3


def test_rectangular(out_m=4, in_m=7):
    layer = PhotonicLinear(in_m, out_m)
    x = torch.randn(in_m, dtype=torch.complex64)
    y = layer(x)
    ok = tuple(y.shape) == (out_m,)
    err = (layer(x, staged=True) - y).abs().max().item()
    print(f"(c) rectangular     {in_m}->{out_m} shape {tuple(y.shape)}, staged err {err:.2e}  -> {'PASS' if ok and err < 1e-4 else 'FAIL'}")
    return ok and err < 1e-4


def test_passive(n=6, trials=200):
    layer = PhotonicLinear(n, n, passive=True)
    worst = 0.0
    for _ in range(trials):
        x = torch.randn(n, dtype=torch.complex64)
        ratio = photodetect(layer(x)).sum().sqrt() / photodetect(x).sum().sqrt()
        worst = max(worst, ratio.item())
    print(f"(d) passive         max energy gain over {trials} inputs = {worst:.4f}  -> {'PASS' if worst <= 1.0 + 1e-4 else 'FAIL'}")
    return worst <= 1.0 + 1e-4


if __name__ == "__main__":
    print("=" * 60)
    print("STEP 2  --  arbitrary photonic linear layer (SVD mesh)")
    print("=" * 60)
    results = [
        test_staged_equals_matrix(),
        test_universal(),
        test_rectangular(),
        test_passive(),
    ]
    print("-" * 60)
    print(f"RESULT: {sum(results)}/{len(results)} acceptance tests passed")