"""T139: finite quantum sensing instruments, inaccessible memory, approximate QEC.

Own finite models, NOT a rerun of an experimental data set or the PRX SDP code.
Units: hbar=1; omega is an angular-frequency parameter. QFI/FI refer to local
estimation at omega0. All syndrome outcomes are kept. Control latency switches
the signal off while the specified probe/environment noise continues; ideal
instantaneous syndrome/recovery and a noiseless auxiliary are still assumptions.
The final readout is an explicitly constructed local SLD projective measurement
with symmetric binary readout flips. Intermediate syndromes are ideal.
"""
from pathlib import Path
import csv
import json
import math

import numpy as np
from scipy.linalg import expm
from scipy.optimize import minimize, minimize_scalar

ROOT = Path(__file__).resolve().parents[2]
SLUG = "2026-10-13-error-corrected-sensing"
OUT = ROOT / "research/assets" / SLUG
I = np.eye(2, dtype=complex)
X = np.array([[0, 1], [1, 0]], complex)
Y = np.array([[0, -1j], [1j, 0]], complex)
Z = np.diag([1., -1.]).astype(complex)
PAULIS = [X, Y, Z]
PLUS = np.array([1, 1], complex) / np.sqrt(2)
MINUS = np.array([1, -1], complex) / np.sqrt(2)
V = np.column_stack([np.kron(PLUS, [1, 0]), np.kron(MINUS, [0, 1])])
P = V @ V.conj().T
XP, ZP = np.kron(X, I), np.kron(Z, I)
RECOVERY = [P, ZP @ (np.eye(4) - P)]


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


def mat(a, d):
    return a.reshape((d, d), order="F")


def commutator(G):
    d = G.shape[0]
    return -1j * (np.kron(np.eye(d), G) - np.kron(G.T, np.eye(d)))


def dissipator(A):
    d = A.shape[0]
    B = A.conj().T @ A
    return np.kron(A.conj(), A) - .5 * (
        np.kron(np.eye(d), B) + np.kron(B.T, np.eye(d)))


def channel_and_derivative(L, D, t):
    """Frechet derivative via an exact block exponential, no finite differences."""
    d = L.shape[0]
    A = np.block([[L, np.zeros_like(L)], [D, L]])
    E = expm(t * A)
    return E[:d, :d], E[d:, :d]


def qfi(rho, drho):
    """Normalized-state SLD QFI; relative support tolerance, no renormalization."""
    rho = (rho + rho.conj().T) / 2
    drho = (drho + drho.conj().T) / 2
    lam, U = np.linalg.eigh(rho)
    assert lam.min() > -2e-9, lam
    lam = np.maximum(lam, 0)
    D = U.conj().T @ drho @ U
    den = lam[:, None] + lam[None, :]
    use = den > 1e-12 * max(float(lam.max()), 1e-300)
    return float(np.sum(2 * np.abs(D[use]) ** 2 / den[use]))


def local_readout(rho, drho, flip):
    """Achievable local FI, with an explicit measurement for noisy binary readout.

    Flipping either result of a qubit projective measurement with probability e
    is equivalent, for these probabilities, to measuring rho_e=(1-2e)rho+e I.
    Its SLD eigenbasis is the optimal noisy binary projective measurement at the
    specified local parameter. The returned probabilities include the flip.
    """
    c = 1 - 2 * flip
    a = c * rho + flip * I
    da = c * drho
    lam, U = np.linalg.eigh((a + a.conj().T) / 2)
    D = U.conj().T @ da @ U
    den = lam[:, None] + lam[None, :]
    z = np.zeros_like(D)
    np.divide(2 * D, den, out=z, where=den > 1e-12)
    sld = U @ z @ U.conj().T
    _, basis = np.linalg.eigh((sld + sld.conj().T) / 2)
    probs = np.real(np.diag(basis.conj().T @ a @ basis))
    deriv = np.real(np.diag(basis.conj().T @ da @ basis))
    fi = float(sum(d*d/p for p, d in zip(probs, deriv) if p > 1e-14))
    return fi, probs, deriv, basis


