"""
diagnostic_propulsion_reaudit.py
==================================

ELSD/v3 propulsion characterization re-audit.

Re-runs the five historical propulsion-page systems (AlH3_minimal,
AlH3_desorption, CH4_minimal, Ni_CH4_minimal, Ni_CH4_activated) under
the patched v1+v2+v3+v3-followup engine, captures every audit-relevant
quantity into a JSON manifest, and saves per-seed statevectors and
history.csv files for independent verification.

DISCIPLINE
----------
This is CHARACTERIZATION, not a rescue of the prior propulsion page's
framing. The script reports exactly what happens — sector status, ansatz
warnings, σ recomputation, correlation captured, dominant bitstrings —
with no retries, no hyperparameter tuning, no rescue logic. Whatever
the engine produces is what gets recorded.

EXACT REPRODUCTION
------------------
All five systems use the geometries, basis (LANL2DZ), charge, multiplicity,
active space (10e, 10o), and original ansatz_type ("uccsd") baked into the
engine registry at lines 5921–5968 and 559–645. This is exact reproduction
of the historical operating point, not best-effort reconstruction.

ANSATZ COMPARISON
-----------------
Each system can be run under either or both of:

  legacy_uccsd : The original page's stated operating point.
                 The engine audit established that the legacy 'uccsd' flag
                 routes to HEAWithGivensAnsatz, whose Givens block is
                 dormant on half-filled active spaces. All five propulsion
                 systems are half-filled at 20 qubits, so the legacy path
                 should fire [ANSATZ-WARN] for Givens dormancy on every
                 run. This run reproduces what the original page numbers
                 actually came from.

  plain_hea    : The engine's currently-demonstrated correlation-recovery
                 path. The molecule ladder showed plain HEA reaches
                 chemical accuracy on H2 and partial on LiH, then fails
                 on BeH2 and H2O at 14 qubits. Whether it recovers
                 correlation on 20-qubit propulsion systems is the
                 empirical question this run answers.

The script saves both runs to the manifest; the deposit author chooses
how to present them.

USAGE
-----
    # Smoke test: seed 0 only, both ansätze (10 runs total)
    python diagnostic_propulsion_reaudit.py \\
        --mode smoke --ansatz both --seeds 0 --save-statevectors

    # Legacy-only full reproduction (25 runs)
    python diagnostic_propulsion_reaudit.py \\
        --mode full --ansatz legacy_uccsd --seeds 0 1 2 3 4 --save-statevectors

    # Full characterization, both ansätze (50 runs)
    python diagnostic_propulsion_reaudit.py \\
        --mode full --ansatz both --seeds 0 1 2 3 4 --save-statevectors

OUTPUTS
-------
  - propulsion_reaudit_results.json    : Full manifest
  - propulsion_reaudit_<sys>_<ansatz>_seed<N>_statevector.npz : Per-seed
  - propulsion_reaudit_<sys>_<ansatz>_seed<N>_history.csv     : Per-seed
  - Full stdout log (recommend tee to file)

The .npz schema matches the metalloenzyme deposit's so the existing
`elsd_metalloenzyme_audit.py` verifier can audit these files unchanged.
"""

from __future__ import annotations

# [PROPULSION-REAUDIT-PENALTY-FIX] APPLIED

import argparse
import json
import os
import subprocess
import sys
import time
import traceback
from pathlib import Path
from typing import Optional

import numpy as np

sys.path.insert(0, str(Path(__file__).resolve().parent))


# ---------------------------------------------------------------------------
# System registry (recovered from engine, exact reproduction)
# ---------------------------------------------------------------------------

