"""T139 Codex author audit, not independent review.

Cross-check reduced instruments against full physical density matrices, hidden
environment blocks against an 8-dimensional joint quantum model, Frechet
derivatives against central differences, QFI against explicit noisy measurement
probabilities, Fock evolution against Kraus/ODE and photon balance.
"""
import csv
import itertools
import json
import math

import numpy as np
from scipy.integrate import solve_ivp
from scipy.linalg import expm, solve_sylvester

import error_corrected_sensing as m

checks = []


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


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


err("code_isometry", np.max(abs(m.V.conj().T@m.V-m.I)))
err("noise_KL_first_moment", np.max(abs(m.V.conj().T@m.ZP@m.V)))
err("noise_KL_second_moment", np.max(abs(m.V.conj().T@m.ZP@m.ZP@m.V-m.I)))
err("good_signal_survives", np.max(abs(m.V.conj().T@m.XP@m.V-m.Z)))
err("parallel_signal_compression_zero", np.max(abs(m.V.conj().T@m.ZP@m.V)))
err("complete_recovery", np.max(abs(sum(K.conj().T@K for K in m.RECOVERY)-np.eye(4))))


def direct_physical(omega=.4, delta=.23, latency=.02, n=4, nu=.2, p0=.5, signal="X"):
    ne = 1 if nu is None else 2
    dim = 4*ne
    G = m.XP if signal == "X" else m.ZP
    GG = np.kron(np.eye(ne), G)
    if ne == 1:
        jumps = [np.sqrt(.2)*m.ZP]
    else:
        e0, e1 = np.diag([1., 0.]), np.diag([0., 1.])
        hop = np.array([[0., 1.], [0., 0.]])
        jumps = [np.sqrt(.2)*np.kron(e0, m.ZP), np.sqrt(.2)*np.kron(e1, m.XP),
                 np.sqrt(nu)*np.kron(hop, np.eye(4)), np.sqrt(nu)*np.kron(hop.T, np.eye(4))]
    D = m.commutator(GG)
    Lnoise = sum(m.dissipator(j) for j in jumps)
    E, DE = m.channel_and_derivative(Lnoise+omega*D, D, delta)
    C = expm(latency*Lnoise)
    E, DE = C@E, C@DE
    psi = m.V@m.PLUS
    rho = np.outer(psi, psi.conj())
    if ne == 2:
        rho = np.kron(np.diag([p0, 1-p0]), rho)
    states = [(rho, np.zeros_like(rho))]
    for _ in range(n):
        evolved = []
        for a, da in states:
            evolved.append((m.mat(E@m.vec(a), dim), m.mat(E@m.vec(da)+DE@m.vec(a), dim)))
        new = []
        for K in m.RECOVERY:
            KK = np.kron(np.eye(ne), K)
            for a, da in evolved:
                new.append((KK@a@KK.conj().T, KK@da@KK.conj().T))
        states = new
    out, dout = [], []
    for a, da in states:
        out.append([m.vec(m.V.conj().T@a[4*k:4*k+4, 4*k:4*k+4]@m.V) for k in range(ne)])
        dout.append([m.vec(m.V.conj().T@da[4*k:4*k+4, 4*k:4*k+4]@m.V) for k in range(ne)])
    return np.array(out), np.array(dout)


for nu, initial in [(None, None), (.2, [.5, .5]), (0., [1., 0.]), (2., [.5, .5])]:
    for s in ["X", "Z"] if nu is None else ["X"]:
        J, DJ = m.reduced_instruments(m.XP if s == "X" else m.ZP, .4, .2, .23, .02, nu)
        a, da = m.propagate(J, DJ, 4, initial, False)
        b, db = direct_physical(nu=nu, p0=1. if initial == [1., 0.] else .5, signal=s)
        err(f"physical_joint_model_states_nu{nu}_{s}", np.max(abs(a-b)))
        err(f"physical_joint_model_derivatives_nu{nu}_{s}", np.max(abs(da-db)))
        eps = 1e-5
        Jp, Dp = m.reduced_instruments(m.XP if s == "X" else m.ZP, .4+eps, .2, .23, .02, nu)
        Jn, Dn = m.reduced_instruments(m.XP if s == "X" else m.ZP, .4-eps, .2, .23, .02, nu)
        ap, _ = m.propagate(Jp, Dp, 4, initial, False)
        an, _ = m.propagate(Jn, Dn, 4, initial, False)
        err(f"all_branch_parameter_derivatives_nu{nu}_{s}", np.max(abs((ap-an)/(2*eps)-da)), 2e-8)

for s in ["X", "Z"]:
    J, DJ = m.reduced_instruments(m.XP if s == "X" else m.ZP, .4, .2, .2, .02)
    a, da = m.propagate(J, DJ, 8, None, False)
    ac, dc = m.propagate(J, DJ, 8, None, True)
    count_a = np.zeros_like(ac); count_d = np.zeros_like(dc)
    for k in range(2**8):
        count_a[k.bit_count()] += a[k]
        count_d[k.bit_count()] += da[k]
    err(f"count_aggregation_states_{s}", np.max(abs(count_a-ac)))
    err(f"count_aggregation_derivatives_{s}", np.max(abs(count_d-dc)))
    full, counted = m.branch_metrics(a, da), m.branch_metrics(ac, dc)
    for key in ["CQFI", "record_FI", "actual_local_readout_FI", "record_ignored_QFI"]:
        err(f"count_aggregation_{s}_{key}", abs(full[key]-counted[key]))