def reduced_instruments(G, omega, gamma, delta, latency=0., nu=None):
    """Physical 4-d map, reduced only AFTER the complete syndrome/recovery.

    For a hidden environment, keep two unnormalized physical density blocks.
    Block 0 has gamma D[Z], block 1 gamma D[X], symmetric hopping rate nu.
    No environment measurement or reset is performed between correction rounds.
    """
    d = 16
    ne = 1 if nu is None else 2
    D0 = commutator(G)
    if ne == 1:
        Lnoise = gamma * dissipator(ZP)
        D = D0
    else:
        Lnoise = np.block([
            [gamma*dissipator(ZP)-nu*np.eye(d), nu*np.eye(d)],
            [nu*np.eye(d), gamma*dissipator(XP)-nu*np.eye(d)]])
        D = np.kron(np.eye(2), D0)
    E, DE = channel_and_derivative(Lnoise + omega*D, D, delta)
    C = expm(latency * Lnoise)
    E, DE = C @ E, C @ DE
    embed = np.kron(V.conj(), V)
    J, DJ = [], []
    for K in RECOVERY:
        B = V.conj().T @ K
        reduce = np.kron(B.conj(), B)
        left = np.kron(np.eye(ne), reduce)
        right = np.kron(np.eye(ne), embed)
        J.append(left @ E @ right)
        DJ.append(left @ DE @ right)
    if ne == 1:
        # Probe-only dynamics preserve the auxiliary label. Thus these maps are
        # EXACTLY diagonal on logical matrix units. Leaving floating-point
        # symmetry-forbidden entries (~1e-17) in repeated rare branches creates
        # artificial coherences larger than their tiny probabilities. Enforce
        # this analytic symmetry after bounding its residual, rather than
        # discarding physical outcomes or normalizing a nonpositive state.
        parallel = np.max(abs(G-ZP)) < 1e-14
        damp = np.exp(-2*gamma*(delta+latency))
        p1 = -.5*np.expm1(-2*gamma*(delta+latency))
        dp1 = 0.
        if parallel:
            p1 += damp*np.sin(omega*delta)**2
            dp1 = damp*delta*np.sin(2*omega*delta)
        for b, (prob, derivative) in enumerate([(1-p1, -dp1), (p1, dp1)]):
            coh = prob if parallel else (J[b][2, 2]+J[b][1, 1].conjugate())/2
            dcoh = derivative if parallel else (DJ[b][2, 2]+DJ[b][1, 1].conjugate())/2
            canonical = np.diag([prob, coh.conjugate(), coh, prob])
            dcanonical = np.diag([derivative, np.conj(dcoh), dcoh, derivative])
            assert np.max(abs(J[b]-canonical)) < 5e-13
            assert np.max(abs(DJ[b]-dcanonical)) < 5e-13
            J[b], DJ[b] = canonical, dcanonical
    elif np.max(abs(G-XP)) < 1e-14:
        # The same auxiliary-label symmetry holds for the classical hidden
        # environment. Obtain populations from its exact (environment, parity)
        # CTMC, avoiding spurious nonzero "impossible" records at nu=0.
        Q = np.zeros((4, 4))
        for e in [0, 1]:
            for parity in [0, 1]:
                i = 2*e+parity
                Q[2*(1-e)+parity, i] += nu
                Q[i, i] -= nu
                if e == 0:
                    Q[2*e+(1-parity), i] += gamma
                    Q[i, i] -= gamma
        population = expm((delta+latency)*Q)
        for b in [0, 1]:
            canonical, dcanonical = np.zeros_like(J[b]), np.zeros_like(DJ[b])
            for eo in [0, 1]:
                for ei in [0, 1]:
                    prob = population[2*eo+b, 2*ei]
                    block = J[b][4*eo:4*eo+4, 4*ei:4*ei+4]
                    dblock = DJ[b][4*eo:4*eo+4, 4*ei:4*ei+4]
                    coh = (block[2, 2]+block[1, 1].conjugate())/2 if prob > 0 else 0.
                    dcoh = (dblock[2, 2]+dblock[1, 1].conjugate())/2 if prob > 0 else 0.
                    canonical[4*eo:4*eo+4, 4*ei:4*ei+4] = np.diag([prob, coh.conjugate(), coh, prob])
                    dcanonical[4*eo:4*eo+4, 4*ei:4*ei+4] = np.diag([0., np.conj(dcoh), dcoh, 0.])
            assert np.max(abs(J[b]-canonical)) < 5e-12
            assert np.max(abs(DJ[b]-dcanonical)) < 5e-12
            J[b], DJ[b] = canonical, dcanonical
    return J, DJ