SYSTEMS = [
    {
        "label":              "AlH3_minimal",
        "molecule_name":      "AlH3_minimal",
        "basis":              "lanl2dz",
        "charge":             0,
        "multiplicity":       1,
        "n_active_electrons": 10,
        "n_active_orbitals":  10,
        "original_sigma":     0.0527,
        "original_classification": "Rigid Stability",
        "description":        "AlH3 resting state, alpha-AlH3 D3h planar (Al-H=1.575 A)",
        "expected_half_filled": True,
    },
    {
        "label":              "AlH3_desorption",
        "molecule_name":      "AlH3_desorption",
        "basis":              "lanl2dz",
        "charge":             0,
        "multiplicity":       1,
        "n_active_electrons": 10,
        "n_active_orbitals":  10,
        "original_sigma":     0.0872,
        "original_classification": "Rigid Stability",
        "description":        "AlH3 desorption geometry, H3 at Al-H=2.20 A (+z axis)",
        "expected_half_filled": True,
    },
    {
        "label":              "CH4_minimal",
        "molecule_name":      "CH4_minimal",
        "basis":              "lanl2dz",
        "charge":             0,
        "multiplicity":       1,
        "n_active_electrons": 10,
        "n_active_orbitals":  10,
        "original_sigma":     0.0262,
        "original_classification": "Rigid Stability (dom_p=1.0000 every seed)",
        "description":        "Pure methane Td baseline, C-H=1.087 A, nothing frozen",
        "expected_half_filled": True,
    },
    {
        "label":              "Ni_CH4_minimal",
        "molecule_name":      "Ni_CH4_minimal",
        "basis":              "lanl2dz",
        "charge":             0,
        "multiplicity":       1,
        "n_active_electrons": 10,
        "n_active_orbitals":  10,
        "original_sigma":     0.4147,
        "original_classification": "Within-Family Ruggedness",
        "description":        "Ni-CH4 physisorption (eta1), Ni-H=1.90 A",
        "expected_half_filled": True,
    },
    {
        "label":              "Ni_CH4_activated",
        "molecule_name":      "Ni_CH4_activated",
        "basis":              "lanl2dz",
        "charge":             0,
        "multiplicity":       1,
        "n_active_electrons": 10,
        "n_active_orbitals":  10,
        "original_sigma":     0.1648,
        "original_classification": "Rigid Stability",
        "description":        "Ni-CH4 activation, C-H1 stretched to 1.55 A",
        "expected_half_filled": True,
    },
]

# Fixed engine config (matches the original page's stated operating point)
SHARED_CONFIG = {
    "ansatz_depth":          6,
    "learning_rate":         0.05,
    "max_iterations":        1000,
    "patience":              200,
    "convergence_threshold": 1e-8,
    "use_amp":               False,
    # [PROPULSION-REAUDIT-PENALTY-FIX] These were already here but unused;
    # the VQEConfig call below now actually propagates them.
    "sector_penalty":        2.0,
    "lam_n":                 2.0,
    "lam_sz":                0.0,   # singlet systems; no Sz penalty needed
    "sz_target":             0.0,
}

# Banding (matches the molecule ladder script)
def band_classify(gap_fci_kcal: Optional[float]) -> str:
    if gap_fci_kcal is None:
        return "n/a (no exact reference)"
    g = abs(gap_fci_kcal)
    if g < 1.0:
        return "chemical_accuracy"
    if g < 5.0:
        return "useful"
    if g <= 20.0:
        return "partial"
    return "essentially_hf"


def _cuda_available() -> bool:
    try:
        import torch
        return torch.cuda.is_available()
    except Exception:
        return False


def _git_commit() -> str:
    try:
        r = subprocess.run(["git", "rev-parse", "HEAD"], capture_output=True, text=True, timeout=2)
        return r.stdout.strip() if r.returncode == 0 else "not_a_git_repo"
    except Exception:
        return "unknown"


