src/engine.py

"""
engine.py — transpile + execute + statistics.

The execution backend for the deterministic pipeline:
  * transpile each circuit (optimisation level 1, CPU-friendly),
  * run it on the local AerSimulator,
  * compute per-action statistics (expectation values, success probability,
    zero-noise extrapolation where applicable).

VQE is special: it runs a grid search over the ansatz parameter θ, picking the
lowest-energy configuration, then reports the converged energy. This is a
faithful port of the original `run_vqe_circuit`.

Quantum hardware submission (IBM / Quantinuum / Braket) is a later phase; here
everything runs on the local CPU simulator, and the result always carries a
`backend` label for the receipt.
"""

from __future__ import annotations

import math
from dataclasses import dataclass, field

import numpy as np
from qiskit import QuantumCircuit, transpile
from qiskit_aer import AerSimulator

from .spec import QuantumSpec
from . import circuits


BACKEND_LABEL = "AerSimulator (local CPU)"
DEFAULT_SHOTS = 1000
VQE_SHOTS = 2000


@dataclass
class ExecutionResult:
    """Uniform result envelope across all actions."""
    action: str
    backend: str
    shots: int
    qubits: int
    raw_counts: dict
    statistics: dict = field(default_factory=dict)
    extra: dict = field(default_factory=dict)  # action-specific payload


# ── Helpers ──────────────────────────────────────────────────────────────────

def _run(qc: QuantumCircuit, shots: int = DEFAULT_SHOTS) -> dict:
    """Transpile + run a circuit once, return counts."""
    sim = AerSimulator()
    tqc = transpile(qc, sim, optimization_level=1)
    return sim.run(tqc, shots=shots).result().get_counts()


def _bitstring_len(counts: dict) -> int:
    """Number of bits in each measurement outcome key (handle spaces)."""
    key = next(iter(counts), "0")
    return len(key.replace(" ", ""))


def _parity(bitstring: str) -> int:
    """+1 if even number of 1s, else -1. Used for ZZ expectation values."""
    return 1 if (bitstring.count("1") % 2 == 0) else -1


def _expectation_zz(counts: dict) -> float:
    """⟨Z⊗Z⟩ over all qubits from a counts dict."""
    total = sum(counts.values())
    return sum(c * _parity(b) for b, c in counts.items()) / total


# ── Per-action execution ─────────────────────────────────────────────────────

def _exec_random(spec, qc, shots) -> ExecutionResult:
    counts = _run(qc, shots)
    # Shannon entropy of the distribution, in bits — measures "how random".
    total = sum(counts.values())
    probs = np.array([c / total for c in counts.values()])
    entropy = float(-(probs * np.log2(np.maximum(probs, 1e-12))).sum())
    max_entropy = float(spec.qubits)  # uniform over 2^n states
    return ExecutionResult(
        action="RANDOM", backend=BACKEND_LABEL, shots=shots, qubits=spec.qubits,
        raw_counts=counts,
        statistics={"entropy_bits": entropy,
                    "max_entropy_bits": max_entropy,
                    "uniformity": entropy / max_entropy if max_entropy else 0.0},
    )


def _exec_bell(spec, qc, shots) -> ExecutionResult:
    counts = _run(qc, shots)
    # For ideal Bell pairs, only |00..0011..11> style correlated outcomes appear.
    # Correlation: fraction of shots where each pair measured equal bits.
    n_pairs = spec.pairs
    total = sum(counts.values())
    correlated = 0
    for b, c in counts.items():
        bits = b.replace(" ", "")
        ok = True
        for p in range(n_pairs):
            a = bits[-(2 * p + 1)]
            bb = bits[-(2 * p + 2)]
            if a != bb:
                ok = False
                break
        if ok:
            correlated += c
    return ExecutionResult(
        action="BELL", backend=BACKEND_LABEL, shots=shots, qubits=spec.pairs * 2,
        raw_counts=counts,
        statistics={"correlation": correlated / total,
                    "ideal_correlation": 1.0,
                    "entangled_pairs": n_pairs},
    )


