"""Extract and validate the as-implemented neutral-sector core.

The neutral-sector core is defined operationally from the archived
Experiment-4 payloads (the depth series ``E4f_neutral_*`` mapped to the
dynamics-operator table by ``../dataset/data/manifest/job_manifest.csv``):

    R_N = S^dag U_1 U_0^dag S                                 (Eq. for R_N)

where ``S`` is the BALANCE separator whose first output mode is the
normalized uniform vector and ``U_d`` is the compiled interferometer unitary
of the depth-``d`` job.  This is the operator that actually ran; it is
characterized directly from the hardware submissions rather than assumed from
the nominal gate parameters of Ref. [washburn2026], which define a different
member of the same 1-fixing invariance family.  All manuscript claims depend
only on the invariance property R_N 1 = 1, which the implemented operator
satisfies.

The script verifies and prints the manuscript's validation checks:
  * the core is real and orthogonal (unitary to roundoff);
  * it fixes the uniform vector with eigenvalue +1 (so it preserves N);
  * det R_N = -1, with the full eigenphase spectrum;
  * the archived depth-2 and depth-3 circuits equal S R_N^d S^dag U_0 to
    machine precision (a single fixed core applied d times).

Outputs (written to ``data/executed_core/``):
  * ``U_depth0.csv`` ... ``U_depth3.csv`` -- compiled 8x8 interferometers U_d
  * ``implemented_core_operator.csv``     -- the core R_N (real, imag columns)
  * ``implemented_core_eigenphases.csv``  -- eigenphases of R_N (deg)
  * ``summary.txt``                       -- the validation report

Dependency-free (decodes the Perceval payloads directly).  Run from
``analysis/``:  ``python extract_core.py``
"""

from __future__ import annotations

import base64
import csv
import json
import struct
import zlib
from pathlib import Path

import numpy as np

HERE = Path(__file__).resolve().parent
RAW = HERE.parent / "dataset" / "data" / "raw" / "quandela_exports"
OUT = HERE / "data" / "executed_core"
N = 8

# Experiment 4 depth series, per ../dataset/data/manifest/job_manifest.csv
DEPTH_PAYLOADS = {
    0: "E4f_neutral_no__db3477ea.payload.json",
    1: "E4f_neutral_1x__7969ff5f.payload.json",
    2: "E4f_neutral_2x__69c4a47c.payload.json",
    3: "E4f_neutral_3x__6c093146.payload.json",
}


# --------------------------------------------------------------------------
# Minimal protobuf reader for the Perceval ":PCVL:zip:" experiment payload.
# --------------------------------------------------------------------------
def _read_varint(b: bytes, i: int) -> tuple[int, int]:
    shift = result = 0
    while True:
        byte = b[i]
        i += 1
        result |= (byte & 0x7F) << shift
        if not (byte & 0x80):
            break
        shift += 7
    return result, i


def _fields(b: bytes) -> list[tuple[int, int, object]]:
    i, n, out = 0, len(b), []
    while i < n:
        key, i = _read_varint(b, i)
        field, wire = key >> 3, key & 7
        if wire == 0:
            v, i = _read_varint(b, i)
        elif wire == 1:
            v = struct.unpack("<d", b[i:i + 8])[0]
            i += 8
        elif wire == 2:
            ln, i = _read_varint(b, i)
            v = b[i:i + ln]
            i += ln
        elif wire == 5:
            v = struct.unpack("<f", b[i:i + 4])[0]
            i += 4
        else:
            break
        out.append((field, wire, v))
    return out


def _try_matrix(blob: bytes) -> np.ndarray | None:
    try:
        fs = _fields(blob)
    except Exception:
        return None
    if not fs or any(w != 2 for _, w, _ in fs):
        return None
    entries = []
    for _, _, v in fs:
        re = im = 0.0
        ok = False
        for ff, ww, vv in _fields(v):
            if ww == 1 and ff == 1:
                re, ok = vv, True
            elif ww == 1 and ff == 2:
                im = vv
        if not ok:
            return None
        entries.append(complex(re, im))
    k = len(entries)
    r = int(round(k ** 0.5))
    if r * r != k or r < 2:
        return None
    return np.array(entries).reshape(r, r)


