"""T140 own finite, one-message quantum storage comparison.

Unknown PHYSICAL input to a PHYSICAL qubit output at deadline W, target identity.
No heralding/postselection/retry. Extra data and four-level controller START
maximally mixed; finite preparation and noisy encoder/decoder are included.
Two protocols share the same conditional coupling and finite controller reset:
continuous simultaneous coupling/reset, or switched coupling/reset/idle.
The two control integrals are SEPARATE capacity measures, not joules or a score.
Search envelopes are over this explicit finite catalogue, not general optima.
"""
from pathlib import Path
from functools import lru_cache
from dataclasses import dataclass, asdict
import csv
import itertools
import json
import math

import numpy as np
from scipy.linalg import expm
from scipy.sparse import csr_matrix, eye, kron

ROOT = Path(__file__).resolve().parents[2]
SLUG = "2026-10-13-matched-protection-budgets"
OUT = ROOT/"research/assets"/SLUG
I = np.eye(2, dtype=complex)
X = np.array([[0, 1], [1, 0]], complex)
Z = np.diag([1., -1.]).astype(complex)
IDX = np.arange(32)
ERROR = [0, 1, 2, 4]  # controller label -> data flip corrected by its coupling
PREP_RATE = 40.
PREP_TIMES = [.02, .05, .1, .2]
DEADLINES = [.5, 1., 2., 4.]
Z_RATES = [0., .02, .1, .2]
G_VALUES = [0., 2., 5., 10.]
GAMMA_VALUES = [0., 2., 5., 10., 20., 40.]
G_BUDGETS = [0., 5., 10., 20., 40., 80., 200.]
R_BUDGETS = [0., 2., 5., 10., 20., 40., 80.]


def dissipator(A):
    A = csr_matrix(A)
    B = A.conj().T@A
    ident = eye(A.shape[0], format="csr")
    return (kron(A.conj(), A)-.5*kron(ident, B)-.5*kron(B.T, ident)).tocsr()


def operators():
    xs, zs, rs = [], [], []
    for bit in [1, 2, 4]:
        a = np.zeros((8, 8))
        for e in range(8):
            a[e^bit, e] = 1
        r = np.zeros((8, 8)); r[0, bit] = r[7, 7^bit] = 1
        xs.append(a); zs.append(np.diag([1-2*bool(e&bit) for e in range(8)])); rs.append(r)
    H = np.zeros((32, 32), complex)
    LA = csr_matrix((1024, 1024), dtype=complex)
    for k, r in enumerate(rs, 1):
        a = np.zeros((4, 4)); a[k, 0] = 1
        H += np.kron(r, a)+np.kron(r.T, a.T)
        LA += dissipator(np.kron(np.eye(8), a.T))
    ident = eye(32, format="csr")
    LH = -1j*(kron(ident, csr_matrix(H))-kron(csr_matrix(H.T), ident))
    LX = sum(dissipator(np.kron(a, np.eye(4))) for a in xs)
    LZ = sum(dissipator(np.kron(a, np.eye(4))) for a in zs)
    return xs, zs, rs, H, LX, LZ, LH.tocsr(), LA


XS, ZS, RS, H, LX, LZ, LH, LA = operators()
EXCITED = np.diag([0., 1., 1., 1.])
PEXC = np.kron(np.eye(8), EXCITED)
FLUX = PEXC.T.reshape(-1, order="F")


def sector(difference):
    ids = []
    for data in range(8):
        for a in range(4):
            for b in range(4):
                other = data^ERROR[a]^ERROR[b]^difference
                row, col = 4*data+a, 4*other+b
                ids.append(row+32*col)
    return np.array(ids)


PIDS, CIDS = sector(0), sector(7)
assert len(set(PIDS)) == len(set(CIDS)) == 128 and not set(PIDS)&set(CIDS)
SECTORS = [PIDS, CIDS]
COMPONENTS = []
for ids in SECTORS:
    others = np.array(sorted(set(range(1024))-set(ids)))
    for L in [LX, LZ, LH, LA]:
        outside = L[others][:, ids]
        assert not outside.nnz or np.max(abs(outside.data)) < 1e-13
    COMPONENTS.append([L[ids][:, ids].toarray() for L in [LX, LZ, LH, LA]])


def full_generator(kx, kz, g=0., gamma=0.):
    return (kx*LX+kz*LZ+g*LH+gamma*LA).tocsr()


def expand(p, c):
    """Two diagonal input operators, one off-diagonal input operator."""
    B = np.zeros((1024, 3), complex)
    B[PIDS, :2] = p
    B[CIDS, 2] = c
    return B.reshape((32, 32, 3), order="F")


