"""
sector_check_postprocess.py
============================

Post-processes the statevectors saved by diagnostic_propulsion_reaudit.py
to perform the independent sector check that failed during the original
run due to a statevector-shape loading bug in the in-script verifier.

This script:
  - Reads each <run_label>.statevector.npz in the output directory.
  - Defensively handles 1-D, 2-D, or batched statevector storage formats.
  - Computes ⟨N⟩, ⟨2Sz⟩, P(N=N_t ∧ 2Sz=Sz_t), dom_p, dominant bitstring.
  - Reads the existing propulsion_reaudit_results.json manifest.
  - Updates each run's `independent_sector_check` entry in place.
  - Writes the corrected manifest back to disk (with .bak preserved).

It does NOT re-run any VQE. The .npz files on disk already contain
everything needed; this is pure post-processing.

Spin convention used: qubit 2i = β (beta), qubit 2i+1 = α (alpha).
Matches `elsd_metalloenzyme_audit.py` (odd_alpha).

Usage
-----
    python sector_check_postprocess.py \\
        --output-dir propulsion_reaudit_smoke

    # dry-run: show what would change without modifying the manifest
    python sector_check_postprocess.py \\
        --output-dir propulsion_reaudit_smoke --dry-run

    # write to a different manifest file (keep original untouched)
    python sector_check_postprocess.py \\
        --output-dir propulsion_reaudit_smoke \\
        --out propulsion_reaudit_smoke/results_with_sector_check.json
"""

from __future__ import annotations

import argparse
import json
import shutil
import sys
from pathlib import Path
from typing import Optional

import numpy as np


def _load_statevector(npz_path: Path) -> np.ndarray:
    """Load a statevector .npz defensively.

    Handles all observed shapes:
      - (2^n,)        : already flat 1-D (expected)
      - (1, 2^n)      : batched first dim of size 1 (torchquantum get_states_1d)
      - (2, 2, ..., 2): n-D tensor of 2s (raw q_device.states), n axes
      - (b, 2, 2, ..) : batched n-D, b=1

    Returns a flat 1-D complex numpy array of length 2^n.
    """
    npz = np.load(str(npz_path))
    if "statevector" in npz:
        arr = npz["statevector"]
    else:
        # fallback: try the first non-metadata key
        keys = [k for k in npz.files if k not in (
            "system", "basis", "vqe_energy", "vqe_energy_physical",
            "hf_energy", "n_qubits", "n_active_electrons", "n_active_orbitals"
        )]
        if not keys:
            raise ValueError(f"No statevector array found in {npz_path}; keys: {npz.files}")
        arr = npz[keys[0]]

    sv = np.asarray(arr).squeeze()

    # If still multi-dimensional, flatten — but require power-of-2 total size
    if sv.ndim > 1:
        sv = sv.reshape(-1)

    total = sv.shape[0]
    # Sanity: total should be a power of two
    if total <= 1 or (total & (total - 1)) != 0:
        raise ValueError(
            f"{npz_path}: flattened size {total} is not a power of 2; cannot interpret as statevector"
        )

    return sv