def propagate(J, DJ, n, initial_environment=None, aggregate_counts=False):
    ne = J[0].shape[0] // 4
    weights = [1.] if initial_environment is None else initial_environment
    assert len(weights) == ne and abs(sum(weights)-1) < 1e-12
    a = np.concatenate([w*vec(np.outer(PLUS, PLUS.conj())) for w in weights])[None, :]
    da = np.zeros_like(a)
    if aggregate_counts:
        # Counts suffice only after explicitly verifying commuting instruments
        # AND their differentiated products. Never use this for finite-nu HMMs.
        assert ne == 1
        assert np.max(abs(J[0]@J[1]-J[1]@J[0])) < 1e-11
        assert np.max(abs(DJ[0]@J[1]+J[0]@DJ[1]-DJ[1]@J[0]-J[1]@DJ[0])) < 1e-11
    for _ in range(n):
        b = [a @ j.T for j in J]
        db = [da @ j.T + a @ dj.T for j, dj in zip(J, DJ)]
        if aggregate_counts:
            aa = np.zeros((len(a)+1, a.shape[1]), complex)
            dd = np.zeros_like(aa)
            aa[:-1] += b[0]; aa[1:] += b[1]
            dd[:-1] += db[0]; dd[1:] += db[1]
            a, da = aa, dd
        else:
            a, da = np.concatenate(b), np.concatenate(db)
    return a.reshape(-1, ne, 4), da.reshape(-1, ne, 4)


def branch_metrics(states, derivatives, read_flip=.02):
    # The inaccessible environment is traced out here, not after every interval.
    a = states.sum(axis=1)
    da = derivatives.sum(axis=1)
    cq, rec, local, expose_env = 0., 0., 0., 0.
    rows = []
    for k, (v, dv) in enumerate(zip(a, da)):
        r, dr = mat(v, 2), mat(dv, 2)
        p, dp = float(np.trace(r).real), float(np.trace(dr).real)
        if p <= 1e-280:
            continue
        rr, drr = r/p, dr/p-r*(dp/p)/p
        f = dp*dp/p
        fq = qfi(rr, drr)
        achievable, probs, deriv, _ = local_readout(rr, drr, read_flip)
        cq += f + p*fq
        rec += f
        local += f + p*achievable
        rows.append({"branch": k, "p": p, "dp": dp, "conditional_QFI": fq,
                     "conditional_noisy_readout_FI": achievable,
                     "final_p0": float(probs[0]), "final_dp0": float(deriv[0])})
    for v, dv in zip(states.reshape(-1, 4), derivatives.reshape(-1, 4)):
        r, dr = mat(v, 2), mat(dv, 2)
        p, dp = float(np.trace(r).real), float(np.trace(dr).real)
        if p > 1e-280:
            expose_env += dp*dp/p + p*qfi(r/p, dr/p-r*(dp/p)/p)
    total, dtotal = mat(a.sum(axis=0), 2), mat(da.sum(axis=0), 2)
    mass = float(np.trace(total).real)
    assert abs(mass-1) < 2e-10 and abs(np.trace(dtotal)) < 2e-10
    ignored = qfi(total, dtotal)
    assert cq >= ignored-2e-9 and expose_env >= cq-2e-9 and local <= cq+2e-9
    return {"CQFI": cq, "record_FI": rec, "record_ignored_QFI": ignored,
            "actual_local_readout_FI": local, "environment_exposed_QFI_bound": expose_env,
            "mass": mass, "branches": len(a), "read_flip": read_flip, "branch_rows": rows}


