"""
diagnostic_inspect_statevector.py
==================================

Forensic inspection of a saved .npz statevector. Designed to confirm what
happened in the unconstrained propulsion smoke test:

  1. Loads a .statevector.npz file
  2. Computes <N>, <2Sz> two independent ways:
       (a) bit-counting (the method in sector_check_postprocess.py)
       (b) OpenFermion number_operator + 2Sz operator expectation
       Confirms (a) and (b) agree to numerical precision.
  3. Prints the top 10 most-populated basis states with their (N, 2Sz) labels.
  4. Reports the population distribution across (N, 2Sz) sectors.
  5. Compares against the engine's claimed sector_clean from the manifest
     if --manifest is provided.

Use this BEFORE deciding whether the sector contamination seen in the
post-processing pass is a real engine issue or a statevector-load artifact.

Standalone — no dependencies on other audit scripts. Just numpy + optionally
openfermion (for the cross-check; falls back gracefully if not installed).

USAGE
-----
    # Inspect a single statevector
    python diagnostic_inspect_statevector.py \\
        --npz propulsion_reaudit_smoke/AlH3_minimal_uccsd_seed0.statevector.npz \\
        --n-qubits 20 --n-target 10 --sz-target 0

    # Cross-reference against the manifest
    python diagnostic_inspect_statevector.py \\
        --npz propulsion_reaudit_smoke/AlH3_minimal_uccsd_seed0.statevector.npz \\
        --manifest propulsion_reaudit_smoke/propulsion_reaudit_results.json

    # Inspect all 10 smoke statevectors and report sector distributions
    python diagnostic_inspect_statevector.py \\
        --inspect-all propulsion_reaudit_smoke \\
        --top-states 5
"""

from __future__ import annotations

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

import numpy as np


# ---------------------------------------------------------------------------
# Statevector loading (same as sector_check_postprocess.py)
# ---------------------------------------------------------------------------

def load_statevector(npz_path: Path) -> np.ndarray:
    npz = np.load(str(npz_path))
    if "statevector" in npz:
        arr = npz["statevector"]
    else:
        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 in {npz_path}")
        arr = npz[keys[0]]

    sv = np.asarray(arr).squeeze()
    if sv.ndim > 1:
        sv = sv.reshape(-1)
    total = sv.shape[0]
    if total <= 1 or (total & (total - 1)) != 0:
        raise ValueError(f"{npz_path}: size {total} is not a power of 2")
    return sv


def load_metadata(npz_path: Path) -> dict:
    """Get system / basis / n_qubits etc. from the .npz metadata."""
    npz = np.load(str(npz_path), allow_pickle=True)
    meta = {}
    for k in ("system", "basis", "n_qubits", "n_active_electrons",
              "n_active_orbitals", "vqe_energy", "vqe_energy_physical", "hf_energy"):
        if k in npz.files:
            v = npz[k]
            try:
                meta[k] = v.item() if v.shape == () else v.tolist()
            except Exception:
                meta[k] = str(v)
    return meta


# ---------------------------------------------------------------------------
# Independent sector computation - method (a): vectorized bit-counting
# ---------------------------------------------------------------------------

def _popcount(x: np.ndarray) -> np.ndarray:
    x = x - ((x >> 1) & 0x5555555555555555)
    x = (x & 0x3333333333333333) + ((x >> 2) & 0x3333333333333333)
    x = (x + (x >> 4)) & 0x0f0f0f0f0f0f0f0f
    return ((x * 0x0101010101010101) >> 56) & 0xff


def bitcount_sector_check(sv: np.ndarray, n_qubits: int,
                          n_target: int, sz_target: int) -> dict:
    """Method (a): basis-state expansion via bit counting.

    Spin convention: qubit 2i = β (even), qubit 2i+1 = α (odd).
    """
    probs = (np.abs(sv) ** 2).astype(np.float64)
    n_states = sv.shape[0]
    idxs = np.arange(n_states, dtype=np.int64)
    N_per_state = _popcount(idxs)

    alpha_mask = 0
    beta_mask = 0
    for q in range(n_qubits):
        if q % 2 == 1:
            alpha_mask |= (1 << q)
        else:
            beta_mask |= (1 << q)
    n_alpha = _popcount(idxs & alpha_mask)
    n_beta = _popcount(idxs & beta_mask)
    twoSz_per_state = n_alpha.astype(np.int64) - n_beta.astype(np.int64)

    N_expect = float(np.sum(probs * N_per_state))
    twoSz_expect = float(np.sum(probs * twoSz_per_state))
    target_mask = (N_per_state == n_target) & (twoSz_per_state == sz_target)
    joint_P = float(np.sum(probs[target_mask]))
    norm = float(probs.sum())

    return {
        "method": "bit_counting",
        "N_expect": N_expect,
        "twoSz_expect": twoSz_expect,
        "joint_P_target": joint_P,
        "norm": norm,
    }


