"""T140 Codex author audit: full physical models, exact channels, real budgets.

Not independent peer review or a replication of the original ion apparatus.
"""
import csv
import itertools
import json
from pathlib import Path

import numpy as np
from scipy.integrate import solve_ivp
from scipy.sparse import bmat, csr_matrix
from scipy.sparse.linalg import expm_multiply

import matched_protection_budgets as m

checks = []


def ck(name, ok, **extra):
    checks.append({"name": name, "passed": bool(ok), **extra})
    assert ok, (name, extra)


def err(name, error, tol=2e-9):
    ck(name, error < tol, error=float(error), tolerance=tol)


def vec(B):
    return B.reshape(1024, -1, order="F")


def mats(v):
    return v.reshape(32, 32, -1, order="F")


def full_initial():
    B = []
    for a, b in [(0, 0), (1, 1), (0, 1)]:
        q = np.zeros((2, 2), complex); q[a, b] = 1
        B.append(np.kron(np.eye(4)/4, np.kron(q, np.eye(4)/4)))
    return np.stack(B, axis=2)


def full_noise(B, kx, kz, t):
    L = m.full_generator(kx, kz)
    return mats(expm_multiply(t*L, vec(B), traceA=t*L.diagonal().sum()))


def explicit_pauli(B, q, p):
    ps = [np.eye(2), np.array([[0, 1], [1, 0]]),
          np.array([[0, -1j], [1j, 0]]), np.diag([1, -1])]
    ops = [np.eye(2), np.eye(2), np.eye(2)]
    ops[q] = ps[p]
    U = np.kron(ops[2], np.kron(ops[1], np.kron(ops[0], np.eye(4))))
    return np.stack([U@B[:, :, k]@U.conj().T for k in range(B.shape[2])], axis=2)


def physical_gate(B, kind, kx, kz, p2=.001, p3=.002):
    t, p = (.04, p3) if kind == "TOF" else (.02, p2)
    B = full_noise(B, kx, kz, t)
    perm = m.PERMS[kind]
    U = np.eye(32)[perm]
    B = np.stack([U@B[:, :, k]@U.conj().T for k in range(3)], axis=2)
    support = [0, 1] if kind == "CX01" else ([0, 2] if kind == "CX02" else [0, 1, 2])
    C = (1-p)*B
    for labels in itertools.product(range(4), repeat=len(support)):
        if not any(labels):
            continue
        term = B.copy()
        for q, label in zip(support, labels):
            term = explicit_pauli(term, q, label)
        C += p/(4**len(support)-1)*term
    return C


def physical_prep(kx, kz, tau, active):
    L = m.full_generator(kx, kz)
    sm = np.array([[0, 1], [0, 0]], complex)
    counts_data = np.zeros((32, 32), complex)
    for q in [1, 2]:
        ops = [np.eye(2), np.eye(2), np.eye(2)]; ops[q] = sm
        J = np.kron(ops[2], np.kron(ops[1], np.kron(ops[0], np.eye(4))))
        L += m.PREP_RATE*m.dissipator(J)
        counts_data += m.PREP_RATE*J.conj().T@J
    if active:
        L += m.PREP_RATE*m.LA
    flux = csr_matrix(np.vstack([counts_data.T.reshape(-1, order="F"),
        m.PREP_RATE*m.FLUX if active else np.zeros(1024)]))
    A = bmat([[L, csr_matrix((1024, 2))], [flux, csr_matrix((2, 2))]], format="csr")
    y0 = np.vstack([vec(full_initial()), np.zeros((2, 3))])
    y = expm_multiply(tau*A, y0, traceA=tau*A.diagonal().sum())
    B = mats(y[:1024])
    B = physical_gate(physical_gate(B, "CX01", kx, kz), "CX02", kx, kz)
    return B, y[1024:, :2].real.mean(axis=1)


for d, ids in zip([0, 7], m.SECTORS):
    ck(f"sector_{d}_dimension128", len(ids) == 128 and len(set(ids)) == 128)
    outside = np.array(sorted(set(range(1024))-set(ids)))
    for name, L in zip(["X", "Z", "coupling", "reset"], [m.LX, m.LZ, m.LH, m.LA]):
        block = L[outside][:, ids]
        err(f"sector_{d}_{name}_exact_invariance", np.max(abs(block.data)) if block.nnz else 0.)

