"""
VQE for H2 Molecule Ground State
================================
Finds the ground state energy of the hydrogen molecule (H2) using the
Variational Quantum Eigensolver (VQE) — a hybrid quantum-classical algorithm
that pairs a parameterized quantum circuit with a classical optimizer to
minimize the energy expectation value.

H2 is the simplest molecular system (2 electrons, 1 bond) and serves as the
canonical benchmark for quantum chemistry on near-term quantum computers.
The qubit Hamiltonian is derived via the STO-3G minimal basis set and
parity mapping with two-qubit reduction.

Category: Chemistry
Difficulty: Intermediate
Framework: Qiskit
Qubits: 2
Depth: Variable (depends on ansatz layers)
Gates: RY, CX, H

Hamiltonian (STO-3G, parity mapping, bond length 0.735 A):
    H_elec = c0*I + c1*Z0 + c2*Z1 + c3*Z0Z1 + c4*X0X1

    Coefficients (Hartree):
        c0 = -1.0523    (constant electronic offset)
        c1 =  0.3979    (single-qubit Z on qubit 0)
        c2 = -0.3979    (single-qubit Z on qubit 1)
        c3 = -0.0112    (two-qubit ZZ correlation)
        c4 =  0.1809    (XX coupling — electron exchange)

    Nuclear repulsion energy: V_nn = 0.7199 Ha (at R = 0.735 A)
    Total molecular energy:   E = <H_elec> + V_nn
    Exact ground state:       E0 = -1.1373 Ha (FCI/STO-3G)

Ansatz (hardware-efficient, 1 layer):
    |00> --RY(t0)--*----RY(t2)--
                   |
    |00> --RY(t1)--X----RY(t3)--

    4 parameters: [t0, t1, t2, t3]

Measurement strategy:
    - ZZ basis (computational): Gives <Z0>, <Z1>, <Z0Z1> from one circuit
    - XX basis (H before measure): Gives <X0X1>
    Only 2 circuits needed per energy evaluation (no YY term in this mapping).

Hardware notes:
    - 2 qubits with CX connectivity (any topology)
    - 2 circuits per energy evaluation (ZZ, XX)
    - Depth 4 (RY-CX-RY), well within NISQ coherence budgets
    - Sensitive to readout error — mitigate with calibration matrices
    - COBYLA optimizer is gradient-free and noise-tolerant

References:
    - Peruzzo et al. (2014). "A variational eigenvalue solver on a photonic
      quantum processor." Nature Communications 5, 4213.
    - O'Malley et al. (2016). "Scalable quantum simulation of molecular
      energies on a superconducting qubit processor." Physical Review X 6, 031007.
    - Nielsen & Chuang, "Quantum Computation and Quantum Information", Ch. 4
"""

import numpy as np
from qiskit import QuantumCircuit
from qiskit_aer import AerSimulator
from scipy.optimize import minimize


# ==========================================================================
# Hamiltonian coefficients — H2 in STO-3G basis at equilibrium (0.735 A)
# ==========================================================================
# Derived from parity mapping with two-qubit reduction of the 4-qubit
# Jordan-Wigner Hamiltonian. These coefficients encode the electronic
# structure: kinetic energy, electron-nuclear attraction, and
# electron-electron repulsion. Nuclear repulsion is added separately.

H2_COEFFICIENTS = {
    "I": -1.0523,       # Identity — constant electronic energy offset
    "Z0": 0.3979,       # Qubit 0 occupation: maps to spin-orbital energy
    "Z1": -0.3979,      # Qubit 1 occupation: maps to spin-orbital energy
    "Z0Z1": -0.0112,    # Two-body correlation: electron-electron interaction
    "X0X1": 0.1809,     # Exchange integral: drives electron delocalization
}

NUCLEAR_REPULSION = 0.7199           # Ha — proton-proton repulsion at R = 0.735 A
EXACT_GROUND_STATE_ENERGY = -1.1373  # Ha — FCI/STO-3G total molecular energy
CHEMICAL_ACCURACY = 0.0016           # 1 kcal/mol = 0.0016 Ha