def corrected(signal="X", omega=.4, gamma=.2, delta=.2, n=5, latency=0.,
              nu=None, initial_environment=None, read_flip=.02, aggregate=True):
    J, DJ = reduced_instruments(XP if signal == "X" else ZP, omega, gamma, delta, latency, nu)
    a, da = propagate(J, DJ, n, initial_environment, aggregate and nu is None)
    ans = branch_metrics(a, da, read_flip)
    ans.update(signal=signal, omega=omega, gamma=gamma, delta=delta, cycles=n,
               sensing_time=n*delta, latency_per_cycle=latency,
               control_time=n*latency, preparation_readout_time=.1,
               wall_time=n*(delta+latency)+.1, nu=nu,
               initial_environment=initial_environment, count_aggregation=aggregate and nu is None,
               noisy_probes=1, noiseless_auxiliaries=1, ideal_syndrome_measurements=n)
    ans["actual_FI_rate"] = ans["actual_local_readout_FI"] / ans["wall_time"]
    return ans


def bare_channel(signal, omega, gamma, t):
    G = X if signal == "X" else Z
    D = commutator(G)
    E, DE = channel_and_derivative(omega*D+gamma*dissipator(Z), D, t)
    # r' = A r for this unital channel; derivative excludes an initial-state derivative.
    A = np.array([[np.trace(B@mat(E@vec(C/2), 2)).real for C in PAULIS] for B in PAULIS])
    DA = np.array([[np.trace(B@mat(DE@vec(C/2), 2)).real for C in PAULIS] for B in PAULIS])
    return A, DA


def bare_optimized(signal, omega, gamma, t, read_flip=.02):
    """Optimize initial PURE one-probe states, no ancillary or adaptive controls.

    Sphere grid seeds local optimization. Final binary readout is optimal locally
    within noisy projective measurements. This is a declared finite baseline
    family, not an upper bound on every possible unencoded adaptive strategy.
    """
    A, DA = bare_channel(signal, omega, gamma, t)
    c = 1-2*read_flip
    def value(par):
        th, ph = par
        r0 = np.array([np.sin(th)*np.cos(ph), np.sin(th)*np.sin(ph), np.cos(th)])
        r, dr = c*(A@r0), c*(DA@r0)
        return float(dr@dr+(r@dr)**2/max(1-r@r, 1e-14))
    grid = [(value((th, ph)), th, ph) for th in np.linspace(0, np.pi, 9)
            for ph in np.linspace(0, 2*np.pi, 17)[:-1]]
    grid.sort(reverse=True)
    candidates = [(v, np.array([th, ph])) for v, th, ph in grid[:8]]
    for _, th, ph in grid[:8]:
        opt = minimize(lambda par: -value(par), [th, ph], method="L-BFGS-B",
                       bounds=[(0, np.pi), (0, 2*np.pi)], options={"ftol": 1e-14, "gtol": 1e-9})
        candidates.append((-float(opt.fun), opt.x))
    best, par = max(candidates, key=lambda z: z[0])
    return {"signal": signal, "sensing_time": t, "wall_time": t+.1,
            "read_flip": read_flip, "actual_local_readout_FI": best,
            "actual_FI_rate": best/(t+.1), "theta": float(par[0]), "phi": float(par[1]),
            "grid_best": grid[0][0], "noisy_probes": 1, "noiseless_auxiliaries": 0}