# Commuting parallel signal: the logical state is independent of omega, but
# records retain nonzero LOCAL FI at omega0 != 0, with a global sign ambiguity.
for delta, n, latency in [(.2, 5, 0.), (.1, 10, .02), (.5, 4, .05)]:
    d = m.corrected("Z", delta=delta, n=n, latency=latency)
    damp = np.exp(-2*.2*(delta+latency))
    p = (1-damp*np.cos(2*.4*delta))/2
    dp = damp*delta*np.sin(2*.4*delta)
    analytic = n*dp*dp/(p*(1-p))
    err(f"parallel_record_analytic_delta{delta}", abs(d["record_FI"]-analytic))
    err(f"parallel_logical_erased_delta{delta}", d["record_ignored_QFI"])
    err(f"parallel_all_information_in_record_delta{delta}", abs(d["CQFI"]-analytic))

rng = np.random.default_rng(139)
for k in range(8):
    A = rng.normal(size=(2, 2))+1j*rng.normal(size=(2, 2))
    rho = A@A.conj().T; rho /= np.trace(rho)
    B = rng.normal(size=(2, 2))+1j*rng.normal(size=(2, 2))
    drho = (B+B.conj().T)/2; drho -= np.trace(drho)*m.I/2
    sld = solve_sylvester(rho, rho, 2*drho)
    err(f"SLD_linear_equation_QFI_{k}", abs(m.qfi(rho, drho)-np.trace(rho@sld@sld).real))
    for flip in [0., .02, .2]:
        fi, probs, deriv, basis = m.local_readout(rho, drho, flip)
        raw = np.real(np.diag(basis.conj().T@rho@basis))
        draw = np.real(np.diag(basis.conj().T@drho@basis))
        err(f"explicit_binary_flip_probabilities_{k}_{flip}", np.max(abs(probs-((1-flip)*raw+flip*raw[::-1]))))
        err(f"explicit_binary_flip_derivatives_{k}_{flip}", np.max(abs(deriv-(1-2*flip)*draw)))
        err(f"readout_FI_equals_depolarized_QFI_{k}_{flip}", abs(fi-m.qfi((1-2*flip)*rho+flip*m.I, (1-2*flip)*drho)))

for t in [.1, 1., 2., 4.]:
    b = m.bare_optimized("Z", .4, .2, t)
    analytic = (1-.04)**2*4*t*t*np.exp(-4*.2*t)
    err(f"bare_parallel_optimized_analytic_t{t}", abs(b["actual_local_readout_FI"]-analytic), 2e-9)
    # The noncommuting baseline must beat the finite initial-state grid it refines.
    b = m.bare_optimized("X", .4, .2, t)
    ck(f"bare_X_refines_grid_t{t}", b["actual_local_readout_FI"] >= b["grid_best"]-1e-10)

# Inaccessible environment never resets. Good fixed sector matches single-noise
# Markov code, whereas the other sector carries noise along the signal.
for n in [2, 4, 8]:
    good = m.corrected(n=n, delta=1/n, read_flip=0.)
    hidden = m.corrected(n=n, delta=1/n, nu=0., initial_environment=[1., 0.], aggregate=False, read_flip=0.)
    err(f"fixed_good_sector_matches_Markov_n{n}", abs(good["CQFI"]-hidden["CQFI"]))
    bad = m.corrected(n=n, delta=1/n, nu=0., initial_environment=[0., 1.], aggregate=False, read_flip=0.)
    err(f"fixed_parallel_sector_unprotected_QFI_n{n}", abs(bad["CQFI"]-4*np.exp(-.8)), 2e-9)
fast = m.corrected(n=8, delta=.25, latency=.02, nu=200., initial_environment=[.5, .5], aggregate=False)
slow = m.corrected(n=8, delta=.25, latency=.02, nu=0., initial_environment=[.5, .5], aggregate=False)
ck("same_marginal_hidden_memory_changes_result", abs(fast["CQFI"]-slow["CQFI"]) > .1)

# Pure photon loss Kraus solution, including coherences in the error subspace.
N = np.diag(np.arange(5)).astype(complex)
for t in [.1, 1., 2.]:
    eta = np.exp(-t)
    psi = np.zeros(5, complex); psi[[1, 4]] = 1/np.sqrt(2)
    r0 = np.outer(psi, psi.conj()); rho = np.zeros_like(r0)
    for ell in range(5):
        K = np.zeros((5, 5), complex)
        for n in range(ell, 5):
            K[n-ell, n] = np.sqrt(math.comb(n, ell)*(1-eta)**ell*eta**(n-ell))
        rho += K@r0@K.conj().T
    phase = np.diag(np.exp(-1j*.4*t*np.arange(5)))
    rho = phase@rho@phase.conj().T
    d = m.fock(t)
    err(f"Fock_pure_loss_Kraus_t{t}", np.max(abs(d["rho"]-rho)))
    err(f"Fock_full_QFI_pair_measurement_t{t}", abs(d["QFI"]-d["all_pairs_Ramsey_FI"]))