# ==========================================================================
# Exact diagonalization (classical reference)
# ==========================================================================

def compute_exact_ground_state(
    coefficients: dict[str, float] | None = None,
    nuclear_repulsion: float = NUCLEAR_REPULSION,
) -> dict:
    """Compute the exact ground state by diagonalizing the 4x4 Hamiltonian matrix.

    This is the classical reference that VQE aims to match. For H2 (2 qubits)
    we can build and diagonalize the full matrix — but for larger molecules
    this becomes exponentially expensive, which is why VQE exists.

    Args:
        coefficients: Hamiltonian coefficients. Defaults to H2_COEFFICIENTS.
        nuclear_repulsion: Nuclear repulsion energy to add to electronic energy.

    Returns:
        dict with ground_state_energy, all eigenvalues, and ground state vector.
    """
    if coefficients is None:
        coefficients = H2_COEFFICIENTS

    # Pauli matrices
    i2 = np.eye(2)
    z = np.array([[1, 0], [0, -1]])
    x = np.array([[0, 1], [1, 0]])

    # Build 4x4 Hamiltonian matrix in computational basis {|00>, |01>, |10>, |11>}
    h_matrix = (
        coefficients["I"] * np.kron(i2, i2)
        + coefficients["Z0"] * np.kron(i2, z)   # Z on qubit 0 (rightmost)
        + coefficients["Z1"] * np.kron(z, i2)   # Z on qubit 1 (leftmost)
        + coefficients["Z0Z1"] * np.kron(z, z)
        + coefficients["X0X1"] * np.kron(x, x)
    )

    eigenvalues, eigenvectors = np.linalg.eigh(h_matrix)

    return {
        "electronic_energy": eigenvalues[0],
        "ground_state_energy": eigenvalues[0] + nuclear_repulsion,
        "eigenvalues": (eigenvalues + nuclear_repulsion).tolist(),
        "ground_state_vector": eigenvectors[:, 0].tolist(),
    }


# ==========================================================================
# Circuit construction
# ==========================================================================

def create_h2_ansatz(theta: list[float]) -> QuantumCircuit:
    """Create a hardware-efficient ansatz for the H2 molecule.

    The ansatz uses a single entangling layer:
      1. RY rotations on each qubit (parameterized single-qubit gates)
      2. CX entangling gate (creates correlation between qubits)
      3. Second RY rotation layer (adds expressibility)

    This is the minimal ansatz that can reach the exact H2 ground state.
    With 4 free parameters and 1 CX gate, it covers the relevant subspace
    of the 2-qubit Hilbert space.

    Args:
        theta: List of 4 rotation angles [t0, t1, t2, t3] in radians.
               t0, t1 set the initial single-qubit states.
               t2, t3 refine the state after entanglement.

    Returns:
        QuantumCircuit with parameterized ansatz (no measurements).

    Raises:
        ValueError: If theta does not contain exactly 4 elements.
    """
    if len(theta) != 4:
        raise ValueError(f"Expected 4 parameters, got {len(theta)}")

    qc = QuantumCircuit(2)

    # Layer 1: Single-qubit rotations — set initial superposition
    # RY(t) rotates |0> to cos(t/2)|0> + sin(t/2)|1>
    qc.ry(theta[0], 0)  # Qubit 0: parameterized rotation
    qc.ry(theta[1], 1)  # Qubit 1: parameterized rotation

    # Entangling layer: CX creates correlations between qubits
    # This is essential — without entanglement, we can only represent
    # product states, missing the electron correlation energy
    qc.cx(0, 1)

    # Layer 2: Post-entanglement rotations — refine the entangled state
    qc.ry(theta[2], 0)  # Fine-tune qubit 0 after entanglement
    qc.ry(theta[3], 1)  # Fine-tune qubit 1 after entanglement

    return qc