def _engine_sentinels(engine_path: Path) -> dict:
    """Extract which audit-patch sentinels are present in the engine file."""
    if not engine_path.exists():
        return {"engine_file_found": False}
    text = engine_path.read_text()
    return {
        "engine_file_found": True,
        "v1_applied": "[ELSD-UNASSAILABILITY-PATCH-V1] APPLIED" in text,
        "v2_applied": "[ELSD-ANSATZ-HONESTY-PATCH-V2] APPLIED" in text,
        "v3_applied": "[ELSD-REPORTING-HARDENING-PATCH-V3] APPLIED" in text,
        "v3_followup_applied": "[ELSD-V3-FOLLOWUP-AUDIT-PROPAGATION] APPLIED" in text,
    }


def _save_statevector(solver, results: dict, out_prefix: Path) -> Optional[str]:
    """Save the final statevector to .npz using the engine's canonical schema."""
    try:
        import torch
        with torch.no_grad():
            if hasattr(solver.q_device, "get_states_1d"):
                sv = solver.q_device.get_states_1d().cpu().numpy()
            else:
                sv = solver.q_device.states.reshape(-1).cpu().numpy()

        sv_path = out_prefix.with_suffix(".statevector.npz")
        np.savez_compressed(
            str(sv_path),
            statevector=sv,
            system=solver.config.molecule_name,
            basis=solver.config.basis,
            vqe_energy=results.get("vqe_energy", 0.0),
            vqe_energy_physical=results.get("vqe_energy_physical", 0.0),
            hf_energy=solver.ham_info.get("hf_energy", solver.ham_info.get("hf_mf_energy", 0.0)),
            n_qubits=solver.ham_info["n_qubits"],
            n_active_electrons=solver.ham_info["n_active_electrons"],
            n_active_orbitals=solver.config.n_active_orbitals,
        )
        return str(sv_path)
    except Exception as e:
        print(f"  [STATEVECTOR-SAVE-FAILED] {e}")
        return None


def _save_history(results: dict, out_prefix: Path) -> Optional[str]:
    """Save the per-iteration energy trace to a CSV matching the deposit schema."""
    try:
        history = results.get("history") or results.get("energy_history")
        if history is None:
            return None
        path = out_prefix.with_suffix(".history.csv")
        with open(path, "w") as f:
            f.write("iter,energy_ha,vqe_minus_hf_kcal\n")
            hf = results.get("hf_energy", 0.0)
            for i, e in enumerate(history):
                kcal = (float(e) - float(hf)) * 627.509
                f.write(f"{i},{float(e):.10f},{kcal:+.6f}\n")
        return str(path)
    except Exception as e:
        print(f"  [HISTORY-SAVE-FAILED] {e}")
        return None


def _independent_sector_check(sv_path: Optional[str], n_qubits: int, n_target: int, sz_target: int) -> dict:
    """Reproduce the metalloenzyme-audit sector check on the saved statevector.

    Returns ⟨N⟩, ⟨2Sz⟩, P(N=N_t and 2Sz=Sz_t), dom_p, dominant bitstring.
    Spin convention: qubit 2i = β, qubit 2i+1 = α (matches metalloenzyme audit).
    """
    if sv_path is None:
        return {"sector_check_status": "no_statevector"}
    try:
        npz = np.load(sv_path)
        sv = npz["statevector"]
        probs = (np.abs(sv) ** 2).astype(np.float64)

        # Dominant bitstring
        i_max = int(np.argmax(probs))
        dom_p = float(probs[i_max])
        dom_bits = format(i_max, f"0{n_qubits}b")

        # ⟨N⟩, ⟨2Sz⟩
        N_expect = 0.0
        Sz2_expect = 0.0
        joint_target = 0.0
        for idx in range(len(probs)):
            p = float(probs[idx])
            if p < 1e-14:
                continue
            bits = format(idx, f"0{n_qubits}b")
            N = 0
            twoSz = 0
            for q in range(n_qubits):
                b = int(bits[n_qubits - 1 - q])
                if b == 1:
                    N += 1
                    # odd_alpha: qubit 2i+1 is alpha, qubit 2i is beta
                    if q % 2 == 1:
                        twoSz += 1   # alpha
                    else:
                        twoSz -= 1   # beta
            N_expect += p * N
            Sz2_expect += p * twoSz
            if N == n_target and twoSz == sz_target:
                joint_target += p

        return {
            "sector_check_status":   "ok",
            "N_expect":              N_expect,
            "twoSz_expect":          Sz2_expect,
            "joint_P_target_sector": joint_target,
            "dom_p":                 dom_p,
            "dominant_bitstring":    dom_bits,
            "spin_convention":       "odd_alpha (qubit 2i=β, 2i+1=α)",
        }
    except Exception as e:
        return {"sector_check_status": "error", "error": str(e)}