for repair in [1., 10., 100.]:
    t = 1.
    a = np.diag(np.sqrt(np.arange(1, 5)), 1).astype(complex)
    R = np.zeros((5, 5), complex); R[1, 0] = 1; R[4, 3] = 1
    def rhs(_, y):
        rho = m.mat(y[:25], 5)
        drho = np.zeros_like(rho)
        for J in [a, np.sqrt(repair)*R]:
            B = J.conj().T@J
            drho += J@rho@J.conj().T-(B@rho+rho@B)/2
        return np.r_[m.vec(drho), np.trace(N@rho), repair*np.trace(R.conj().T@R@rho)]
    psi = np.zeros(5, complex); psi[[1, 4]] = 1/np.sqrt(2)
    y0 = np.r_[m.vec(np.outer(psi, psi.conj())), 0., 0.]
    sol = solve_ivp(rhs, (0, t), y0, rtol=1e-10, atol=1e-12)
    ck(f"Fock_independent_ODE_success_{repair}", sol.success)
    phase = np.diag(np.exp(-1j*.4*t*np.arange(5)))
    rho = phase@m.mat(sol.y[:25, -1], 5)@phase.conj().T
    d = m.fock(t, repair)
    err(f"Fock_ODE_density_{repair}", np.max(abs(d["rho"]-rho)), 2e-8)
    err(f"Fock_ODE_lost_photons_{repair}", abs(d["lost_photons"]-sol.y[25, -1].real), 2e-8)
    err(f"Fock_ODE_pump_photons_{repair}", abs(d["pump_photons"]-sol.y[26, -1].real), 2e-8)
    err(f"Fock_repaired_full_QFI_pair_measurement_{repair}", abs(d["QFI"]-d["all_pairs_Ramsey_FI"]))
    ck(f"Fock_code_only_measurement_loses_information_{repair}", d["code_only_Ramsey_FI"] <= d["QFI"]+1e-10)
    eps = 1e-5
    rp, rn = m.fock(t, repair, omega=.4+eps)["rho"], m.fock(t, repair, omega=.4-eps)["rho"]
    err(f"Fock_parameter_derivative_{repair}", np.max(abs((rp-rn)/(2*eps)-d["derivative"])), 1e-8)
for t in [.5, 1., 2.]:
    target = 9*t*t*np.exp(-t)
    d1, d2 = m.fock(t, 1000.), m.fock(t, 10000.)
    e1, e2 = abs(d1["QFI"]-target), abs(d2["QFI"]-target)
    ck(f"fast_repair_limit_converges_t{t}", e2 < e1/5 and e2/target < .003,
       error1000=e1, error10000=e2, target=target)

source = m.experiment_recalculation()
ck("printed_source_same_pair_gain_3point3dB", abs(source["same_pair_gain_dB"]-3.3) < .08)
ck("printed_source_different_pair_gain_6point3dB", abs(source["different_pair_gain_dB"]-6.3) < .08)
err("source_wall_time_464point21us", abs(source["rows"][0]["wall_us"]-464.21))
for row in source["rows"]:
    ck(row["label"]+"_exact_phase_max_near_quadrature",
       0 <= row["max_vs_quadrature_relative"] < .005,
       relative_difference=row["max_vs_quadrature_relative"])

for name in ["markov-control-scan.csv", "bare-initial-time-scan.csv", "hidden-environment-scan.csv", "fock-repair-scan.csv"]:
    rows = list(csv.DictReader((m.OUT/name).open()))
    bad = []
    for i, r in enumerate(rows):
        if "CQFI" in r:
            a, b, c = [float(r[k]) for k in ["CQFI", "actual_local_readout_FI", "record_ignored_QFI"]]
            if not (a >= -1e-10 and b <= a+1e-8 and a >= c-1e-8 and abs(float(r["mass"])-1)<2e-9):
                bad.append(i)
        elif "QFI" in r:
            if abs(float(r["QFI"])-float(r["all_pairs_Ramsey_FI"])) > 2e-8:
                bad.append(i)
    ck(name+"_all_rows_physical_invariants", not bad, rows=len(rows), failing_rows=bad)

result = {"by": "Codex", "type": "author-self-review", "passed": all(c["passed"] for c in checks),
          "count": len(checks), "checks": checks,
          "limitations": ["not an independent review or a rerun of the authors' experimental fits/SDPs",
                          "finite local FI checks do not establish general asymptotic scaling",
                          "ideal ancilla/intermediate projective recovery; no hardware gate circuit",
                          "bare optimization is a declared pure-input, no-control family"]}
(m.OUT/"audit.json").write_text(json.dumps(result, ensure_ascii=False, indent=2)+"\n")
print(json.dumps({"passed": result["passed"], "count": len(checks)}))