# ---------------------------------------------------------------------------
# Independent sector computation - method (b): OpenFermion Pauli operator
# ---------------------------------------------------------------------------

def openfermion_sector_check(sv: np.ndarray, n_qubits: int,
                             n_target: int, sz_target: int) -> Optional[dict]:
    """Method (b): build number_operator and 2Sz operator via OpenFermion,
    convert to sparse matrix, compute <ψ|N|ψ> and <ψ|2Sz|ψ>.

    Returns None if openfermion isn't installed (graceful fallback).
    """
    try:
        from openfermion.ops import FermionOperator, QubitOperator
        from openfermion.transforms import jordan_wigner
        from openfermion.utils import number_operator
        from openfermion.linalg import get_sparse_operator
    except ImportError:
        return None

    # number operator: Σ_i a†_i a_i  (sums over all spin-orbitals)
    N_fermion = number_operator(n_qubits, parity=-1)
    N_qubit = jordan_wigner(N_fermion)
    N_sparse = get_sparse_operator(N_qubit, n_qubits=n_qubits)

    # 2Sz operator: sum over orbitals of (n_alpha - n_beta)
    # convention: qubit 2i = β, qubit 2i+1 = α (matches our bit-counting)
    twoSz_fermion = FermionOperator()
    for q in range(n_qubits):
        if q % 2 == 1:   # alpha qubit
            twoSz_fermion += FermionOperator(f"{q}^ {q}", 1.0)
        else:            # beta qubit
            twoSz_fermion += FermionOperator(f"{q}^ {q}", -1.0)
    twoSz_qubit = jordan_wigner(twoSz_fermion)
    twoSz_sparse = get_sparse_operator(twoSz_qubit, n_qubits=n_qubits)

    sv_complex = sv.astype(np.complex128)
    N_expect = float(np.real(np.conj(sv_complex) @ (N_sparse @ sv_complex)))
    twoSz_expect = float(np.real(np.conj(sv_complex) @ (twoSz_sparse @ sv_complex)))

    return {
        "method": "openfermion_pauli",
        "N_expect": N_expect,
        "twoSz_expect": twoSz_expect,
        "note": "openfermion sparse matrix expectation",
    }


# ---------------------------------------------------------------------------
# Distribution analysis
# ---------------------------------------------------------------------------

def top_basis_states(sv: np.ndarray, n_qubits: int, k: int = 10) -> list:
    """Return the top-k most-populated basis states with their (N, 2Sz)."""
    probs = (np.abs(sv) ** 2).astype(np.float64)
    top_idx = np.argsort(probs)[-k:][::-1]
    out = []
    for idx in top_idx:
        bits = format(int(idx), f"0{n_qubits}b")
        # count N and 2Sz under odd_alpha convention
        N = bits.count("1")
        twoSz = 0
        for q in range(n_qubits):
            # bit at string position (n_qubits - 1 - q) is bit q
            b = int(bits[n_qubits - 1 - q])
            if b == 1:
                twoSz += 1 if (q % 2 == 1) else -1
        out.append({
            "index": int(idx),
            "bitstring": bits,
            "probability": float(probs[int(idx)]),
            "N": N,
            "twoSz": twoSz,
        })
    return out