def run_one(system: dict, ansatz_type: str, seed: int, save_statevectors: bool,
            output_dir: Path) -> dict:
    """Run one (system, ansatz, seed) configuration. Never raises."""
    label = system["label"]
    run_label = f"{label}_{ansatz_type}_seed{seed}"
    out_prefix = output_dir / run_label

    out = {
        "run_label":          run_label,
        "system":             label,
        "ansatz_type":        ansatz_type,
        "seed":               seed,
        "molecule_name":      system["molecule_name"],
        "basis":              system["basis"],
        "charge":             system["charge"],
        "multiplicity":       system["multiplicity"],
        "n_active_electrons": system["n_active_electrons"],
        "n_active_orbitals":  system["n_active_orbitals"],
        "config":             dict(SHARED_CONFIG, ansatz_type=ansatz_type, seed=seed),
        "original_sigma":     system.get("original_sigma"),
        "original_classification": system.get("original_classification"),
        "status":             "pending",
    }

    print()
    print("#" * 78)
    print(f"# {run_label}")
    print(f"#   system: {system['molecule_name']} / {system['basis']}")
    print(f"#   active: ({system['n_active_electrons']}e, {system['n_active_orbitals']}o)")
    print(f"#   ansatz: {ansatz_type}  seed: {seed}")
    print(f"#   original σ on page: {system.get('original_sigma')}")
    print("#" * 78)

    t0 = time.time()
    try:
        import prometheus_vqe_engine_penalized_latest as engine

        cfg = engine.VQEConfig(
            molecule_name=system["molecule_name"],
            basis=system["basis"],
            charge=system["charge"],
            multiplicity=system["multiplicity"],
            n_active_electrons=system["n_active_electrons"],
            n_active_orbitals=system["n_active_orbitals"],
            ansatz_type=ansatz_type,
            ansatz_depth=SHARED_CONFIG["ansatz_depth"],
            learning_rate=SHARED_CONFIG["learning_rate"],
            max_iterations=SHARED_CONFIG["max_iterations"],
            patience=SHARED_CONFIG["patience"],
            convergence_threshold=SHARED_CONFIG["convergence_threshold"],
            seed=seed,
            output_prefix=str(out_prefix),
            device="cuda" if _cuda_available() else "cpu",
            use_amp=SHARED_CONFIG["use_amp"],
            # [PROPULSION-REAUDIT-PENALTY-FIX] Critical: actually pass these.
            # Without these, the engine runs UNCONSTRAINED — sector penalties
            # exist in SHARED_CONFIG only as documentation. This was the bug
            # that produced the 6/10 sector contaminations in the first
            # smoke test.
            sector_penalty=SHARED_CONFIG["sector_penalty"],
            lam_sz=SHARED_CONFIG.get("lam_sz", 0.0),
        )
        # Set the per-system N_target and Sz_target on the config.
        # (These are not always in VQEConfig's __init__ signature, so we
        # set them as attributes after construction. The engine reads them
        # via getattr() with defaults.)
        cfg.n_target = system["n_active_electrons"]
        cfg.sz_target = system.get("sz_target", SHARED_CONFIG.get("sz_target", 0.0))
        solver = engine.VQESolverV2(cfg)
        solver.setup()
        results = solver.optimize()

        # Save artifacts
        sv_path = None
        hist_path = _save_history(results, out_prefix)
        if save_statevectors:
            sv_path = _save_statevector(solver, results, out_prefix)

        # Energies
        hf = results.get("hf_energy")
        vqe_phys = results.get("vqe_energy_physical") or results.get("vqe_energy")
        casci = results.get("casci_energy")
        fci = results.get("fci_energy")
        ref = fci if fci is not None else casci

        gap_fci_kcal = None
        gap_hf_kcal = None
        total_corr_kcal = None
        captured_pct = None
        if ref is not None and vqe_phys is not None:
            gap_fci_kcal = (vqe_phys - ref) * 627.509
        if hf is not None and vqe_phys is not None:
            gap_hf_kcal = (vqe_phys - hf) * 627.509
        if hf is not None and ref is not None:
            total_corr_kcal = (hf - ref) * 627.509
            if vqe_phys is not None and total_corr_kcal != 0:
                captured_pct = 100.0 * (hf - vqe_phys) / (hf - ref)

        # Independent sector check
        n_qubits = results.get("n_qubits")
        sec_check = _independent_sector_check(
            sv_path, n_qubits,
            n_target=system["n_active_electrons"],
            sz_target=0  # singlets, target 2Sz=0
        )

        out.update({
            "status":                "ok",
            "elapsed_seconds":       round(time.time() - t0, 2),
            "n_qubits":              n_qubits,
            "iterations":            results.get("iterations"),
            # Engine v3 reporting fields
            "ansatz_legacy_flag":         results.get("ansatz_legacy_flag"),
            "ansatz_effective_type":      results.get("ansatz_effective_type"),
            "ansatz_givens_block_active": results.get("ansatz_givens_block_active"),
            "ansatz_givens_expected_calls": results.get("ansatz_givens_expected_calls"),
            "audit_warnings":             results.get("audit_warnings", []),
            "active_space_requested":     results.get("active_space_requested"),
            "active_space_effective":     results.get("active_space_effective"),
            "active_space_was_clamped":   results.get("active_space_was_clamped"),
            "eri_alpha":                  results.get("eri_alpha"),
            "reference_comparable":       results.get("reference_comparable"),
            "qubit_identity_imag_abs":    results.get("qubit_identity_imag_abs"),
            "has_exact_reference":        results.get("has_exact_reference"),
            "ref_energy_name":            results.get("ref_energy_name"),
            "h1_active_l1":               results.get("h1_active_l1"),
            "eri_active_l1":              results.get("eri_active_l1"),
            "truncation_n_dropped":       results.get("truncation_n_dropped"),
            "truncation_l1_dropped":      results.get("truncation_l1_dropped"),
            "sector_clean":               results.get("sector_clean"),
            "sector_purity_threshold":    results.get("sector_purity_threshold"),
            "sector_max_abs_penalty_ha":  results.get("sector_max_abs_penalty_ha"),
            "sector_sum_abs_penalty_ha":  results.get("sector_sum_abs_penalty_ha"),
            # Energies
            "hf_energy":                  hf,
            "vqe_energy_physical":        vqe_phys,
            "vqe_energy_minimized":       results.get("vqe_energy"),
            "casci_energy":               casci,
            "fci_energy":                 fci,
            "total_correlation_kcal":     total_corr_kcal,
            "vqe_minus_hf_kcal":          gap_hf_kcal,
            "vqe_minus_fci_kcal":         gap_fci_kcal,
            "captured_correlation_pct":   captured_pct,
            "band":                       band_classify(gap_fci_kcal),
            # Independent sector check
            "independent_sector_check":   sec_check,
            # Artifacts
            "statevector_path":           sv_path,
            "history_path":               hist_path,
        })

        # Compact summary
        print()
        print("-" * 72)
        print(f"SUMMARY  {run_label}")
        print(f"  n_qubits                 : {n_qubits}")
        print(f"  ansatz_effective_type    : {out['ansatz_effective_type']}")
        print(f"  ansatz_givens_block_active: {out['ansatz_givens_block_active']}")
        print(f"  has_exact_reference      : {out['has_exact_reference']}")
        print(f"  sector_clean             : {out['sector_clean']}")
        print(f"  reference_comparable     : {out['reference_comparable']}")
        print(f"  active_space_was_clamped : {out['active_space_was_clamped']}")
        print(f"  HF energy                : {hf}")
        print(f"  VQE energy (physical)    : {vqe_phys}")
        print(f"  FCI energy               : {ref}")
        if gap_fci_kcal is not None:
            print(f"  VQE - FCI                : {gap_fci_kcal:+.4f} kcal/mol")
        if captured_pct is not None:
            print(f"  captured                 : {captured_pct:+.1f}%")
        print(f"  band                     : {out['band']}")
        if out.get("audit_warnings"):
            print(f"  audit_warnings           : {len(out['audit_warnings'])}")
            for w in out['audit_warnings']:
                print(f"     - {w}")
        if sec_check.get("sector_check_status") == "ok":
            print(f"  indep. ⟨N⟩               : {sec_check['N_expect']:.4f}")
            print(f"  indep. ⟨2Sz⟩             : {sec_check['twoSz_expect']:+.4f}")
            print(f"  indep. dom_p             : {sec_check['dom_p']:.4f}")
            print(f"  indep. dominant bits     : {sec_check['dominant_bitstring']}")
        print(f"  elapsed                  : {out['elapsed_seconds']}s")
        print("-" * 72)

    except Exception as e:
        out.update({
            "status":          "error",
            "error_message":   str(e),
            "traceback":       traceback.format_exc(),
            "elapsed_seconds": round(time.time() - t0, 2),
        })
        print()
        print("!" * 72)
        print(f"FAILED: {run_label}")
        print(f"  {e}")
        print(f"  (continuing)")
        print("!" * 72)

    return out