def create_measurement_circuits(
    ansatz: QuantumCircuit,
) -> list[tuple[str, QuantumCircuit]]:
    """Create measurement circuits for each Pauli term in the H2 Hamiltonian.

    The Hamiltonian has 5 terms, but we only need 2 measurement circuits:
      - ZZ basis (computational): Extracts <Z0>, <Z1>, and <Z0Z1> simultaneously
      - XX basis (Hadamard transform): Extracts <X0X1>

    Each Pauli basis requires rotating into the eigenbasis before measuring
    in the computational basis:
      - Z eigenbasis: no rotation needed (already computational)
      - X eigenbasis: apply H (Hadamard)

    Args:
        ansatz: The parameterized circuit (without measurements).

    Returns:
        List of (basis_name, circuit) tuples ready for execution.
    """
    circuits = []

    # ZZ basis — computational basis measurement
    # This single circuit gives us <Z0>, <Z1>, and <Z0Z1>
    zz_circuit = ansatz.copy()
    zz_circuit.measure_all()
    circuits.append(("ZZ", zz_circuit))

    # XX basis — rotate both qubits from X eigenbasis to Z eigenbasis
    # H|+> = |0>, H|-> = |1>, so measuring in Z after H gives X eigenvalues
    xx_circuit = ansatz.copy()
    xx_circuit.h(0)
    xx_circuit.h(1)
    xx_circuit.measure_all()
    circuits.append(("XX", xx_circuit))

    return circuits


# ==========================================================================
# Expectation value computation
# ==========================================================================

def compute_pauli_expectation(counts: dict, shots: int) -> float:
    """Compute the expectation value of a 2-qubit Pauli operator from counts.

    For a Pauli string P = P0 (x) P1, the expectation value is:
        <P> = sum_b (-1)^parity(b) * count(b) / shots

    where parity(b) is the number of 1s in bitstring b. This works because
    Pauli eigenvalues are +/-1, and the parity determines the eigenvalue.

    Args:
        counts: Measurement counts dict, e.g. {'00': 512, '11': 512}.
        shots: Total number of measurements.

    Returns:
        Expectation value in [-1.0, +1.0].
    """
    expectation = 0.0
    for bitstring, count in counts.items():
        parity = (-1) ** bitstring.count("1")
        expectation += parity * count / shots
    return expectation


def compute_energy(
    measurement_results: dict[str, dict],
    shots: int,
    coefficients: dict[str, float] | None = None,
    nuclear_repulsion: float = NUCLEAR_REPULSION,
) -> float:
    """Compute the total H2 energy from measurement results in all Pauli bases.

    Combines expectation values from ZZ and XX measurement circuits with
    the Hamiltonian coefficients, then adds nuclear repulsion:

        E_elec = c0 + c1*<Z0> + c2*<Z1> + c3*<Z0Z1> + c4*<X0X1>
        E_total = E_elec + V_nn

    The ZZ circuit gives three expectation values simultaneously:
        <Z0> = sum_b (-1)^b0 * count(b) / shots
        <Z1> = sum_b (-1)^b1 * count(b) / shots
        <Z0Z1> = sum_b (-1)^(b0+b1) * count(b) / shots

    Args:
        measurement_results: Dict mapping basis name ("ZZ", "XX") to counts.
        shots: Number of shots used per circuit.
        coefficients: Hamiltonian coefficients. Defaults to H2_COEFFICIENTS.
        nuclear_repulsion: Nuclear repulsion energy (Hartree).

    Returns:
        Total molecular energy estimate in Hartree.
    """
    if coefficients is None:
        coefficients = H2_COEFFICIENTS

    # Start with the identity (constant) term
    energy = coefficients["I"]

    # Extract Z0, Z1, Z0Z1 from computational basis measurement
    # Qiskit uses little-endian bit ordering: bitstring "ba" means
    # qubit 0 = a, qubit 1 = b
    zz_counts = measurement_results["ZZ"]
    z0_exp = 0.0
    z1_exp = 0.0
    z0z1_exp = 0.0

    for bitstring, count in zz_counts.items():
        b0 = int(bitstring[-1])  # Last char = qubit 0 (Qiskit convention)
        b1 = int(bitstring[-2])  # Second-to-last = qubit 1
        z0 = 1 - 2 * b0         # Map 0->+1, 1->-1
        z1 = 1 - 2 * b1
        z0_exp += z0 * count / shots
        z1_exp += z1 * count / shots
        z0z1_exp += z0 * z1 * count / shots

    energy += coefficients["Z0"] * z0_exp
    energy += coefficients["Z1"] * z1_exp
    energy += coefficients["Z0Z1"] * z0z1_exp

    # Extract XX expectation value
    xx_exp = compute_pauli_expectation(measurement_results["XX"], shots)
    energy += coefficients["X0X1"] * xx_exp

    # Add nuclear repulsion to get total molecular energy
    energy += nuclear_repulsion

    return energy