rng = np.random.default_rng(140)
test = rng.normal(size=(32, 32))+1j*rng.normal(size=(32, 32))
jumps = [np.sqrt(.2)*np.kron(a, np.eye(4)) for a in m.XS]
jumps += [np.sqrt(.02)*np.kron(a, np.eye(4)) for a in m.ZS]
for i in [1, 2, 3]:
    a = np.zeros((4, 4)); a[0, i] = 1
    jumps.append(np.sqrt(20)*np.kron(np.eye(8), a))
direct = -1j*5*(m.H@test-test@m.H)
for J in jumps:
    A = J.conj().T@J
    direct += J@test@J.conj().T-(A@test+test@A)/2
predicted = (m.full_generator(.2, .02, 5, 20)@test.reshape(-1, order="F")).reshape(32, 32, order="F")
err("full_generator_manual_matrix_RHS", np.max(abs(predicted-direct)))

for active in [False, True]:
    for tau in [.02, .1, .2]:
        B, events = physical_prep(.2, .02, tau, active)
        p, c, data_events, aux_events, _ = m.prepared(.2, .02, tau, active)
        err(f"finite_preparation_full_physical_active{active}_t{tau}", np.max(abs(B-m.expand(p, c))))
        err(f"finite_preparation_data_flux_active{active}_t{tau}", abs(events[0]-data_events))
        err(f"finite_preparation_aux_flux_active{active}_t{tau}", abs(events[1]-aux_events))
for name in ["CX01", "CX02", "TOF"]:
    B = full_initial()
    direct = physical_gate(B, name, .2, .02)
    pred = m.gate(B, name, .2, .02)
    err(name+"_explicit_all_Pauli_gate_branches", np.max(abs(direct-pred)))

# Construct all four correctable error classes, then explicitly decode by gates.
B = np.zeros((32, 32, 3), complex)
for k, (a, b) in enumerate([(0, 0), (7, 7), (0, 7)]):
    B[4*a, 4*b, k] = 1
for bit in [0, 1, 2, 4]:
    U = np.eye(32) if bit == 0 else np.kron(m.XS[[1, 2, 4].index(bit)], np.eye(4))
    C = np.stack([U@B[:, :, k]@U.conj().T for k in range(3)], axis=2)
    C = m.decoder(C, 0., 0., 0., 0.)
    metrics, _ = m.output_channel(C)
    err(f"coherent_decoder_all_single_X_bit{bit}", abs(metrics["Fe"]-1))

# Independent dense ODE on a positive physical state, including actual reset flux.
p, c, *_ = m.prepared(.2, .02, .1, True)
B = m.expand(p, c); rho0 = (B[:, :, 0]+B[:, :, 1])/2
def rhs(_, y):
    rho = y[:1024].reshape(32, 32, order="F")
    drho = -1j*5*(m.H@rho-rho@m.H)
    for J in jumps:
        A = J.conj().T@J
        drho += J@rho@J.conj().T-(A@rho+rho@A)/2
    return np.r_[drho.reshape(-1, order="F"), 20*np.trace(m.PEXC@rho)]
sol = solve_ivp(rhs, (0, .23), np.r_[rho0.reshape(-1, order="F"), 0.], rtol=2e-10, atol=2e-12)
ck("independent_dense_ODE_success", sol.success)
pp, cc, flow = m.evolve(p, c, .2, .02, 5., 20., .23)
pred = m.expand(pp, cc)
err("dense_ODE_vs_reduced_density", np.max(abs(sol.y[:1024, -1].reshape(32, 32, order="F")-(pred[:, :, 0]+pred[:, :, 1])/2)), 2e-8)
err("dense_ODE_vs_integrated_reset_events", abs(sol.y[1024, -1].real-flow), 2e-8)
pp, cc, flow = m.free_segment(p, c, .2, .02, 20., .23)
full_L = m.full_generator(.2, .02, 0., 20.)
full = mats(expm_multiply(.23*full_L, vec(B), traceA=.23*full_L.diagonal().sum()))
err("finite_reset_Kraus_vs_full_Lindblad", np.max(abs(m.expand(pp, cc)-full)))
err("finite_reset_actual_event_formula", abs(flow-m.controller_excitation(B)*(1-np.exp(-20*.23))))