def fock(t, repair=0., kappa=1., omega=.4, pair=(1, 4)):
    """Exact closed Fock0..4 model, shared jump R=|1><0|+|4><3|.

    All loss sectors remain. No truncation at a populated upper boundary: both
    jumps and the diagonal Hamiltonian leave this five-state space invariant.
    The repair pump adds ONE photon per jump; neither its energy nor its clock
    is matched to an actual pulsed cavity experiment. Integrated flux is saved.
    """
    d = 5
    a = np.diag(np.sqrt(np.arange(1, d)), 1).astype(complex)
    N = np.diag(np.arange(d)).astype(complex)
    R = np.zeros((d, d), complex); R[1, 0] = 1; R[4, 3] = 1
    L = kappa*dissipator(a)+repair*dissipator(R)
    psi = np.zeros(d, complex); psi[list(pair)] = 1/np.sqrt(2)
    r0 = vec(np.outer(psi, psi.conj()))
    # Augment by cumulative lost and supplied photons, with zero feedback.
    A = np.zeros((27, 27), complex); A[:25, :25] = L
    A[25, :25] = kappa*vec(N.T)
    A[26, :25] = repair*vec((R.conj().T@R).T)
    y = expm(t*A) @ np.r_[r0, 0., 0.]
    phase = np.diag(np.exp(-1j*omega*t*np.arange(d)))
    rho = phase @ mat(y[:25], d) @ phase.conj().T
    drho = -1j*t*(N@rho-rho@N)
    full = qfi(rho, drho)
    m, n = pair
    mass = float((rho[m, m]+rho[n, n]).real)
    code_fi = float(4*(n-m)**2*t*t*abs(rho[m, n])**2/mass) if mass > 1e-14 else 0.
    all_pair_fi = 0.
    for left, right in [(0, 3), (1, 4)]:
        p = float((rho[left, left]+rho[right, right]).real)
        if p > 1e-14:
            all_pair_fi += 4*9*t*t*abs(rho[left, right])**2/p
    mean = float(np.trace(N@rho).real)
    loss, pump = float(y[25].real), float(y[26].real)
    assert abs(mean-sum(pair)/2-(pump-loss)) < 5e-9
    assert abs(np.trace(rho)-1) < 5e-10
    return {"t": t, "repair_over_kappa": repair/kappa, "kappa": kappa, "omega": omega,
            "pair": list(pair), "QFI": full, "code_only_Ramsey_FI": code_fi,
            "all_pairs_Ramsey_FI": all_pair_fi if pair == (1, 4) else None,
            "initial_mean_photons": sum(pair)/2, "final_mean_photons": mean,
            "code_probability": mass, "lost_photons": loss, "pump_photons": pump,
            "wall_time": t+.1, "FI_rate": full/(t+.1),
            "rho": rho, "derivative": drho}


def experiment_recalculation():
    """Recalculate only printed rounded fit parameters and times in S30-S32.

    No raw fitting, uncertainty re-estimation or experiment replication.
    Authors label the quadrature formula max_phi. With B != .5 its exact
    numerical maximum is slightly displaced; both values are retained.
    """
    rows = []
    for label, A, B, gap, cycles, fixed_us, control_us in [
        ("corrected_Fock14", .296, .516, 3, 3, 6.08, 2.71),
        ("uncorrected_Fock14", .246, .482, 3, 2, 6.08, 0.),
        ("uncorrected_Fock01", .261, .525, 1, 8, 5.08, 0.)]:
        sense = 150e-6*cycles
        wall = (150e-6+control_us*1e-6)*cycles+fixed_us*1e-6
        quad = A*A*gap*gap*sense*sense/(B*(1-B))/wall
        def fi(cos):
            p = B+A*cos
            return A*A*gap*gap*sense*sense*(1-cos*cos)/(p*(1-p))/wall
        opt = minimize_scalar(lambda x: -fi(x), bounds=(-1, 1), method="bounded",
                              options={"xatol": 1e-14})
        rows.append({"label": label, "A_rounded": A, "B_rounded": B, "gap": gap,
                     "cycles": cycles, "sensing_us": sense*1e6, "wall_us": wall*1e6,
                     "quadrature_FI_rate_kHz_inverse": quad*1000,
                     "exact_phase_max_FI_rate_kHz_inverse": -float(opt.fun)*1000,
                     "optimal_cos_phase": float(opt.x),
                     "max_vs_quadrature_relative": -float(opt.fun)/quad-1})
    a, b, c = [r["quadrature_FI_rate_kHz_inverse"] for r in rows]
    return {"source": "2509.26042v1 SM IV B S30-S32, printed rounded A/B and times only",
            "rows": rows, "same_pair_gain_dB": 10*math.log10(a/b),
            "different_pair_gain_dB": 10*math.log10(a/c),
            "raw_experimental_data_refitted": False}


def concise(d):
    return {k: v for k, v in d.items() if k not in ["branch_rows", "rho", "derivative"]}


def write_csv(name, rows):
    keys = list(rows[0])
    with (OUT/name).open("w", newline="") as f:
        w = csv.DictWriter(f, fieldnames=keys, lineterminator="\n")
        w.writeheader(); w.writerows(rows)