def reduce(B):
    v = B.reshape((1024, 3), order="F")
    residual = v.copy(); residual[PIDS, :2] = 0; residual[CIDS, 2] = 0
    assert np.max(abs(residual)) < 2e-12
    return v[PIDS, :2].copy(), v[CIDS, 2].copy()


def conjugate_pauli(B, q, z=False):
    bit = q+2  # data bits above the two auxiliary-index bits
    if z:
        s = 1-2*((IDX>>bit)&1)
        return s[:, None, None]*B*s[None, :, None]
    perm = IDX^(1<<bit)
    return B[perm][:, perm]


def data_noise(B, kx, kz, t):
    # Independent Pauli white-noise channels commute, exact finite-time map.
    px, pz = -.5*np.expm1(-2*kx*t), -.5*np.expm1(-2*kz*t)
    for q in range(3):
        if px:
            B = (1-px)*B+px*conjugate_pauli(B, q)
        if pz:
            B = (1-pz)*B+pz*conjugate_pauli(B, q, True)
    return B


def reset_aux(B, exposure):
    if exposure == 0:
        return B
    eta = np.exp(-exposure)
    a = np.tile([1., np.sqrt(eta), np.sqrt(eta), np.sqrt(eta)], 8)
    C = a[:, None, None]*B*a[None, :, None]
    ground = np.arange(0, 32, 4)
    for k in [1, 2, 3]:
        C[ground[:, None], ground[None, :]] += (1-eta)*B[(ground+k)[:, None], (ground+k)[None, :]]
    return C


def controller_excitation(B):
    return float(sum(np.trace(PEXC@B[:, :, k]).real for k in [0, 1])/2)


def permutation(kind):
    data = np.arange(8)
    if kind == "CX01":
        data = data^((data&1)<<1)
    elif kind == "CX02":
        data = data^((data&1)<<2)
    else:
        assert kind == "TOF"
        data = data^(((data>>1)&1)*((data>>2)&1))
    return np.array([4*data[e]+a for e in range(8) for a in range(4)])


PERMS = {k: permutation(k) for k in ["CX01", "CX02", "TOF"]}


def local_twirl(B, support):
    for q in support:
        xx = conjugate_pauli(B, q)
        zz = conjugate_pauli(B, q, True)
        yy = conjugate_pauli(xx, q, True)
        B = (B+xx+zz+yy)/4
    return B


def gate(B, kind, kx, kz, p2=.001, p3=.002):
    support = [0, 1] if kind == "CX01" else ([0, 2] if kind == "CX02" else [0, 1, 2])
    t, p = (.04, p3) if kind == "TOF" else (.02, p2)
    B = data_noise(B, kx, kz, t)
    perm = PERMS[kind]
    B = B[perm][:, perm]
    # Uniform nonidentity Pauli channel on the gate's support.
    factor = p*4**len(support)/(4**len(support)-1)
    return (1-factor)*B+factor*local_twirl(B, support)


@lru_cache(maxsize=256)
def prepared(kx, kz, tau, active):
    """Finite preparation, unknown data0 waits while other resources reset.

    Unused data and auxiliary start independently maximally mixed. Data1/2
    reset at PREP_RATE with the same X/Z noise; controller has no intrinsic
    noise in this model and its three excited levels reset at PREP_RATE.
    Encoded-no-repair needs no auxiliary preparation or auxiliary inventory.
    """
    rate = PREP_RATE+2*kx
    steady = kx/rate
    excited = steady+(.5-steady)*np.exp(-rate*tau)
    d = np.diag([1-excited, excited])
    if active:
        e = .25*np.exp(-PREP_RATE*tau)
        aux = np.diag([1-3*e, e, e, e])
    else:
        aux = np.eye(4)/4  # unused mathematical tensor factor, no controller cost
    B = []
    for a, b in [(0, 0), (1, 1), (0, 1)]:
        input = np.zeros((2, 2), complex); input[a, b] = 1
        px = -.5*np.expm1(-2*kx*tau)
        pz = -.5*np.expm1(-2*kz*tau)
        input = (1-px)*input+px*X@input@X
        input = (1-pz)*input+pz*Z@input@Z
        B.append(np.kron(d, np.kron(d, np.kron(input, aux))))
    B = np.stack(B, axis=2)
    B = gate(gate(B, "CX01", kx, kz), "CX02", kx, kz)
    p, c = reduce(B)
    prep_events_data = 2*PREP_RATE*(steady*tau+(.5-steady)*(-np.expm1(-rate*tau))/rate)
    prep_events_aux = .75*(-np.expm1(-PREP_RATE*tau)) if active else 0.
    return p, c, float(prep_events_data), float(prep_events_aux), float(excited)