def aggregate_per_system(runs: list, system_label: str, ansatz_type: str) -> dict:
    """Aggregate seed-ensemble results for one (system, ansatz) pair.

    Computes σ both ways (ddof=0 matches the page's stated procedure, ddof=1
    matches some VQE conventions). Reports dominant-bitstring agreement and
    sector cleanliness across the ensemble.
    """
    matched = [r for r in runs if r["system"] == system_label
               and r["ansatz_type"] == ansatz_type and r["status"] == "ok"]
    if not matched:
        return {"system": system_label, "ansatz_type": ansatz_type,
                "n_seeds_completed": 0, "status": "no_successful_seeds"}

    energies = [r["vqe_energy_physical"] for r in matched
                if r.get("vqe_energy_physical") is not None]
    if not energies:
        return {"system": system_label, "ansatz_type": ansatz_type,
                "n_seeds_completed": len(matched), "status": "no_physical_energies"}

    energies_kcal = [e * 627.509 for e in energies]
    sigma_ddof0_kcal = float(np.std(energies_kcal, ddof=0))
    sigma_ddof1_kcal = float(np.std(energies_kcal, ddof=1)) if len(energies_kcal) > 1 else 0.0
    mean_kcal = float(np.mean(energies_kcal))

    # Dominant bitstring agreement
    dom_bits = [r.get("independent_sector_check", {}).get("dominant_bitstring")
                for r in matched if r.get("independent_sector_check", {}).get("sector_check_status") == "ok"]
    bitstring_agree = (len(set(dom_bits)) == 1) if dom_bits else None

    # Sector cleanliness (engine's view + independent view)
    sector_clean_all = all(r.get("sector_clean") for r in matched)
    indep_sec_clean = None
    indep_secs = [r.get("independent_sector_check", {}) for r in matched]
    if all(s.get("sector_check_status") == "ok" for s in indep_secs):
        n_target = matched[0]["n_active_electrons"]
        indep_sec_clean = all(
            abs(s["N_expect"] - n_target) < 0.1 and abs(s["twoSz_expect"]) < 0.1
            for s in indep_secs
        )

    # Original σ comparison
    original_sigma = matched[0].get("original_sigma")
    sigma_match_exact = (
        abs(sigma_ddof0_kcal - original_sigma) < 0.0001
        if original_sigma is not None else None
    )

    out = {
        "system":                  system_label,
        "ansatz_type":             ansatz_type,
        "n_seeds_completed":       len(matched),
        "energies_ha":             energies,
        "mean_energy_kcal":        mean_kcal,
        "sigma_ddof0_kcal":        sigma_ddof0_kcal,
        "sigma_ddof1_kcal":        sigma_ddof1_kcal,
        "original_sigma_kcal":     original_sigma,
        "sigma_match_exact":       sigma_match_exact,
        "sigma_diff_kcal":         (sigma_ddof0_kcal - original_sigma) if original_sigma is not None else None,
        "dominant_bitstrings":     dom_bits,
        "bitstring_agreement":     bitstring_agree,
        "engine_sector_clean_all": sector_clean_all,
        "independent_sector_clean_all": indep_sec_clean,
        "audit_warnings_per_seed": [r.get("audit_warnings", []) for r in matched],
        "ansatz_givens_block_active_per_seed": [
            r.get("ansatz_givens_block_active") for r in matched
        ],
        "band_per_seed":           [r.get("band") for r in matched],
    }
    return out