r = json.loads((m.OUT/"results.json").read_text())
for key, ex in r["examples"].items():
    for kind, protocol in ex["protocols"].items():
        a = protocol["metrics"]
        B, _ = physical_prep(a["kx"], a["kz"], a["prep"], kind != "encoded_no_repair")
        reset_flow = 0.
        for seg in protocol["segments"]:
            t = seg["t"]
            L = m.full_generator(a["kx"], a["kz"], seg["g"], seg["gamma"])
            flux = csr_matrix((seg["gamma"]*m.FLUX)[None, :])
            A = bmat([[L, csr_matrix((1024, 1))], [flux, csr_matrix((1, 1))]], format="csr")
            y = expm_multiply(t*A, np.vstack([vec(B), np.zeros((1, 3))]),
                              traceA=t*A.diagonal().sum())
            B = mats(y[:1024]); reset_flow += float(y[1024, :2].real.mean())
        core_exc = m.controller_excitation(B) if kind != "encoded_no_repair" else 0.
        for g in ["CX02", "CX01", "TOF"]:
            B = physical_gate(B, g, a["kx"], a["kz"])
        metric, choi = m.output_channel(B)
        saved = np.load(m.OUT/f"{key}-{kind}.npz")
        err(key+"_"+kind+"_full_physical_output_channel", np.max(abs(choi-saved["choi"])))
        err(key+"_"+kind+"_full_reset_events", abs(reset_flow-a["core_aux_reset_events"]))
        err(key+"_"+kind+"_controller_excitation", abs(core_exc-a["final_aux_excitation"]))
        g_used = sum(seg["g"]**2*seg["t"] for seg in protocol["segments"])
        r_used = sum(seg["gamma"]*seg["t"] for seg in protocol["segments"])
        r_used += m.PREP_RATE*a["prep"] if kind != "encoded_no_repair" else 0.
        err(key+"_"+kind+"_coupling_capacity", abs(g_used-a["coupling_budget_used"]))
        err(key+"_"+kind+"_reset_capacity_includes_preparation", abs(r_used-a["aux_reset_budget_used"]))
        # Six cardinal pure states are a 2-design: evaluate the output explicitly.
        E0, E1, E01 = 2*choi[:2, :2], 2*choi[2:, 2:], 2*choi[:2, 2:]
        fidelities = []
        for psi in [np.array([1, 0]), np.array([0, 1]), np.array([1, 1])/np.sqrt(2),
                    np.array([1, -1])/np.sqrt(2), np.array([1, 1j])/np.sqrt(2),
                    np.array([1, -1j])/np.sqrt(2)]:
            rho = np.outer(psi, psi.conj())
            output = rho[0, 0]*E0+rho[1, 1]*E1+rho[0, 1]*E01+rho[1, 0]*E01.conj().T
            fidelities.append(np.vdot(psi, output@psi).real)
        err(key+"_"+kind+"_average_fidelity_six_actual_inputs",
            abs(np.mean(fidelities)-(2*metric["Fe"]+1)/3))

for W, kz in itertools.product(m.DEADLINES, m.Z_RATES):
    px, pz = (1-np.exp(-.4*W))/2, (1-np.exp(-2*kz*W))/2
    ps = [m.I, m.X, m.Z, m.X@m.Z]
    weights = [(1-px)*(1-pz), px*(1-pz), (1-px)*pz, px*pz]
    bell = np.array([1, 0, 0, 1])/np.sqrt(2)
    choi = sum(w*np.outer(np.kron(m.I, P)@bell, (np.kron(m.I, P)@bell).conj())
               for w, P in zip(weights, ps))
    err(f"bare_actual_complete_channel_W{W}_Z{kz}", abs(np.vdot(bell, choi@bell).real-m.bare(W, kz=kz)))
    A = rng.normal(size=(2, 2))+1j*rng.normal(size=(2, 2))
    U, _ = np.linalg.qr(A)
    transformed_Fe = sum(w*abs(np.trace(U.conj().T@P@U))**2/4 for w, P in zip(weights, ps))
    err(f"bare_static_encoding_basis_invariance_W{W}_Z{kz}", abs(transformed_Fe-m.bare(W, kz=kz)))