@lru_cache(maxsize=192)
def propagation(kx, kz, g, gamma, t):
    mats = [kx*a+kz*b+g*c+gamma*d for a, b, c, d in COMPONENTS]
    # Integral of actual emitted reset events for each input population operator.
    A = np.zeros((129, 129), complex); A[:128, :128] = mats[0]
    A[128, :128] = gamma*FLUX[PIDS]
    return expm(t*A), expm(t*mats[1])


def evolve(p, c, kx, kz, g, gamma, t):
    if t <= 1e-14:
        return p, c, 0.
    A, C = propagation(kx, kz, g, gamma, float(t))
    y = A@np.vstack([p, np.zeros((1, 2))])
    return y[:128], C@c, float(y[128].real.mean())


def free_segment(p, c, kx, kz, gamma, t):
    if t <= 1e-14:
        return p, c, 0.
    B = expand(p, c)
    events = controller_excitation(B)*(-np.expm1(-gamma*t))
    B = reset_aux(data_noise(B, kx, kz, t), gamma*t)
    p, c = reduce(B)
    return p, c, events


def decoder(B, kx, kz, p2=.001, p3=.002):
    for gate_name in ["CX02", "CX01", "TOF"]:
        B = gate(B, gate_name, kx, kz, p2, p3)
    return B


def output_channel(B):
    out = []
    for k in range(3):
        data = np.einsum("iaja->ij", B[:, :, k].reshape(8, 4, 8, 4))
        out.append(np.einsum("aiaj->ij", data.reshape(4, 2, 4, 2)))
    A, B, C = out
    choi = np.block([[A, C], [C.conj().T, B]])/2
    bell = np.array([1, 0, 0, 1])/np.sqrt(2)
    fe = float(np.vdot(bell, choi@bell).real)
    trace = float(np.trace(choi).real)
    tp = np.trace(choi.reshape(2, 2, 2, 2), axis1=1, axis2=3)
    tp_error = float(np.max(abs(tp-I/2)))
    minimum = float(np.linalg.eigvalsh(choi).min())
    assert abs(trace-1) < 2e-9 and tp_error < 2e-9 and minimum > -2e-9
    return {"Fe": fe, "average_state_fidelity": (2*fe+1)/3,
            "classical_Z_success": float((A[0, 0].real+B[1, 1].real)/2),
            "channel_trace": trace, "TP_error": tp_error, "Choi_min_eigenvalue": minimum}, choi


def bare(W, kx=.2, kz=.02):
    # All fixed input-basis changes preserve this unknown-state Fe. Pauli echo
    # pulses only flip Lindblad operator signs and leave white noise unchanged.
    return float((1+np.exp(-2*kx*W))*(1+np.exp(-2*kz*W))/4)


@dataclass(frozen=True)
class Protocol:
    kind: str
    prep: float = .1
    g: float = 0.
    gamma: float = 0.
    angle: float = 0.
    reset_exposure: float = 0.
    cycle_wait: float = 0.
    offset: float = 0.


def catalogue():
    for tau in PREP_TIMES:
        yield Protocol("encoded_no_repair", tau)
        for g in G_VALUES:
            for gamma in GAMMA_VALUES:
                if g == 0 and gamma == 0:
                    continue
                yield Protocol("continuous", tau, g, gamma)
        for g, gamma, angle, reset, wait, offset in itertools.product(
                [2., 5., 10.], [5., 20., 40.],
                [np.pi/4, np.pi/2, 3*np.pi/4], [.5, 1.5, 3.], [0., .1], [0., .5]):
            yield Protocol("switched", tau, g, gamma, angle, reset, wait, offset)