def main():
    ap = argparse.ArgumentParser(
        description=__doc__,
        formatter_class=argparse.RawDescriptionHelpFormatter,
    )
    ap.add_argument("--mode", choices=["smoke", "full"], default="smoke",
                    help="smoke = seed 0 only (default); full = all listed seeds")
    ap.add_argument("--ansatz", choices=["legacy_uccsd", "plain_hea", "both"], default="both",
                    help="Which ansatz path(s) to run")
    ap.add_argument("--seeds", type=int, nargs="+", default=[0, 1, 2, 3, 4],
                    help="Seed list (used in full mode; smoke forces seed 0)")
    ap.add_argument("--save-statevectors", action="store_true",
                    help="Save final statevector .npz per run (recommended)")
    ap.add_argument("--output-dir", type=Path,
                    default=Path("propulsion_reaudit_out"),
                    help="Directory for statevectors, history, manifest")
    ap.add_argument("--systems", type=str, nargs="*", default=None,
                    help="Optional list of system labels to filter (default: all 5)")
    args = ap.parse_args()

    seeds = [0] if args.mode == "smoke" else args.seeds
    if args.ansatz == "both":
        ansatz_list = ["uccsd", "hea"]
    elif args.ansatz == "legacy_uccsd":
        ansatz_list = ["uccsd"]
    else:
        ansatz_list = ["hea"]

    systems = SYSTEMS
    if args.systems:
        systems = [s for s in SYSTEMS if s["label"] in args.systems]
        if not systems:
            print(f"[ERROR] No systems match {args.systems}")
            return 2

    args.output_dir.mkdir(parents=True, exist_ok=True)

    engine_path = Path(__file__).parent / "prometheus_vqe_engine_penalized_latest.py"

    print("=" * 78)
    print("PROPULSION RE-AUDIT  —  ELSD/v3 characterization")
    print("=" * 78)
    print(f"  mode:              {args.mode}")
    print(f"  ansätze:           {ansatz_list}")
    print(f"  seeds:             {seeds}")
    print(f"  systems:           {[s['label'] for s in systems]}")
    print(f"  save statevectors: {args.save_statevectors}")
    print(f"  output dir:        {args.output_dir}")
    print(f"  total runs:        {len(systems) * len(ansatz_list) * len(seeds)}")
    print(f"  engine sentinels:  {_engine_sentinels(engine_path)}")
    print()

    t0 = time.time()
    all_runs = []
    for sys_def in systems:
        for ansatz_type in ansatz_list:
            for seed in seeds:
                result = run_one(sys_def, ansatz_type, seed,
                                 args.save_statevectors, args.output_dir)
                all_runs.append(result)

    total_elapsed = round(time.time() - t0, 2)

    # Aggregate per (system, ansatz)
    aggregates = []
    for sys_def in systems:
        for ansatz_type in ansatz_list:
            aggregates.append(aggregate_per_system(all_runs, sys_def["label"], ansatz_type))

    # Final table
    print()
    print("=" * 78)
    print("FINAL TABLE")
    print("=" * 78)
    header = f"{'system':>22s} {'ansatz':>10s} {'n':>3s} {'σ(this)':>10s} {'σ(orig)':>10s} {'match':>6s} {'eng_clean':>9s} {'ind_clean':>9s} {'givens':>7s}"
    print(header)
    print("-" * len(header))
    for a in aggregates:
        sigma = a.get("sigma_ddof0_kcal")
        orig = a.get("original_sigma_kcal")
        match = a.get("sigma_match_exact")
        eng_cl = a.get("engine_sector_clean_all")
        ind_cl = a.get("independent_sector_clean_all")
        givens_list = a.get("ansatz_givens_block_active_per_seed", [])
        any_givens = any(v for v in givens_list) if givens_list else None
        print(f"{a['system']:>22s} {a['ansatz_type']:>10s} "
              f"{a['n_seeds_completed']:>3d} "
              f"{(f'{sigma:.4f}' if sigma is not None else 'n/a'):>10s} "
              f"{(f'{orig:.4f}' if orig is not None else 'n/a'):>10s} "
              f"{str(match):>6s} {str(eng_cl):>9s} {str(ind_cl):>9s} {str(any_givens):>7s}")
    print("-" * len(header))
    print(f"Total elapsed: {total_elapsed}s ({total_elapsed/60:.1f} min)")
    print()

    # Write manifest
    manifest = {
        "diagnostic": "propulsion_reaudit",
        "reproduction_mode": "exact",
        "geometry_source": "prometheus_vqe_engine_penalized_latest.py lines 559-645 (geometries) and 5921-5968 (config registry)",
        "original_sigma_source": "inline code comments and quantum-clarity.com/propulsion-reliability page (May 2026)",
        "claim_status": "historical propulsion page under re-audit; this deposit is ELSD/v3 characterization, not propulsion validation",
        "shared_config": SHARED_CONFIG,
        "engine_path": str(engine_path),
        "engine_sentinels": _engine_sentinels(engine_path),
        "git_commit": _git_commit(),
        "mode": args.mode,
        "ansatz_list": ansatz_list,
        "seeds": seeds,
        "systems_run": [s["label"] for s in systems],
        "total_runs": len(all_runs),
        "total_elapsed_seconds": total_elapsed,
        "runs": all_runs,
        "aggregates": aggregates,
        "interpretation_rules": [
            "If both ansätze converge to essentially_hf, the engine has not demonstrated correlation recovery on these 20-qubit systems under either path.",
            "If legacy_uccsd matches original σ to 4 decimals AND ansatz_givens_block_active is False, the original numbers are reproduced and the Givens block was dormant on those runs (same regime as Fe production).",
            "If independent_sector_clean disagrees with engine_sector_clean, the engine's audit and the independent statevector check disagree — investigate.",
            "Original σ values 0.0527 / 0.0872 / 0.0262 / 0.4147 / 0.1648 are seed-ensemble standard deviations characterized as such; their interpretation as physical-landscape ruggedness remains open pending in-sector exact reference solver.",
        ],
    }
    manifest_path = args.output_dir / "propulsion_reaudit_results.json"
    manifest_path.write_text(json.dumps(manifest, indent=2, default=str))
    print(f"Manifest written to: {manifest_path}")

    return 0


if __name__ == "__main__":
    sys.exit(main())