def main():
    OUT.mkdir(parents=True, exist_ok=True)
    opening = [corrected(signal=s, read_flip=0., aggregate=False) for s in ["X", "Z"]]
    markov = []
    for wall in [1.1, 2.1, 4.1]:
        for latency in [0., .005, .02, .05]:
            for n in [1, 2, 4, 8, 16, 32, 64, 128]:
                delta = (wall-.1-n*latency)/n
                if delta <= 0:
                    continue
                for s in ["X", "Z"]:
                    markov.append(concise(corrected(s, delta=delta, n=n, latency=latency)))
    write_csv("markov-control-scan.csv", markov)
    bare = []
    for s in ["X", "Z"]:
        for t in sorted(set([1., 2., 4.]+list(np.linspace(.05, 8, 81)))):
            bare.append(bare_optimized(s, .4, .2, t))
    write_csv("bare-initial-time-scan.csv", bare)
    hmm = []
    for p0 in [0., .5, 1.]:
        for nu in [0., .02, .2, 2., 20., 200.]:
            for t in [1., 2., 4.]:
                hmm.append(concise(corrected(delta=t/8, n=8, latency=.02, nu=nu,
                                             initial_environment=[p0, 1-p0], aggregate=False)))
    write_csv("hidden-environment-scan.csv", hmm)
    fock_rows = []
    for repair in [0., 1., 3., 10., 30., 100., 1000.]:
        for t in [.05, .1, .2, .5, 1., 2., 4., 6., 8., 10.]:
            fock_rows.append(concise(fock(t, repair)))
    write_csv("fock-repair-scan.csv", fock_rows)
    optima = []
    for repair in [0., 1., 3., 10., 30., 100., 1000.]:
        grid = [(fock(t, repair)["FI_rate"], t) for t in np.linspace(.005, 15., 201)]
        _, best = max(grid)
        result = minimize_scalar(lambda t: -fock(t, repair)["FI_rate"],
                                 bounds=(max(.005, best-.08), min(15., best+.08)), method="bounded",
                                 options={"xatol": 1e-10})
        optima.append(concise(fock(float(result.x), repair)))
    result = minimize_scalar(lambda t: -fock(t, 0., pair=(0, 1))["FI_rate"],
                             bounds=(.005, 15.), method="bounded")
    optima.append(concise(fock(float(result.x), 0., pair=(0, 1))))
    write_csv("fock-time-optima.csv", optima)
    detailed = corrected(delta=.23, n=4, latency=.02, nu=.2,
                         initial_environment=[.5, .5], aggregate=False)
    write_csv("hidden-example-branches.csv", detailed["branch_rows"])
    for label, repair in [("unprotected", 0.), ("repair100", 100.)]:
        d = fock(1., repair)
        np.savez_compressed(OUT/f"fock-{label}.npz", rho=d["rho"], derivative=d["derivative"])
    summary = {"scope": __doc__, "opening": [concise(x) for x in opening],
               "experiment_printed_parameter_recalculation": experiment_recalculation(),
               "hidden_example": concise(detailed), "markov_rows": len(markov),
               "bare_rows": len(bare), "hidden_rows": len(hmm), "fock_rows": len(fock_rows),
               "fock_time_optima": optima,
               "assumptions": ["local omega0=.4, angular-frequency convention",
                               "ideal noiseless ancilla and intermediate syndrome/recovery",
                               "signal off during correction latency; noise and hidden hopping continue",
                               "all records retained; no postselection or invisible-environment reset",
                               "final projective SLD readout depends on a local prior, binary flip .02",
                               "bare family optimizes pure input, final projective readout, scanned times",
                               "Fock continuous shared-pump model is not the pulsed cavity experiment",
                               "resource ledgers are not a whole-device energy budget"]}
    (OUT/"results.json").write_text(json.dumps(summary, ensure_ascii=False, indent=2)+"\n")
    print(json.dumps({"markov_rows": len(markov), "bare_rows": len(bare), "hidden_rows": len(hmm),
                      "fock_rows": len(fock_rows), "experiment": summary["experiment_printed_parameter_recalculation"]},
                     ensure_ascii=False))


if __name__ == "__main__":
    main()