"""
QAOA simulation engine using Qiskit statevector simulation.

Key trick: instead of decomposing the Hamiltonian into individual Rz/Rzz gates
for simulation purposes, we pre-compute the classical cost for every basis state,
convert to phases via exp(-i*gamma*cost), and apply the full diagonal unitary in
one shot using qiskit.circuit.library.Diagonal. This is mathematically identical
but far more efficient and noise-free for statevector simulation.
"""

import numpy as np
from qiskit import QuantumCircuit
from qiskit.circuit.library import Diagonal
from qiskit.quantum_info import Statevector


def simulate_qaoa(gamma: float, beta: float, cost_function, num_qubits: int,
                  shots: int = 1024, seed: int | None = None):
    """
    Simulate a 1-layer (p=1) QAOA circuit and return measurement statistics
    together with the exact complex statevector.

    Parameters
    ----------
    gamma : float
        Cost-layer variational angle.
    beta : float
        Mixer-layer variational angle.
    cost_function : callable
        Accepts `num_qubits` binary integer arguments (MSB first) and returns
        a scalar cost value.
    num_qubits : int
        Number of qubits (3 for standard/equality, 4 for inequality).
    shots : int
        Number of measurement shots for probability estimation.
    seed : int | None
        Random seed for shot sampling (ensures reproducibility).

    Returns
    -------
    probs : np.ndarray, shape (2**num_qubits,)
        Empirical probability of each basis state.
    expected_energy : float
        Weighted average cost over all basis states.
    cost_values : np.ndarray, shape (2**num_qubits,)
        Classical cost for each basis state (diagonal of H_C).
    statevector : Statevector
        Full complex statevector after circuit execution (needed for disc viz).
    """
    num_states = 2 ** num_qubits

    # --- Pre-compute diagonal of H_C ---
    cost_values = np.array([
        cost_function(*[int(b) for b in format(i, f"0{num_qubits}b")])
        for i in range(num_states)
    ])

    # --- Build circuit ---
    phases = np.exp(-1j * gamma * cost_values)
    qc = QuantumCircuit(num_qubits)
    qc.h(range(num_qubits))                    # equal superposition
    qc.append(Diagonal(phases), range(num_qubits))  # cost layer
    qc.rx(2 * beta, range(num_qubits))         # mixer layer

    statevector = Statevector.from_instruction(qc)

    # --- Sample shots ---
    if seed is not None:
        np.random.seed(seed)
    counts = statevector.sample_counts(shots=shots)

    probs = np.zeros(num_states)
    for bitstring, count in counts.items():
        probs[int(bitstring, 2)] = count / shots

    expected_energy = float(np.sum(probs * cost_values))

    return probs, expected_energy, cost_values, statevector


def get_statevector_after_init(num_qubits: int) -> Statevector:
    """Return the statevector after Hadamard initialization only (no cost/mixer)."""
    qc = QuantumCircuit(num_qubits)
    qc.h(range(num_qubits))
    return Statevector.from_instruction(qc)


def get_statevector_after_cost(gamma: float, cost_function, num_qubits: int) -> Statevector:
    """Return the statevector after init + cost layer (beta = 0)."""
    return simulate_qaoa(gamma, 0.0, cost_function, num_qubits, shots=1)[3]