# ==========================================================================
# Execution
# ==========================================================================

def run_circuit(theta: list[float] | None = None, shots: int = 1024) -> dict:
    """Execute the VQE ansatz and compute the H2 energy estimate.

    Builds the ansatz with the given parameters, creates measurement circuits
    for all Pauli bases, executes them on the simulator, and combines the
    results into a single energy estimate.

    Args:
        theta: Ansatz parameters (4 angles in radians).
               If None, uses near-optimal values [pi, pi, 0, 0] which prepare
               the Hartree-Fock state |01> (single-determinant approximation).
        shots: Number of measurement shots per basis circuit (default: 1024).

    Returns:
        dict with keys:
            - theta: Parameter values used
            - energy: Total molecular energy estimate (Hartree)
            - exact_ground_state: Reference energy (Hartree)
            - error: |computed - exact| (Hartree)
            - measurements: Raw counts per basis
    """
    # Default: Hartree-Fock state |01> — good starting point, misses
    # ~20 mHa of correlation energy that optimization will recover
    if theta is None:
        theta = [np.pi, np.pi, 0.0, 0.0]

    ansatz = create_h2_ansatz(theta)
    circuits = create_measurement_circuits(ansatz)

    # Execute all measurement circuits
    backend = AerSimulator()
    measurement_results = {}
    for basis_name, circuit in circuits:
        job = backend.run(circuit, shots=shots)
        counts = dict(job.result().get_counts(circuit))
        measurement_results[basis_name] = counts

    energy = compute_energy(measurement_results, shots)

    return {
        "theta": theta,
        "energy": energy,
        "exact_ground_state": EXACT_GROUND_STATE_ENERGY,
        "error": abs(energy - EXACT_GROUND_STATE_ENERGY),
        "measurements": measurement_results,
    }


# ==========================================================================
# Optimization
# ==========================================================================