def simulate(spec, W=2., kx=.2, kz=.02, save_channel=False):
    core = W-spec.prep-.12
    assert core >= -1e-12
    active = spec.kind != "encoded_no_repair"
    pp, cc, prep_d, prep_a, prep_error = prepared(kx, kz, spec.prep, active)
    p, c = pp.copy(), cc.copy()
    used_g, used_r, reset_events, segments = 0., 0., 0., []

    def segment(typ, t):
        nonlocal p, c, used_g, used_r, reset_events
        if t <= 1e-14:
            return
        gg = spec.g if typ in ["couple", "continuous"] else 0.
        rr = spec.gamma if typ in ["reset", "continuous"] else 0.
        if gg:
            p, c, events = evolve(p, c, kx, kz, gg, rr, t)
        else:
            p, c, events = free_segment(p, c, kx, kz, rr, t)
        used_g += gg*gg*t
        used_r += rr*t
        reset_events += events
        segments.append({"type": typ, "t": t, "g": gg, "gamma": rr})

    if spec.kind == "continuous":
        segment("continuous", core)
    elif spec.kind == "switched":
        pulse, reset = spec.angle/spec.g, spec.reset_exposure/spec.gamma
        cycle = pulse+reset+spec.cycle_wait
        remaining = core
        first = min(remaining, spec.offset*cycle)
        segment("idle", first); remaining -= first
        while remaining > 1e-12:
            for typ, duration in [("idle", spec.cycle_wait), ("couple", pulse), ("reset", reset)]:
                t = min(remaining, duration)
                segment(typ, t); remaining -= t
                if remaining <= 1e-12:
                    break
    else:
        assert spec.kind == "encoded_no_repair"
        segment("idle", core)
    B = expand(p, c)
    aux_end = controller_excitation(B) if active else 0.
    B = decoder(B, kx, kz)
    metrics, choi = output_channel(B)
    ans = {"deadline": W, "kx": kx, "kz": kz, **asdict(spec), **metrics,
           "bare_Fe": bare(W, kx, kz), "storage_core_time": core,
           "coupling_budget_used": used_g,
           "aux_reset_budget_used": used_r+(PREP_RATE*spec.prep if active else 0.),
           "core_reset_integral": used_r, "data_reset_integral": 2*PREP_RATE*spec.prep,
           "prep_data_reset_events": prep_d, "prep_aux_reset_events": prep_a,
           "core_aux_reset_events": reset_events, "final_aux_excitation": aux_end,
           "prep_data_excitation": prep_error, "physical_data": 3,
           "controller_dimension": 4 if active else 0,
           "CNOT_count": 4, "Toffoli_count": 1, "encoding_decoding_time": .12,
           "unconditional_delivery_probability": 1.}
    assert abs(sum(a["t"] for a in segments)-core) < 2e-10
    assert used_g <= spec.g**2*core+1e-9
    assert used_r <= spec.gamma*core+1e-9
    if save_channel:
        return ans, choi, B, segments
    return ans


def envelope(rows, W, kz, Bg, Br):
    allowed = [r for r in rows if r["deadline"] == W and r["kz"] == kz and
               r["coupling_budget_used"] <= Bg+1e-9 and r["aux_reset_budget_used"] <= Br+1e-9]
    result = {"deadline": W, "kz": kz, "coupling_cap": Bg, "aux_reset_cap": Br,
              "bare_Fe": bare(W, kz=kz)}
    for kind in ["encoded_no_repair", "continuous", "switched"]:
        choices = [r for r in allowed if r["kind"] == kind]
        best = max(choices, key=lambda r: r["Fe"]) if choices else None
        result[kind+"_Fe"] = best["Fe"] if best else None
        result[kind+"_id"] = best["id"] if best else None
    scored = [(result["bare_Fe"], "bare", None)]
    scored += [(r["Fe"], r["kind"], r["id"]) for r in allowed]
    value, kind, identity = max(scored)
    result.update(best_Fe=value, best_kind=kind, best_id=identity)
    return result


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


def main():
    OUT.mkdir(parents=True, exist_ok=True)
    specs = list(catalogue())
    rows = []
    for W in DEADLINES:
        for kz in Z_RATES:
            for spec in specs:
                if spec.prep+.12 > W:
                    continue
                result = simulate(spec, W=W, kz=kz)
                result["id"] = len(rows)
                rows.append(result)
            print("catalogue condition", W, kz, "cumulative", len(rows), flush=True)
    finalize(rows)


