"""
circuits.py — circuit builders per action.

Each builder takes a *validated* `QuantumSpec` and returns a Qiskit
`QuantumCircuit` (un-transpiled, already measured where measurement makes
sense). The Bouncer has already enforced bounds by the time we get here, so
builders assume well-formed inputs.

Notation convention: bitstring readout is little-endian in Qiskit (qubit 0 is
the rightmost / least-significant bit). Counts keys reflect this.
"""

from __future__ import annotations

import math
import numpy as np
from qiskit import QuantumCircuit

from .spec import QuantumSpec


# ── RANDOM ───────────────────────────────────────────────────────────────────

def build_random(spec: QuantumSpec) -> QuantumCircuit:
    """Hadamard on every qubit → measure all. True quantum RNG."""
    n = spec.qubits
    qc = QuantumCircuit(n, n)
    qc.h(range(n))
    qc.measure(range(n), range(n))
    return qc


# ── BELL ─────────────────────────────────────────────────────────────────────

def build_bell(spec: QuantumSpec) -> QuantumCircuit:
    """One Bell pair per requested pair: H(0) → CNOT(0,1), measured."""
    pairs = spec.pairs
    n = pairs * 2
    qc = QuantumCircuit(n, n)
    for p in range(pairs):
        a, b = 2 * p, 2 * p + 1
        qc.h(a)
        qc.cx(a, b)
    qc.measure(range(n), range(n))
    return qc


# ── GROVER ───────────────────────────────────────────────────────────────────

def grover_qubits(items: int) -> int:
    """Number of qubits needed to address `items` states."""
    return max(1, int(math.ceil(math.log2(items))))


def grover_iterations(items: int) -> int:
    """
    Optimal Grover iteration count for `items` states (1 marked).

    Uses the exact per-iteration rotation angle θ = arcsin(√(1/N)) and the
    standard optimum k = round(π/(4θ) − ½), clamped to ≥1. The cruder
    `round(π/4·√N)` is off for small N — e.g. it returns 2 for N=4, which
    over-rotates back to the classical baseline; the true optimum there is 1
    (and achieves ~100% success).
    """
    n = grover_qubits(items)
    N = 2 ** n
    theta = math.asin(math.sqrt(1.0 / N))
    k = round(math.pi / (4 * theta) - 0.5)
    return max(1, int(k))


def _diffuser(n: int) -> QuantumCircuit:
    """The Grover diffusion operator on n qubits (in-place)."""
    qc = QuantumCircuit(n, name="diffuser")
    qc.h(range(n))
    qc.x(range(n))

    # Multi-controlled Z built from H + multi-controlled X + H.
    qc.h(n - 1)
    qc.mcx(list(range(n - 1)), n - 1)
    qc.h(n - 1)

    qc.x(range(n))
    qc.h(range(n))
    return qc


def build_grover(spec: QuantumSpec) -> QuantumCircuit:
    """
    Grover search over `items` states.

    We use exactly the qubits needed to address `items` states
    (ceil(log2(items))), mark the highest-index state |11..1⟩ as the target,
    and apply the optimal number of Grover iterations.
    """
    items = spec.items
    n = grover_qubits(items)
    qc = QuantumCircuit(n, n)

    # Initialize uniform superposition.
    qc.h(range(n))

    # Oracle: mark |11..1⟩ via a multi-controlled Z (phase flip).
    qc.h(n - 1)
    qc.mcx(list(range(n - 1)), n - 1)
    qc.h(n - 1)

    # Optimal iteration count via the exact rotation-angle formula.
    iterations = grover_iterations(items)
    diff = _diffuser(n)
    for _ in range(iterations):
        qc.compose(diff, inplace=True)

    qc.measure(range(n), range(n))
    return qc


# ── QAOA ─────────────────────────────────────────────────────────────────────

def build_qaoa(spec: QuantumSpec) -> QuantumCircuit:
    """
    QAOA for MaxCut on a ring graph, depth p = spec.depth.

    The MaxCut Hamiltonian for a ring of `nodes` vertices is a sum of ZZ
    interactions over ring edges. We use the standard QAOA ansatz:

        |γ,β⟩ = e^{-i β H_M} e^{-i γ H_C}  (repeated p times)

    with H_M = Σ X_i (mixer) and H_C = Σ_{(i,j)∈ring} Z_i Z_j (cost).

    Angles γ, β are fixed heuristic values (γ≈0.8, β≈0.4) suitable for a demo.
    A real deployment would optimise them; here we expose a deterministic,
    reproducible circuit.
    """
    nodes = spec.nodes
    p = spec.depth
    qc = QuantumCircuit(nodes, nodes)

    gamma, beta = 0.8, 0.4  # heuristic demo angles

    # Initial state: uniform superposition (H_M ground state).
    qc.h(range(nodes))

    for _ in range(p):
        # Cost unitaries: ZZ on each ring edge.
        for i in range(nodes):
            j = (i + 1) % nodes
            qc.rzz(2 * gamma, i, j)
        # Mixer unitaries.
        qc.rx(2 * beta, range(nodes))

    qc.measure(range(nodes), range(nodes))
    return qc


# ── VQE (2-qubit H₂) ─────────────────────────────────────────────────────────
#
# Ported from the original `quantum_bouncer.py`. This is a self-contained
# 2-qubit variational eigensolver with a UCCSD-inspired ansatz and a grid
# search over θ. The measured circuit returned here is the *ansatz at the
# optimal θ*; the energy optimisation itself lives in `engine.py` because it
# needs to run many circuits and aggregate counts.

def vqe_ansatz(theta: float) -> QuantumCircuit:
    """UCCSD-inspired 2-qubit H₂ trial state, measured."""
    qc = QuantumCircuit(2, 2)
    qc.x(1)                 # Hartree-Fock reference |01⟩
    qc.ry(theta, 0)         # single-excitation rotation
    qc.cx(0, 1)             # entangling gate
    qc.measure([0, 1], [0, 1])
    return qc


def h2_ground_state_energy(distance: float) -> float:
    """Analytical STO-3G H₂ ground-state energy approximation [Ha]."""
    Re = 0.735
    De = 0.1372
    a = 4.0
    E_min = -1.1372
    return E_min + De * (1.0 - np.exp(-a * (distance - Re))) ** 2


# ── Dispatch ─────────────────────────────────────────────────────────────────

# Actions that yield a single, ready-to-run circuit via a builder.
BUILDERS = {
    "RANDOM": build_random,
    "BELL": build_bell,
    "GROVER": build_grover,
    "QAOA": build_qaoa,
}