def optimize_vqe(
    max_iterations: int = 100,
    shots: int = 2048,
    seed: int | None = None,
) -> dict:
    """Run the full VQE optimization loop to find the H2 ground state.

    The classical optimizer (COBYLA) iteratively adjusts the 4 ansatz
    parameters to minimize the energy expectation value. Each iteration
    requires 2 quantum circuit evaluations (ZZ and XX bases).

    COBYLA is used because:
      - Gradient-free: no need for parameter-shift circuits
      - Noise-tolerant: handles shot noise better than gradient methods
      - Low overhead: no auxiliary circuits for gradient estimation

    Args:
        max_iterations: Maximum optimizer iterations (default: 100).
        shots: Shots per circuit per iteration (default: 2048).
        seed: Random seed for reproducible initial parameters.

    Returns:
        dict with keys:
            - optimal_theta: Best parameters found
            - optimal_energy: Lowest energy achieved (Hartree)
            - exact_ground_state: Reference energy (Hartree)
            - error: |optimal - exact| (Hartree)
            - chemical_accuracy: Whether error < 1.6 mHa
            - iterations: Number of function evaluations
            - history: List of energies at each iteration
    """
    rng = np.random.default_rng(seed)
    theta0 = rng.uniform(0, 2 * np.pi, 4)

    history = []

    def cost_function(theta):
        result = run_circuit(theta.tolist(), shots=shots)
        history.append(result["energy"])
        return result["energy"]

    result = minimize(
        cost_function,
        theta0,
        method="COBYLA",
        options={"maxiter": max_iterations, "rhobeg": 0.5},
    )

    error = abs(result.fun - EXACT_GROUND_STATE_ENERGY)

    return {
        "optimal_theta": result.x.tolist(),
        "optimal_energy": result.fun,
        "exact_ground_state": EXACT_GROUND_STATE_ENERGY,
        "error": error,
        "chemical_accuracy": error < CHEMICAL_ACCURACY,
        "iterations": result.nfev,
        "history": history,
    }


# ==========================================================================
# Verification
# ==========================================================================

def verify_vqe(shots: int = 4096, tolerance: float = 0.03) -> dict:
    """Verify the VQE circuit produces a physically correct H2 ground state.

    Runs three checks:
      1. Hartree-Fock energy (default params) is within tolerance of expected
         value (-1.117 Ha = exact - correlation energy)
      2. Variational principle holds: computed energy >= exact ground state
         (within shot noise — we allow a small margin below)
      3. Full optimization converges close to the exact ground state

    Args:
        shots: Measurement shots for statistical confidence (default: 4096).
        tolerance: Energy tolerance in Hartree for check 1 (default: 0.03).

    Returns:
        dict with keys: passed (bool), checks (list), energy, optimization.
    """
    checks = []
    passed = True

    # ---- Check 1: Hartree-Fock parameters give reasonable energy ----
    # [pi, pi, 0, 0] prepares |01> (HF state), expected energy ~ -1.117 Ha
    result = run_circuit(theta=[np.pi, np.pi, 0.0, 0.0], shots=shots)
    energy = result["energy"]
    hf_expected = -1.117  # Hartree-Fock energy (no correlation)
    hf_error = abs(energy - hf_expected)

    check_hf = {
        "name": "hartree_fock_energy",
        "passed": hf_error <= tolerance,
        "detail": (
            f"E_HF = {energy:.4f} Ha, expected ~{hf_expected} Ha, "
            f"error = {hf_error:.4f} Ha (tolerance: {tolerance})"
        ),
    }
    checks.append(check_hf)
    if not check_hf["passed"]:
        passed = False

    # ---- Check 2: Variational principle — E >= E0 (with noise margin) ----
    # Shot noise can push the estimate slightly below E0, so we allow
    # a margin of 2*tolerance below the exact value
    noise_margin = 2 * tolerance
    check_variational = {
        "name": "variational_principle",
        "passed": energy >= EXACT_GROUND_STATE_ENERGY - noise_margin,
        "detail": (
            f"E = {energy:.4f} Ha >= E0 - margin = "
            f"{EXACT_GROUND_STATE_ENERGY - noise_margin:.4f} Ha"
        ),
    }
    checks.append(check_variational)
    if not check_variational["passed"]:
        passed = False

    # ---- Check 3: Optimization converges near exact ground state ----
    opt_result = optimize_vqe(max_iterations=150, shots=shots, seed=42)
    opt_error = opt_result["error"]

    # Allow 10x chemical accuracy (16 mHa) to account for shot noise
    check_optimization = {
        "name": "optimization_convergence",
        "passed": opt_error < CHEMICAL_ACCURACY * 10,
        "detail": (
            f"Optimized E = {opt_result['optimal_energy']:.4f} Ha, "
            f"error = {opt_error:.4f} Ha, "
            f"iterations = {opt_result['iterations']}, "
            f"chemical accuracy (1.6 mHa): "
            f"{'YES' if opt_result['chemical_accuracy'] else 'NO (within 16 mHa threshold)'}"
        ),
    }
    checks.append(check_optimization)
    if not check_optimization["passed"]:
        passed = False

    return {
        "passed": passed,
        "energy": energy,
        "optimization": opt_result,
        "shots": shots,
        "tolerance": tolerance,
        "checks": checks,
    }


