"""
Step 5 acceptance tests: the hybrid photonic computer, with BOTH repos running.

  (A) real CPU slice   -> loaded ADC8 is a bit-exact adder (8-bit and rippled 32-bit)
  (B) hybrid learns    -> optical + electronic net classifies two-moons > 90%
  (C) CPU does readout -> the real ADC8 executes the integer readout accumulation,
                          bit-exact vs a reference -- the electronic half is the CPU
"""

import math
import torch
from photonic.hybrid import HybridPhotonicNet
from photonic.cpu_bridge import ADC8

torch.manual_seed(0)


def make_moons(n=600, noise=0.15):
    n2 = n // 2
    t = torch.rand(n2) * math.pi
    outer = torch.stack([torch.cos(t), torch.sin(t)], 1)
    inner = torch.stack([1 - torch.cos(t), 1 - torch.sin(t) - 0.5], 1)
    X = torch.cat([outer, inner], 0) + noise * torch.randn(n, 2)
    y = torch.cat([torch.zeros(n2), torch.ones(n2)]).long()
    perm = torch.randperm(n)
    return X[perm], y[perm]


def test_cpu_slice():
    adc = ADC8()
    ok8 = all(adc.add(a, b, c) == ((a + b + c) & 0xFF)
              for a, b, c in [(int(torch.randint(0, 256, ()).item()),
                               int(torch.randint(0, 256, ()).item()),
                               int(torch.randint(0, 2, ()).item())) for _ in range(512)])
    ok32 = all(adc.add_word(a, b, 4) == ((a + b) & 0xFFFFFFFF)
               for a, b in [(int(torch.randint(0, 2**32, ()).item()),
                             int(torch.randint(0, 2**32, ()).item())) for _ in range(200)])
    print(f"(A) real ADC8 slice: 8-bit adds {'ok' if ok8 else 'FAIL'} (512), "
          f"rippled 32-bit adds {'ok' if ok32 else 'FAIL'} (200)  -> {'PASS' if ok8 and ok32 else 'FAIL'}")
    return ok8 and ok32


def test_hybrid_learns():
    Xtr, ytr = make_moons(600)
    Xte, yte = make_moons(400)
    net = HybridPhotonicNet(d_in=2, modes=8, n_classes=2, phase_bits=6)
    opt = torch.optim.Adam(net.parameters(), lr=0.02)
    for _ in range(1500):
        opt.zero_grad()
        loss = torch.nn.functional.cross_entropy(net(Xtr), ytr)
        loss.backward()
        opt.step()
    with torch.no_grad():
        acc = (net(Xte).argmax(1) == yte).float().mean().item()
    print(f"(B) hybrid photonic net: two-moons test accuracy = {acc*100:.1f}%  -> {'PASS' if acc > 0.90 else 'FAIL'}")
    return acc > 0.90, net, (Xte, yte)


def test_cpu_does_readout(net, data):
    """Execute the electronic readout accumulation on the REAL ADC8 slice.
    We accumulate the per-mode integer contributions to a class logit using the
    neural CPU adder, and check it is bit-exact vs a plain integer sum."""
    adc = ADC8()
    Xte, _ = data
    with torch.no_grad():
        # quantize the head's contribution for one sample/class to bytes
        x = Xte[:1]
        field = net.optical(net.enc(x).to(torch.complex64))
        inten = (field.abs() ** 2)
        h = torch.relu(inten * net.gain + net.bias)[0]
        contrib = (h * net.head.weight[0]).clamp(-1, 1)      # per-mode logit terms
        bytes_ = torch.round((contrib + 1) * 127).int().tolist()  # map to [0,254]
    # accumulate with the verified slice (mod 256) vs reference
    acc_cpu, acc_ref = 0, 0
    for v in bytes_:
        acc_cpu = adc.add(acc_cpu, v)
        acc_ref = (acc_ref + v) & 0xFF
    ok = acc_cpu == acc_ref
    print(f"(C) CPU executes readout: sum of {len(bytes_)} mode-bytes  "
          f"ADC8={acc_cpu}  ref={acc_ref}  -> {'PASS' if ok else 'FAIL'}")
    return ok


if __name__ == "__main__":
    print("=" * 60)
    print("STEP 5  --  hybrid photonic computer (optics + neural CPU)")
    print("=" * 60)
    r_a = test_cpu_slice()
    r_b, net, data = test_hybrid_learns()
    r_c = test_cpu_does_readout(net, data)
    print("-" * 60)
    print(f"RESULT: {sum([r_a, r_b, r_c])}/3 acceptance tests passed")