def sector_population_distribution(sv: np.ndarray, n_qubits: int) -> dict:
    """Aggregate probability mass by (N, 2Sz) sector.

    Returns a list of {N, twoSz, total_prob} entries sorted by probability,
    with the top sectors shown.
    """
    probs = (np.abs(sv) ** 2).astype(np.float64)
    n_states = sv.shape[0]
    idxs = np.arange(n_states, dtype=np.int64)
    N_per_state = _popcount(idxs)

    alpha_mask = 0
    beta_mask = 0
    for q in range(n_qubits):
        if q % 2 == 1:
            alpha_mask |= (1 << q)
        else:
            beta_mask |= (1 << q)
    n_alpha = _popcount(idxs & alpha_mask)
    n_beta = _popcount(idxs & beta_mask)
    twoSz_per_state = n_alpha.astype(np.int64) - n_beta.astype(np.int64)

    # Group by (N, twoSz)
    sector_keys = N_per_state.astype(np.int64) * 1000 + (twoSz_per_state + 500)
    sector_probs = {}
    for k, p in zip(sector_keys.tolist(), probs.tolist()):
        sector_probs[k] = sector_probs.get(k, 0.0) + p

    out = []
    for key, total in sector_probs.items():
        N = key // 1000
        twoSz = (key % 1000) - 500
        if total > 1e-6:
            out.append({"N": int(N), "twoSz": int(twoSz), "total_prob": float(total)})
    out.sort(key=lambda d: -d["total_prob"])
    return out


# ---------------------------------------------------------------------------
# Per-file inspection
# ---------------------------------------------------------------------------

def inspect_one(npz_path: Path, n_qubits: Optional[int], n_target: Optional[int],
                sz_target: int, top_states: int, run_openfermion: bool) -> dict:
    print()
    print("=" * 78)
    print(f"INSPECTING: {npz_path.name}")
    print("=" * 78)

    meta = load_metadata(npz_path)
    if meta:
        print("Metadata from .npz:")
        for k, v in meta.items():
            print(f"  {k:>22s}: {v}")

    # Resolve n_qubits / n_target from metadata if not given
    n_qubits = n_qubits or int(meta.get("n_qubits") or 20)
    n_target = n_target or int(meta.get("n_active_electrons") or 10)

    sv = load_statevector(npz_path)
    print(f"\nStatevector shape: ({sv.shape[0]},) complex {sv.dtype}")
    print(f"Expected size (2^{n_qubits}): {2**n_qubits}")
    if sv.shape[0] != 2**n_qubits:
        print(f"  WARNING: size mismatch! Got {sv.shape[0]}, expected {2**n_qubits}")
        return {"npz_path": str(npz_path), "status": "size_mismatch"}

    # Method (a): bit counting
    print("\n[Method A] Bit-counting expectation values:")
    a = bitcount_sector_check(sv, n_qubits, n_target, sz_target)
    print(f"  norm        : {a['norm']:.10f}  (should be ~1.0)")
    print(f"  <N>         : {a['N_expect']:.6f}   (target {n_target})")
    print(f"  <2Sz>       : {a['twoSz_expect']:+.6f}   (target {sz_target})")
    print(f"  joint P     : {a['joint_P_target']:.6f}   (probability in target sector)")

    # Method (b): openfermion (cross-check)
    if run_openfermion:
        print("\n[Method B] OpenFermion Pauli operator expectation values:")
        b = openfermion_sector_check(sv, n_qubits, n_target, sz_target)
        if b is None:
            print("  (openfermion not available — skipped)")
        else:
            print(f"  <N>         : {b['N_expect']:.6f}   (target {n_target})")
            print(f"  <2Sz>       : {b['twoSz_expect']:+.6f}   (target {sz_target})")
            n_diff = abs(b['N_expect'] - a['N_expect'])
            sz_diff = abs(b['twoSz_expect'] - a['twoSz_expect'])
            print(f"\n[Cross-check A vs B]")
            print(f"  |<N>_A - <N>_B|     = {n_diff:.3e}")
            print(f"  |<2Sz>_A - <2Sz>_B| = {sz_diff:.3e}")
            if n_diff < 1e-6 and sz_diff < 1e-6:
                print(f"  → MATHEMATICAL AGREEMENT (both methods give same answer)")
            else:
                print(f"  → DISAGREEMENT — investigate further")

    # Top basis states
    print(f"\nTop {top_states} most-populated basis states:")
    states = top_basis_states(sv, n_qubits, k=top_states)
    print(f"  {'rank':>4s}  {'prob':>8s}  {'N':>3s}  {'2Sz':>4s}  bitstring")
    print(f"  {'-'*4}  {'-'*8}  {'-'*3}  {'-'*4}  {'-'*n_qubits}")
    for i, s in enumerate(states):
        marker = ""
        if s["N"] == n_target and s["twoSz"] == sz_target:
            marker = " ← TARGET SECTOR"
        elif s["N"] != n_target:
            marker = " ← wrong N"
        elif s["twoSz"] != sz_target:
            marker = " ← wrong 2Sz"
        print(f"  {i+1:>4d}  {s['probability']:.6f}  {s['N']:>3d}  {s['twoSz']:+4d}  {s['bitstring']}{marker}")

    # Sector population distribution
    print(f"\nSector population distribution (P > 1e-6):")
    dist = sector_population_distribution(sv, n_qubits)
    print(f"  {'(N, 2Sz)':>12s}  {'total_prob':>10s}")
    print(f"  {'-'*12}  {'-'*10}")
    for d in dist[:15]:  # show top 15 sectors
        marker = ""
        if d["N"] == n_target and d["twoSz"] == sz_target:
            marker = " ← target"
        print(f"  ({d['N']:>3d}, {d['twoSz']:+3d})  {d['total_prob']:.6f}{marker}")
    if len(dist) > 15:
        print(f"  ... and {len(dist) - 15} more sectors with smaller weight")

    return {"npz_path": str(npz_path), "status": "ok",
            "bitcount": a, "n_qubits": n_qubits, "n_target": n_target}