def sector_check(npz_path: Path, n_qubits: int, n_target: int, sz_target: int) -> dict:
    """Run the independent sector check on one .npz statevector file."""
    try:
        sv = _load_statevector(npz_path)
    except Exception as e:
        return {"sector_check_status": "load_error", "error": str(e),
                "npz_path": str(npz_path)}

    expected_size = 2 ** n_qubits
    if sv.shape[0] != expected_size:
        return {
            "sector_check_status": "shape_mismatch",
            "error": f"statevector length {sv.shape[0]} != 2^{n_qubits}={expected_size}",
            "actual_size": int(sv.shape[0]),
            "expected_size": int(expected_size),
            "npz_path": str(npz_path),
        }

    probs = (np.abs(sv) ** 2).astype(np.float64)
    norm = float(probs.sum())

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

    # Vectorized ⟨N⟩, ⟨2Sz⟩, joint P(target sector)
    # Build masks once, in vectorized fashion: for each qubit q, the bit at
    # position q in the integer-indexed basis state.
    # We use the same bit ordering as the metalloenzyme audit: position q
    # corresponds to (i >> q) & 1 where q=0 is the lowest-weight bit.
    # Then the "bitstring" is format(i, f'0{n}b') which orders bits MSB-first.
    # Within format(i,'0Nb'), the bit at *string* index (N-1-q) is bit q.
    # Match the metalloenzyme audit convention: q even = β (-1), q odd = α (+1)

    n_states = expected_size
    idxs = np.arange(n_states, dtype=np.int64)

    # Per-basis-state electron count: popcount of idx
    # numpy doesn't have a built-in popcount for int64, but bin().count('1')
    # is too slow for 1M states. Use bit-twiddle:
    def _popcount(x: np.ndarray) -> np.ndarray:
        # x is int64 array
        x = x - ((x >> 1) & 0x5555555555555555)
        x = (x & 0x3333333333333333) + ((x >> 2) & 0x3333333333333333)
        x = (x + (x >> 4)) & 0x0f0f0f0f0f0f0f0f
        return ((x * 0x0101010101010101) >> 56) & 0xff

    N_per_state = _popcount(idxs)

    # 2Sz = (# of alpha set) - (# of beta set)
    # alpha qubits are q=1,3,5,...; beta qubits are q=0,2,4,...
    # Mask for alpha bits: 0b...10101010 for n_qubits bits
    alpha_mask = int(0)
    beta_mask = int(0)
    for q in range(n_qubits):
        if q % 2 == 1:
            alpha_mask |= (1 << q)
        else:
            beta_mask |= (1 << q)

    n_alpha_per_state = _popcount(idxs & alpha_mask)
    n_beta_per_state = _popcount(idxs & beta_mask)
    twoSz_per_state = (n_alpha_per_state.astype(np.int64)
                       - n_beta_per_state.astype(np.int64))

    N_expect = float(np.sum(probs * N_per_state))
    twoSz_expect = float(np.sum(probs * twoSz_per_state))

    # Joint P(N=N_target AND 2Sz=Sz_target)
    target_mask = (N_per_state == n_target) & (twoSz_per_state == sz_target)
    joint_target_p = float(np.sum(probs[target_mask]))

    return {
        "sector_check_status":   "ok",
        "spin_convention":       "odd_alpha (qubit 2i=β, 2i+1=α)",
        "N_target":              int(n_target),
        "twoSz_target":          int(sz_target),
        "N_expect":              N_expect,
        "twoSz_expect":          twoSz_expect,
        "joint_P_target_sector": joint_target_p,
        "dom_p":                 dom_p,
        "dominant_bitstring":    dom_bits,
        "statevector_norm":      norm,
        "n_qubits":              int(n_qubits),
        "npz_path":              str(npz_path),
    }


