"""
Honest benchmark of the photonic core.

Three DIFFERENT speeds, kept strictly separate so a simulation number can never
masquerade as a hardware number (same discipline as the rest of the project):

  1. SIM (rebuild)  -- run the layer the way training does: assemble the mesh
                       unitary from MZIs every call. This is SIMULATION OVERHEAD.
                       It is slower than a plain matmul by construction.
  2. DEPLOY (frozen)-- freeze the trained mesh to its matrix, then it's a normal
                       complex matmul at full GPU speed. This is how you'd ship it.
  3. DIGITAL base   -- torch matmul of the same shape, for reference.

Plus an ANALYTIC projection of what real photonic HARDWARE would achieve --
explicitly a model, not a measurement.
"""

import time
import torch
from photonic.mzi import MZIMesh

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


def timed(fn, iters=50, warmup=10):
    for _ in range(warmup):
        fn()
    if dev == "cuda":
        torch.cuda.synchronize()
    t0 = time.perf_counter()
    for _ in range(iters):
        fn()
    if dev == "cuda":
        torch.cuda.synchronize()
    return (time.perf_counter() - t0) / iters


def bench(N=32, batch=1024):
    print(f"device={dev}  N={N} modes  batch={batch}")
    mesh = MZIMesh(N, seed=0).to(dev)
    x = torch.randn(batch, N, dtype=torch.complex64, device=dev)

    # 1. SIM: rebuild the mesh unitary every forward (training-time path)
    t_sim = timed(lambda: mesh(x))

    # 2. DEPLOY: freeze to a matrix once, then just multiply
    U = mesh.unitary().detach()
    t_deploy = timed(lambda: x @ U.T)

    # 3. DIGITAL baseline: same-shape complex matmul
    W = torch.randn(N, N, dtype=torch.complex64, device=dev)
    t_dig = timed(lambda: x @ W.T)

    macs = batch * N * N                       # complex MACs per forward
    print(f"  1. SIM  (rebuild mesh) : {t_sim*1e3:8.3f} ms/call   "
          f"{macs/t_sim/1e9:7.2f} G-MAC/s   [simulation overhead]")
    print(f"  2. DEPLOY (frozen mat) : {t_deploy*1e3:8.3f} ms/call   "
          f"{macs/t_deploy/1e9:7.2f} G-MAC/s   [ship this]")
    print(f"  3. DIGITAL baseline    : {t_dig*1e3:8.3f} ms/call   "
          f"{macs/t_dig/1e9:7.2f} G-MAC/s")
    print(f"  -> sim is {t_sim/t_deploy:.0f}x slower than frozen deploy "
          f"(that cost is simulating optics, not the optics)")
    return N


def hardware_projection(N=32):
    """MODEL, not a measurement. A real N-mode mesh does N*N MACs in ONE optical
    pass. Latency ~ time of flight through N layers; throughput set by the
    modulation rate. Numbers below are order-of-magnitude with stated assumptions."""
    print("\nANALYTIC HARDWARE PROJECTION (a model, NOT measured):")
    mzi_len_um = 200.0                         # typical MZI stage length
    n_group = 4.2                              # group index in silicon
    c = 3e8
    latency_s = (N * mzi_len_um * 1e-6) * n_group / c
    mod_rate = 50e9                            # 50 GHz modulators -> vectors/s
    macs_per_pass = N * N
    throughput = macs_per_pass * mod_rate
    print(f"  assumptions: {mzi_len_um:.0f}um/stage, n_g={n_group}, {mod_rate/1e9:.0f}GHz modulation")
    print(f"  latency (time of flight, {N} layers) : {latency_s*1e12:7.1f} ps")
    print(f"  throughput (1 matvec / modulation)   : {throughput/1e12:7.1f} T-MAC/s  ({N}x{N} per pass)")
    print(f"  energy: dominated by O(N) modulators, NOT O(N^2) MACs -- the win is")
    print(f"          that the N*N multiply is 'free' in the optical pass.")
    print("  NOTE: real chips lose to insertion loss, phase drift, DAC/ADC energy,")
    print("        and limited precision. This is the ceiling, not a promise.")


if __name__ == "__main__":
    print("=" * 64)
    print("PHOTONIC CORE BENCHMARK")
    print("=" * 64)
    for N in (16, 32, 64):
        bench(N=N)
        print()
    hardware_projection(N=64)