# ---------------------------------------------------------------------------
# Inspect-all mode
# ---------------------------------------------------------------------------

def inspect_all(output_dir: Path, top_states: int, run_openfermion: bool):
    npz_files = sorted(output_dir.glob("*.statevector.npz"))
    if not npz_files:
        print(f"[ERROR] No .statevector.npz files found in {output_dir}")
        return 2
    print(f"[INFO] Found {len(npz_files)} statevector files in {output_dir}")

    summary = []
    for path in npz_files:
        result = inspect_one(path, None, None, 0, top_states, run_openfermion)
        summary.append(result)

    # Compact final table
    print()
    print("=" * 90)
    print("INSPECT-ALL SUMMARY")
    print("=" * 90)
    print(f"  {'file':>40s}  {'<N>':>8s}  {'<2Sz>':>8s}  {'joint_P':>9s}  status")
    print(f"  {'-'*40}  {'-'*8}  {'-'*8}  {'-'*9}  ------")
    for s in summary:
        name = Path(s["npz_path"]).name.replace(".statevector.npz", "")
        if s["status"] == "ok":
            bc = s["bitcount"]
            print(f"  {name:>40s}  {bc['N_expect']:>8.4f}  {bc['twoSz_expect']:+8.4f}  "
                  f"{bc['joint_P_target']:>9.6f}  ok")
        else:
            print(f"  {name:>40s}  {'ERR':>8s}  {'ERR':>8s}  {'ERR':>9s}  {s['status']}")
    print("=" * 90)
    return 0


def main():
    ap = argparse.ArgumentParser(
        description=__doc__,
        formatter_class=argparse.RawDescriptionHelpFormatter,
    )
    ap.add_argument("--npz", type=Path, default=None,
                    help="Inspect a single .statevector.npz")
    ap.add_argument("--inspect-all", type=Path, default=None,
                    help="Inspect all .statevector.npz files in a directory")
    ap.add_argument("--n-qubits", type=int, default=None)
    ap.add_argument("--n-target", type=int, default=None,
                    help="Target electron count (defaults to .npz metadata)")
    ap.add_argument("--sz-target", type=int, default=0,
                    help="Target 2Sz (default 0 for singlets)")
    ap.add_argument("--top-states", type=int, default=10,
                    help="Number of dominant basis states to show")
    ap.add_argument("--no-openfermion", action="store_true",
                    help="Skip the openfermion cross-check method")
    args = ap.parse_args()

    if args.inspect_all is not None:
        return inspect_all(args.inspect_all, args.top_states, not args.no_openfermion)
    elif args.npz is not None:
        inspect_one(args.npz, args.n_qubits, args.n_target, args.sz_target,
                    args.top_states, not args.no_openfermion)
        return 0
    else:
        ap.print_help()
        return 1


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