def update_manifest(manifest_path: Path, output_dir: Path,
                    out_path: Optional[Path] = None, dry_run: bool = False) -> int:
    if not manifest_path.exists():
        print(f"[ERROR] Manifest not found: {manifest_path}", file=sys.stderr)
        return 2

    manifest = json.loads(manifest_path.read_text())
    runs = manifest.get("runs", [])
    aggregates = manifest.get("aggregates", [])

    print(f"[INFO] Loaded {len(runs)} runs from {manifest_path}")
    print(f"[INFO] Spin convention: odd_alpha (qubit 2i=β, 2i+1=α)")
    print()

    new_sector_per_run: dict[str, dict] = {}

    for run in runs:
        label = run.get("run_label", "<unknown>")
        sv_path_str = run.get("statevector_path")
        if not sv_path_str:
            print(f"[SKIP] {label}: no statevector_path in manifest")
            run["independent_sector_check"] = {"sector_check_status": "no_statevector"}
            continue

        sv_path = Path(sv_path_str)
        if not sv_path.exists():
            # try relative to output_dir
            alt = output_dir / sv_path.name
            if alt.exists():
                sv_path = alt
            else:
                print(f"[FAIL] {label}: statevector file not found at "
                      f"{sv_path_str} or {alt}")
                run["independent_sector_check"] = {
                    "sector_check_status": "file_not_found",
                    "path_tried": str(sv_path_str),
                }
                continue

        n_qubits = run.get("n_qubits", 20)
        n_target = run.get("n_active_electrons", 10)
        # Multiplicity 1 (singlet) → 2Sz_target = 0 for all five propulsion systems
        sz_target = 0

        result = sector_check(sv_path, n_qubits, n_target, sz_target)
        run["independent_sector_check"] = result
        new_sector_per_run[label] = result

        status = result.get("sector_check_status")
        if status == "ok":
            print(f"[OK]   {label}:")
            print(f"         ⟨N⟩    = {result['N_expect']:.4f}   (target {n_target})")
            print(f"         ⟨2Sz⟩  = {result['twoSz_expect']:+.4f}   (target {sz_target})")
            print(f"         dom_p  = {result['dom_p']:.6f}")
            print(f"         joint  = {result['joint_P_target_sector']:.6f}")
            print(f"         dominant bitstring = {result['dominant_bitstring']}")
        else:
            print(f"[FAIL] {label}: {status}")
            print(f"         {result.get('error', '')}")

    # Recompute the per-(system, ansatz) aggregate fields that depend on
    # the sector-check results (dominant_bitstrings, bitstring_agreement,
    # independent_sector_clean_all).
    for agg in aggregates:
        sys_label = agg["system"]
        ansatz_type = agg["ansatz_type"]
        matched = [r for r in runs
                   if r["system"] == sys_label
                   and r["ansatz_type"] == ansatz_type
                   and r["status"] == "ok"]

        sec_results = [r.get("independent_sector_check", {}) for r in matched]
        oks = [s for s in sec_results if s.get("sector_check_status") == "ok"]

        if oks:
            dom_bits = [s["dominant_bitstring"] for s in oks]
            agg["dominant_bitstrings"] = dom_bits
            agg["bitstring_agreement"] = (len(set(dom_bits)) == 1)
            n_target = matched[0]["n_active_electrons"]
            agg["independent_sector_clean_all"] = all(
                abs(s["N_expect"] - n_target) < 0.1 and abs(s["twoSz_expect"]) < 0.1
                for s in oks
            )
            # Also surface the per-seed N/2Sz/dom_p/joint for the aggregate
            agg["N_expect_per_seed"] = [s["N_expect"] for s in oks]
            agg["twoSz_expect_per_seed"] = [s["twoSz_expect"] for s in oks]
            agg["dom_p_per_seed"] = [s["dom_p"] for s in oks]
            agg["joint_P_target_per_seed"] = [s["joint_P_target_sector"] for s in oks]
        else:
            agg["dominant_bitstrings"] = []
            agg["bitstring_agreement"] = None
            agg["independent_sector_clean_all"] = None

    # Compact summary table
    print()
    print("=" * 96)
    print("INDEPENDENT SECTOR CHECK — SUMMARY")
    print("=" * 96)
    header = (f"{'run_label':>34s}  {'⟨N⟩':>7s}  {'⟨2Sz⟩':>7s}  "
              f"{'dom_p':>7s}  {'joint_P':>8s}  {'eng_sec':>8s}  {'ind_sec':>8s}")
    print(header)
    print("-" * len(header))
    for run in runs:
        label = run["run_label"]
        sec = run.get("independent_sector_check", {})
        if sec.get("sector_check_status") == "ok":
            n_target = run["n_active_electrons"]
            ind_clean = (abs(sec["N_expect"] - n_target) < 0.1
                         and abs(sec["twoSz_expect"]) < 0.1)
            print(f"{label:>34s}  {sec['N_expect']:7.4f}  {sec['twoSz_expect']:+7.4f}  "
                  f"{sec['dom_p']:7.4f}  {sec['joint_P_target_sector']:8.4f}  "
                  f"{str(run.get('sector_clean')):>8s}  {str(ind_clean):>8s}")
        else:
            print(f"{label:>34s}  {'ERR':>7s}  {'ERR':>7s}  {'ERR':>7s}  {'ERR':>8s}  "
                  f"{str(run.get('sector_clean')):>8s}  {'fail':>8s}")
    print("=" * 96)
    print()

    if dry_run:
        print("[DRY-RUN] Manifest NOT written.")
        return 0

    # Decide write target and backup
    if out_path is None:
        out_path = manifest_path
        # Make a backup of the original
        bak = manifest_path.with_suffix(manifest_path.suffix + ".pre_sector_check.bak")
        if not bak.exists():
            shutil.copy2(manifest_path, bak)
            print(f"[BACKUP] {manifest_path} -> {bak}")

    out_path.write_text(json.dumps(manifest, indent=2, default=str))
    print(f"[WRITE] Updated manifest written to {out_path}")
    return 0


def main():
    ap = argparse.ArgumentParser(
        description=__doc__,
        formatter_class=argparse.RawDescriptionHelpFormatter,
    )
    ap.add_argument("--output-dir", type=Path, required=True,
                    help="Directory containing .statevector.npz files and the results JSON")
    ap.add_argument("--manifest", type=Path, default=None,
                    help="Manifest path (default: <output-dir>/propulsion_reaudit_results.json)")
    ap.add_argument("--out", type=Path, default=None,
                    help="Write updated manifest here (default: in-place, with .bak)")
    ap.add_argument("--dry-run", action="store_true",
                    help="Show results without writing manifest")
    args = ap.parse_args()

    manifest_path = args.manifest or (args.output_dir / "propulsion_reaudit_results.json")
    return update_manifest(manifest_path, args.output_dir, args.out, args.dry_run)


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