for P in [m.X, m.Z]:
    for Q in [m.X, m.Z, m.X@m.Z]:
        transformed = Q.conj().T@P@Q
        err("Pauli_echo_does_not_change_white_noise_"+str(len(checks)),
            np.max(abs(m.dissipator(transformed).toarray()-m.dissipator(P).toarray())))

def read_rows(name):
    out = []
    for raw in csv.DictReader((m.OUT/name).open()):
        row = {}
        for k, v in raw.items():
            if v == "":
                row[k] = None
            else:
                try: row[k] = float(v)
                except ValueError: row[k] = v
        out.append(row)
    return out

rows = read_rows("protocol-catalogue.csv")
by_id = {int(a["id"]): a for a in rows}
fronts = read_rows("matched-budget-envelopes.csv")
bad_physics, bad_resource = [], []
for a in rows:
    if not (-1e-10 <= a["Fe"] <= 1+1e-10 and a["TP_error"] < 2e-9 and
            a["Choi_min_eigenvalue"] > -2e-9 and abs(a["channel_trace"]-1)<2e-9):
        bad_physics.append(int(a["id"]))
    if not (a["data_reset_integral"] <= 16+1e-9 and a["g"] <= 10+1e-9 and
            a["gamma"] <= 40+1e-9 and a["prep"]+.12 <= a["deadline"]+1e-9 and
            a["core_aux_reset_events"] >= -1e-9 and
            a["unconditional_delivery_probability"] == 1):
        bad_resource.append(int(a["id"]))
ck("all_catalogue_channels_physical", not bad_physics, rows=len(rows), failures=bad_physics[:10])
ck("all_catalogue_peak_capacity_and_deadlines", not bad_resource, rows=len(rows), failures=bad_resource[:10])
bad_envelopes = []
for f in fronts:
    choices = [a for a in rows if a["deadline"] == f["deadline"] and a["kz"] == f["kz"] and
               a["coupling_budget_used"] <= f["coupling_cap"]+1e-9 and
               a["aux_reset_budget_used"] <= f["aux_reset_cap"]+1e-9]
    maximum = max([m.bare(f["deadline"], kz=f["kz"])]+[a["Fe"] for a in choices])
    if abs(maximum-f["best_Fe"]) > 1e-10:
        bad_envelopes.append(f)
    for kind in ["encoded_no_repair", "continuous", "switched"]:
        feasible = [a for a in choices if a["kind"] == kind]
        expected = max((a["Fe"] for a in feasible), default=None)
        got = f[kind+"_Fe"]
        if (expected is None) != (got is None) or (expected is not None and abs(expected-got)>1e-10):
            bad_envelopes.append(f)
ck("all_budget_envelopes_equal_feasible_catalogue_maxima", not bad_envelopes,
   budgets=len(fronts), failures=len(bad_envelopes))
resource_fronts = read_rows("quality-resource-frontiers.csv")
bad_targets = []
for point in resource_fronts:
    a = by_id[int(point["id"])]
    if a["Fe"] < point["target_Fe"]-1e-10:
        bad_targets.append(point)
    choices = [b for b in rows if b["deadline"] == point["deadline"] and b["kz"] == point["kz"] and
               b["kind"] == point["kind"] and b["Fe"] >= point["target_Fe"]]
    if any(b["coupling_budget_used"] <= a["coupling_budget_used"] and
           b["aux_reset_budget_used"] < a["aux_reset_budget_used"]-1e-10 for b in choices):
        bad_targets.append(point)
ck("quality_resource_frontier_points_meet_target_and_are_nondominated", not bad_targets,
   frontier_points=len(resource_fronts), failures=len(bad_targets))

result = {"by": "Codex", "type": "author-self-review", "passed": all(c["passed"] for c in checks),
          "count": len(checks), "checks": checks,
          "scope": "full physical Lindblad/Pauli branches/dense ODE, preparation/reset flux, actual output channels and capacity checks",
          "limitations": ["finite declared protocol search, not a proof of all-control optimality",
                          "one-message task, no repeated-service throughput or reset-free reuse guarantee",
                          "control integrals are capacity conventions, not total apparatus energy",
                          "not independent review or replication of the source ion device"]}
(m.OUT/"audit.json").write_text(json.dumps(result, ensure_ascii=False, indent=2)+"\n")
print(json.dumps({"passed": result["passed"], "count": result["count"]}))