def _exec_grover(spec, qc, shots) -> ExecutionResult:
    counts = _run(qc, shots)
    n = circuits.grover_qubits(spec.items)
    iterations = circuits.grover_iterations(spec.items)
    target = "1" * n
    total = sum(counts.values())
    hits = counts.get(target, 0)
    return ExecutionResult(
        action="GROVER", backend=BACKEND_LABEL, shots=shots, qubits=n,
        raw_counts=counts,
        statistics={"target_state": target,
                    "success_probability": hits / total,
                    "classical_baseline": 1 / (2 ** n)},
        extra={"iterations": iterations},
    )


def _exec_qaoa(spec, qc, shots) -> ExecutionResult:
    counts = _run(qc, shots)
    nodes = spec.nodes
    total = sum(counts.values())
    # MaxCut value for each measured bitstring on a ring graph.
    edges = [(i, (i + 1) % nodes) for i in range(nodes)]

    def _cut(bitstring: str) -> int:
        bits = bitstring.replace(" ", "")
        # bitstring is little-endian; index into it accordingly.
        val = 0
        for i, j in edges:
            bi = bits[-(i + 1)]
            bj = bits[-(j + 1)]
            if bi != bj:
                val += 1
        return val

    cuts = {_cut(b): c for b, c in counts.items()}
    best_cut = max(cuts)
    approx_ratio = sum(k * v for k, v in cuts.items()) / total / len(edges)
    return ExecutionResult(
        action="QAOA", backend=BACKEND_LABEL, shots=shots, qubits=nodes,
        raw_counts=counts,
        statistics={"best_cut": best_cut,
                    "max_possible_cut": len(edges),
                    "approximation_ratio": approx_ratio,
                    "depth_p": spec.depth},
        extra={"edges": edges},
    )


# ── VQE (grid search) ────────────────────────────────────────────────────────
# Ported faithfully from the original quantum_bouncer.py: same Hamiltonian
# coefficients, same nuclear-repulsion conversion, same 50-point θ grid.

_VQE_G = (-0.4804, 0.3435, -0.4347, 0.5716)


def _vqe_expectation(counts: dict, distance: float) -> float:
    g0, g1, g2, g3 = _VQE_G
    r_bohr = distance * 1.88973
    v_nuc = 1.0 / r_bohr
    total = sum(counts.values())
    energy = 0.0
    for bs, cnt in counts.items():
        bs = bs.zfill(2)
        z0 = 1 - 2 * int(bs[1])
        z1 = 1 - 2 * int(bs[0])
        energy += (g0 + g1 * z0 + g2 * z1 + g3 * z0 * z1) * cnt
    return energy / total + v_nuc


def _exec_vqe(spec, _qc, shots) -> ExecutionResult:
    thetas = np.linspace(0, 2 * np.pi, 50, endpoint=False)
    best_e, best_counts, best_theta = float("inf"), {}, 0.0
    for theta in thetas:
        counts = _run(circuits.vqe_ansatz(float(theta)), shots)
        e = _vqe_expectation(counts, spec.distance)
        if e < best_e:
            best_e, best_counts, best_theta = e, counts, float(theta)

    return ExecutionResult(
        action="VQE", backend=BACKEND_LABEL, shots=shots, qubits=2,
        raw_counts=best_counts,
        statistics={"vqe_converged_energy_ha": round(best_e, 6),
                    "analytical_energy_ha": round(
                        circuits.h2_ground_state_energy(spec.distance), 6),
                    "optimal_theta_rad": round(best_theta, 4)},
        extra={"bond_distance_angstrom": spec.distance},
    )


# ── Public API ───────────────────────────────────────────────────────────────

_EXECUTORS = {
    "RANDOM": _exec_random,
    "BELL": _exec_bell,
    "GROVER": _exec_grover,
    "QAOA": _exec_qaoa,
    "VQE": _exec_vqe,
}


def execute(spec: QuantumSpec, qc: QuantumCircuit,
            shots: int | None = None) -> ExecutionResult:
    """
    Execute a validated, lowered circuit for the given spec.

    `qc` is the circuit from the lowering pass. For VQE it is the θ=0 ansatz
    template; the executor re-parameterises internally during the grid search.
    """
    if shots is None:
        shots = VQE_SHOTS if spec.action == "VQE" else DEFAULT_SHOTS
    return _EXECUTORS[spec.action](spec, qc, shots)