def finalize(rows, refinement=None):
    write_csv("protocol-catalogue.csv", rows)
    fronts = [envelope(rows, W, kz, Bg, Br) for W, kz, Bg, Br in itertools.product(
        DEADLINES, Z_RATES, G_BUDGETS, R_BUDGETS)]
    write_csv("matched-budget-envelopes.csv", fronts)
    targets, resource_frontiers = [], []
    for W, kz, target in itertools.product(DEADLINES, Z_RATES, [.8, .85, .9, .95]):
        for kind in ["bare", "encoded_no_repair", "continuous", "switched"]:
            if kind == "bare":
                good = bare(W, kz=kz) >= target
                best = None
                Bg = Br = 0.
            else:
                choices = [r for r in rows if r["deadline"] == W and r["kz"] == kz and
                           r["kind"] == kind and r["Fe"] >= target]
                good = bool(choices)
                lowest_reset = float("inf")
                for point in sorted(choices, key=lambda r: (
                        r["coupling_budget_used"], r["aux_reset_budget_used"], -r["Fe"])):
                    if point["aux_reset_budget_used"] < lowest_reset-1e-10:
                        lowest_reset = point["aux_reset_budget_used"]
                        resource_frontiers.append({"deadline": W, "kz": kz, "target_Fe": target,
                            "kind": kind, "id": point["id"], "Fe": point["Fe"],
                            "coupling_integral": point["coupling_budget_used"],
                            "aux_reset_integral": point["aux_reset_budget_used"],
                            "prep_time": point["prep"], "physical_data": point["physical_data"],
                            "controller_dimension": point["controller_dimension"]})
                best = min(choices, key=lambda r: (r["coupling_budget_used"], r["aux_reset_budget_used"])) if good else None
                Bg = best["coupling_budget_used"] if best else None
                Br = best["aux_reset_budget_used"] if best else None
            # Lexicographic choice is a labeled example, not a universal scalar
            # resource score. The complete two-axis frontier remains available.
            targets.append({"deadline": W, "kz": kz, "target_Fe": target, "kind": kind,
                            "catalogue_feasible": good, "example_id": best["id"] if best else None,
                            "example_coupling_integral": Bg, "example_reset_integral": Br,
                            "selection": "least coupling then reset among sampled feasible examples"})
    write_csv("quality-targets.csv", targets)
    write_csv("quality-resource-frontiers.csv", resource_frontiers)
    examples = {}
    for W, kz, Bg, Br in [(2., .02, 20., 20.), (2., .02, 80., 40.),
                          (.5, .02, 20., 20.), (2., .2, 80., 40.)]:
        front = next(r for r in fronts if all(r[k] == v for k, v in [
            ("deadline", W), ("kz", kz), ("coupling_cap", Bg), ("aux_reset_cap", Br)]))
        key = f"W{W:g}_Z{kz:g}_G{Bg:g}_R{Br:g}"
        examples[key] = {"envelope": front, "protocols": {}}
        for kind in ["encoded_no_repair", "continuous", "switched"]:
            identity = front[kind+"_id"]
            if identity is None:
                continue
            row = rows[identity]
            spec = Protocol(**{k: row[k] for k in Protocol.__dataclass_fields__})
            detail, choi, state, segments = simulate(spec, W, kz=kz, save_channel=True)
            examples[key]["protocols"][kind] = {"metrics": detail, "segments": segments}
            np.savez_compressed(OUT/f"{key}-{kind}.npz", choi=choi, output_operator_basis=state)
    summary = {"scope": __doc__, "rows": len(rows), "envelopes": len(fronts),
               "target_examples": len(targets), "resource_frontier_points": len(resource_frontiers),
               "examples": examples, "refinement": refinement,
               "grid": {"prep_rate": PREP_RATE, "prep_times": PREP_TIMES,
                        "deadlines": DEADLINES, "Z_rates": Z_RATES,
                        "continuous_g": G_VALUES, "continuous_gamma": GAMMA_VALUES,
                        "switched_g": [2., 5., 10.], "switched_gamma": [5., 20., 40.],
                        "switched_angles": [np.pi/4, np.pi/2, 3*np.pi/4],
                        "switched_reset_exposures": [.5, 1.5, 3.],
                        "switched_cycle_wait": [0., .1], "switched_offset": [0., .5],
                        "g_caps": G_BUDGETS, "reset_caps": R_BUDGETS},
               "resource_scope": ["one message, no repeated-service throughput claim",
                                  "separate integrals of g^2 and common auxiliary reset rate",
                                  "reset cap includes auxiliary PREPARATION",
                                  "data preparation integral <=16, independent channel rates <=40",
                                  "fixed 4 CNOT + 1 Toffoli, .12 time, uniform gate Pauli faults",
                                  "three data and at most one four-level controller",
                                  "conditional coupling is assumed available, no microscopic hardware-cost fit",
                                  "all branch failures remain in the CPTP output, no retry or postselection"]}
    (OUT/"results.json").write_text(json.dumps(summary, ensure_ascii=False, indent=2)+"\n")
    print(json.dumps({"rows": len(rows), "envelopes": len(fronts), "targets": len(targets)}))


if __name__ == "__main__":
    main()