# ==========================================================================
# Main — interactive exploration
# ==========================================================================

def create_circuit() -> QuantumCircuit:
    """Zero-arg entry point for the QubitHub Qiskit runner.

    Returns the H2 hardware-efficient ansatz at the Hartree-Fock reference
    state |01> (theta=[π, π, 0, 0]). This is the classical mean-field
    baseline that VQE optimization improves on by recovering ~20 mHa of
    electron correlation energy (see ``optimize_vqe()`` in this file).
    """
    return create_h2_ansatz([np.pi, np.pi, 0.0, 0.0])


if __name__ == "__main__":
    print("VQE for H2 Molecule Ground State")
    print("=" * 50)
    print("Finding the ground state energy of hydrogen (H2)")
    print("using a 2-qubit hardware-efficient ansatz.\n")

    # Classical reference
    exact = compute_exact_ground_state()
    print(f"Classical reference (exact diagonalization):")
    print(f"  Ground state energy: {exact['ground_state_energy']:.4f} Ha")
    print(f"  Electronic energy:   {exact['electronic_energy']:.4f} Ha")
    print(f"  Nuclear repulsion:   {NUCLEAR_REPULSION:.4f} Ha")

    # Single evaluation with Hartree-Fock parameters
    print(f"\nStep 1: Hartree-Fock state (no correlation)")
    print("-" * 50)
    result = run_circuit()
    print(f"  Parameters: {[f'{t:.3f}' for t in result['theta']]}")
    print(f"  Computed energy:   {result['energy']:>8.4f} Ha")
    print(f"  Exact ground state:{EXACT_GROUND_STATE_ENERGY:>8.4f} Ha")
    print(f"  Error:             {result['error']:>8.4f} Ha")
    print(f"  (Missing ~20 mHa of correlation energy)")

    # Show measurement breakdown
    print(f"\n  Measurement counts:")
    for basis, counts in result["measurements"].items():
        total = sum(counts.values())
        sorted_counts = sorted(counts.items(), key=lambda x: -x[1])
        states_str = ", ".join(f"|{s}>:{c}" for s, c in sorted_counts[:4])
        print(f"    {basis}: {states_str}  (total: {total})")

    # Run full optimization
    print(f"\nStep 2: Full VQE optimization (COBYLA, max 100 iterations)")
    print("-" * 50)
    opt = optimize_vqe(max_iterations=100, shots=2048, seed=42)
    print(f"  Optimal energy:    {opt['optimal_energy']:>8.4f} Ha")
    print(f"  Exact ground state:{EXACT_GROUND_STATE_ENERGY:>8.4f} Ha")
    print(f"  Error:             {opt['error']:>8.4f} Ha")
    print(f"  Chemical accuracy: {'YES' if opt['chemical_accuracy'] else 'NO'}")
    print(f"  Iterations:        {opt['iterations']}")

    # Energy convergence (ASCII sparkline)
    if opt["history"]:
        hist = opt["history"]
        bars = "".join(
            "#" if abs(e - opt["optimal_energy"]) < 0.01 else "." for e in hist
        )
        print(f"  Convergence:       [{bars[:60]}]")

    # Verification
    print(f"\nStep 3: Verification")
    print("-" * 50)
    v = verify_vqe()
    for check in v["checks"]:
        symbol = "PASS" if check["passed"] else "FAIL"
        print(f"  [{symbol}] {check['name']}: {check['detail']}")
    print(f"\nOverall: {'PASSED' if v['passed'] else 'FAILED'}")