def _find_8x8(b: bytes, found: list | None = None) -> list[np.ndarray]:
    if found is None:
        found = []
    try:
        fs = _fields(b)
    except Exception:
        return found
    for _, w, v in fs:
        if w == 2 and isinstance(v, (bytes, bytearray)):
            m = _try_matrix(v)
            if m is not None and m.shape == (N, N):
                found.append(m)
            else:
                _find_8x8(v, found)
    return found


def compiled_unitary(path: Path) -> np.ndarray:
    d = json.loads(Path(path).read_text())
    exp = d["payload"]["experiment"]
    raw = zlib.decompress(base64.b64decode(exp[len(":PCVL:zip:"):]))
    inner = base64.b64decode(raw[len(":PCVL:Experiment:"):])
    mats = _find_8x8(inner)
    if not mats:
        raise RuntimeError(f"no 8x8 unitary found in {path.name}")
    return mats[0]


def balance_separator() -> np.ndarray:
    """BALANCE separator S: first output mode is the normalized uniform vector."""
    u = np.ones(N, dtype=complex) / np.sqrt(N)
    seed = np.column_stack([u, np.eye(N, dtype=complex)[:, 1:]])
    q, _ = np.linalg.qr(seed)
    if np.real(np.vdot(q[:, 0], u)) < 0:
        q[:, 0] *= -1
    return q.conj().T


# --------------------------------------------------------------------------
def main() -> None:
    OUT.mkdir(parents=True, exist_ok=True)
    lines: list[str] = []

    def log(msg: str = "") -> None:
        print(msg)
        lines.append(msg)

    log("Implemented neutral-sector core: extraction and validation")
    log("=" * 60)

    S = balance_separator()
    U = {}
    for d, name in DEPTH_PAYLOADS.items():
        U[d] = compiled_unitary(RAW / name)
        np.savetxt(OUT / f"U_depth{d}.csv", U[d].real, delimiter=",", fmt="%.15f")

    # Operational core R_N = S^dag U_1 U_0^dag S  (mode basis; matches Eq. for R_N).
    R = S.conj().T @ U[1] @ U[0].conj().T @ S

    imag = float(np.max(np.abs(R.imag)))
    unit = float(np.max(np.abs(R.conj().T @ R - np.eye(N))))
    one = np.ones(N) / np.sqrt(N)
    ev1 = complex(np.vdot(one, R @ one))
    det = float(np.linalg.det(R.real))
    phases = np.degrees(np.sort(np.angle(np.linalg.eigvals(R))))

    log("\nValidation checks (manuscript table):")
    log(f"  field / type               : real (max|Im| = {imag:.2e}), orthogonal")
    log(f"  unitarity residual         : {unit:.2e}")
    log(f"  uniform-vector eigenvalue  : {ev1.real:.6f}{ev1.imag:+.6f}i")
    log(f"  det R_N                    : {det:+.6f}")
    log(f"  eigenphases (deg)          : {np.round(phases, 4)}")
    rec = {}
    for d in (2, 3):
        rec[d] = float(np.max(np.abs(
            U[d] - S @ np.linalg.matrix_power(R, d) @ S.conj().T @ U[0])))
        log(f"  depth-{d} reconstruction res. : {rec[d]:.2e}")

    assert imag < 1e-12 and unit < 1e-12 and abs(ev1 - 1) < 1e-9
    assert rec[2] < 1e-9 and rec[3] < 1e-9

    with open(OUT / "implemented_core_operator.csv", "w", newline="") as f:
        w = csv.writer(f)
        w.writerow(["row", "col", "re", "im"])
        for i in range(N):
            for j in range(N):
                w.writerow([i, j, f"{R[i, j].real:.15f}", f"{R[i, j].imag:.15f}"])
    np.savetxt(OUT / "implemented_core_eigenphases.csv",
               phases, delimiter=",", header="eigenphase_deg", comments="")

    log("\nThe nominal gate parameters of Ref. [washburn2026] (theta=pi/4, "
        "psi=2pi/3)")
    log("define a complex operator that assigns the uniform vector eigenvalue "
        "e^{i pi/2}=i;")
    log("the implemented operator above is real, fixes the uniform vector with "
        "eigenvalue +1,")
    log("and is the archived representative used to exercise the neutral sector.")
    (OUT / "summary.txt").write_text("\n".join(lines) + "\n")
    print(f"\nArtifacts written to {OUT}")


if __name__ == "__main__":
    main()
