"""
Gate-model quantum + exhaustive enumeration on the simulator.

This is the "actually run the quantum algorithm through the simulator, and find
every possible option" layer the local bake-off was missing.

WHY A SECOND FORMULATION
------------------------
geo_optimize.build_coverage_qubo() is exact but uses auxiliary y + slack bits ->
~150+ binary variables for the operational instance. A *statevector* simulator is
2^n amplitudes, so gate-model QAOA caps out around ~24-28 qubits. To run real
QAOA we therefore use a PURE-SITE QUBO: one qubit per ELIGIBLE candidate site, no
auxiliaries. Max-coverage becomes the standard overlap-penalty surrogate:

    maximize  sum_c a_c x_c  -  sum_{c<c'} b_{cc'} x_c x_c'  -  lambda (sum x - p)^2
        a_c     = total risk site c covers
        b_{cc'} = risk covered by BOTH c and c' (overlap; removes double counting)

The surrogate under-counts triple+ overlaps, so we ALWAYS re-score the decoded
selection with the TRUE objective (geo_optimize.coverage_score). The surrogate is
only the search landscape; truth is the judge.

WHAT RUNS HERE (all on the simulator / CPU, laptop-light):
  - exhaustive_true   : enumerate EVERY feasible placement by the true objective
                        -> global optimum + full option landscape
  - exact_qubo        : dimod.ExactSolver on the pure-site QUBO (ground state)
  - qaoa_aer          : real QAOA circuit, Qiskit statevector simulation, p=1,2,3
  - (optional) PennyLane QAOA as a second gate-model path

Real-QPU execution is still geo_optimize/runpod/dwave_qpu.py (D-Wave) — hardware,
not a simulator.
"""
from __future__ import annotations
import sys, os, time, itertools
import numpy as np
from dataclasses import dataclass

sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
from geo_optimize import (PlacementInstance, coverage_score, repair_to_feasible,
                          SolverResult)


# --------------------------------------------------------------------------
# Pure-site QUBO (one qubit per eligible candidate)
# --------------------------------------------------------------------------
@dataclass
class SiteQUBO:
    sites: list          # original candidate indices, one per qubit
    lin: np.ndarray      # (n,)
    quad: np.ndarray     # (n,n) symmetric, zero diag
    const: float
    lam: float


def build_site_qubo(inst: PlacementInstance, lam: float = None) -> SiteQUBO:
    elig = np.where(inst.eligible())[0]
    n = len(elig)
    w = inst.weights
    a = np.array([w[inst.cover[c]].sum() for c in elig])            # own coverage
    b = np.zeros((n, n))
    for i in range(n):
        ci = inst.cover[elig[i]]
        for j in range(i + 1, n):
            overlap = w[ci & inst.cover[elig[j]]].sum()
            b[i, j] = b[j, i] = overlap
    if lam is None:
        lam = float(a.sum()) + 1.0
    p = inst.budget_p
    # E(x) = lin.x + 0.5 x.quad.x + const   (minimize)
    lin = -a + lam * (1.0 - 2.0 * p)
    quad = b + 2.0 * lam
    np.fill_diagonal(quad, 0.0)
    const = lam * p * p
    return SiteQUBO(sites=elig.tolist(), lin=lin, quad=quad, const=const, lam=lam)


def all_energies(sq: SiteQUBO):
    n = len(sq.lin)
    states = np.arange(2 ** n)
    bits = ((states[:, None] >> np.arange(n)[None, :]) & 1).astype(float)  # col i = qubit i
    E = bits @ sq.lin + 0.5 * np.einsum("si,ij,sj->s", bits, sq.quad, bits) + sq.const
    return states, bits, E


def site_qubo_to_ising(sq: SiteQUBO):
    """E(x)=lin.x + 0.5 x.quad.x -> H(s)=sum h_i Z_i + sum_{i<j} J_ij Z_i Z_j, x=(1-s)/2."""
    lin, quad = sq.lin, sq.quad
    h = -0.5 * lin - 0.25 * quad.sum(axis=1)
    J = 0.25 * quad
    return h, J


def decode_bits(bits_row, sq: SiteQUBO) -> list:
    return [sq.sites[i] for i, b in enumerate(bits_row) if b > 0.5]


# --------------------------------------------------------------------------
# 1. Exhaustive enumeration — EVERY feasible option, true objective
# --------------------------------------------------------------------------
def exhaustive_true(inst: PlacementInstance):
    elig = np.where(inst.eligible())[0].tolist()
    p = inst.budget_p
    t0 = time.perf_counter()
    best_sel, best = [], -1.0
    scores = []
    n_options = 0
    for k in range(0, p + 1):
        for combo in itertools.combinations(elig, k):
            sc = coverage_score(inst, combo)
            scores.append(sc)
            n_options += 1
            if sc > best:
                best, best_sel = sc, list(combo)
    ms = (time.perf_counter() - t0) * 1e3
    scores = np.array(scores)
    landscape = {
        "feasible_options_enumerated": n_options,
        "global_optimum_score": round(float(best), 4),
        "n_within_1pct_of_optimum": int((scores >= best * 0.99).sum()),
        "median_option_score": round(float(np.median(scores)), 4),
        "p90_option_score": round(float(np.percentile(scores, 90)), 4),
    }
    res = SolverResult("exhaustive_true", best_sel, float(best), ms, proven_optimal=True,
                       note=f"enumerated EVERY feasible placement ({n_options} options) "
                            f"by the true objective -- global optimum")
    return res, landscape


# --------------------------------------------------------------------------
# 2. dimod ExactSolver on the pure-site QUBO (ground state cross-check)
# --------------------------------------------------------------------------
def exact_qubo(inst: PlacementInstance):
    import dimod
    sq = build_site_qubo(inst)
    t0 = time.perf_counter()
    bqm = dimod.BinaryQuadraticModel("BINARY")
    for i in range(len(sq.lin)):
        bqm.add_variable(i, float(sq.lin[i]))
    for i in range(len(sq.lin)):
        for j in range(i + 1, len(sq.lin)):
            if sq.quad[i, j] != 0:
                bqm.add_interaction(i, j, float(sq.quad[i, j]))
    ss = dimod.ExactSolver().sample(bqm)
    best = ss.first.sample
    bits = np.array([best[i] for i in range(len(sq.lin))])
    sel = repair_to_feasible(inst, decode_bits(bits, sq))
    ms = (time.perf_counter() - t0) * 1e3
    return SolverResult("exact_qubo_groundstate", sel, coverage_score(inst, sel), ms,
                        proven_optimal=False,
                        note="dimod ExactSolver ground state of pure-site QUBO surrogate "
                             "(decoded+repaired, scored on true objective)")


# --------------------------------------------------------------------------
# 3. Real QAOA on the Qiskit statevector simulator
# --------------------------------------------------------------------------
def qaoa_aer(inst: PlacementInstance, p: int = 2, restarts: int = 3,
             maxiter: int = 80, seed: int = 0, max_qubits: int = 18):
    from qiskit import QuantumCircuit
    from qiskit.quantum_info import Statevector
    from scipy.optimize import minimize

    sq = build_site_qubo(inst)
    n = len(sq.lin)
    if n > max_qubits:
        return SolverResult(f"qaoa_p{p}_aer", [], 0.0, 0.0,
                            note=f"skipped: {n} qubits > max_qubits={max_qubits} "
                                 f"(statevector limit; use the annealer/QPU path)"), {}
    h, J = site_qubo_to_ising(sq)
    states, bits, E = all_energies(sq)        # exact QUBO energies for expectation
    gs_idx = int(np.argmin(E))                # surrogate ground state
    jpairs = [(i, j, J[i, j]) for i in range(n) for j in range(i + 1, n) if J[i, j] != 0]

    def circuit(params):
        gammas, betas = params[:p], params[p:]
        qc = QuantumCircuit(n)
        qc.h(range(n))
        for layer in range(p):
            g = gammas[layer]
            for i in range(n):
                if h[i] != 0:
                    qc.rz(2 * g * h[i], i)
            for (i, j, jij) in jpairs:
                qc.rzz(2 * g * jij, i, j)
            for i in range(n):
                qc.rx(2 * betas[layer], i)
        return qc

    def probs(params):
        sv = Statevector.from_instruction(circuit(params))
        return np.abs(sv.data) ** 2           # index s: bit_i = (s>>i)&1

    def expectation(params):
        return float(probs(params) @ E)

    rng = np.random.default_rng(seed)
    t0 = time.perf_counter()
    best_params, best_exp = None, np.inf
    for r in range(restarts):
        x0 = rng.uniform(0, np.pi, size=2 * p) * np.array([0.3] * p + [0.5] * p)
        out = minimize(expectation, x0, method="COBYLA", options={"maxiter": maxiter})
        if out.fun < best_exp:
            best_exp, best_params = out.fun, out.x
    ms = (time.perf_counter() - t0) * 1e3

    pr = probs(best_params)
    p_opt = float(pr[gs_idx])                 # prob mass QAOA puts on surrogate optimum
    # raw best-sampled state (highest probability) and best-by-energy among top states
    top = np.argsort(-pr)[:64]
    raw_sel = decode_bits(bits[int(top[0])], sq)
    raw_score = coverage_score(inst, raw_sel)
    best_sel, best = [], -1.0
    for s in top:
        sel = repair_to_feasible(inst, decode_bits(bits[int(s)], sq))
        sc = coverage_score(inst, sel)
        if sc > best:
            best, best_sel = sc, sel
    diag = {
        "qubits": n,
        "depth_p": p,
        "prob_on_surrogate_optimum": round(p_opt, 4),
        "raw_top1_true_score": round(raw_score, 4),
        "qaoa_plus_repair_true_score": round(best, 4),
        "expectation_energy": round(best_exp, 3),
    }
    return SolverResult(f"qaoa_p{p}_aer", best_sel, float(best), ms,
                        note=f"real QAOA depth-{p}, Qiskit statevector sim, {n} qubits; "
                             f"P(optimum)={p_opt:.3f}; decoded+repaired -> true objective"), diag


# --------------------------------------------------------------------------
# 4. Optional: PennyLane QAOA (second gate-model path)
# --------------------------------------------------------------------------
def qaoa_pennylane(inst: PlacementInstance, p: int = 2, steps: int = 60,
                   max_qubits: int = 16):
    try:
        import pennylane as qml
        from pennylane import numpy as pnp
    except Exception as e:
        return SolverResult(f"qaoa_p{p}_pennylane", [], 0.0, 0.0,
                            note=f"skipped: pennylane unavailable ({e})"), {}
    sq = build_site_qubo(inst)
    n = len(sq.lin)
    if n > max_qubits:
        return SolverResult(f"qaoa_p{p}_pennylane", [], 0.0, 0.0,
                            note=f"skipped: {n} qubits > max_qubits={max_qubits}"), {}
    h, J = site_qubo_to_ising(sq)
    states, bits, E = all_energies(sq)
    dev = qml.device("default.qubit", wires=n)
    coeffs, ops = [], []
    for i in range(n):
        if h[i] != 0:
            coeffs.append(float(h[i])); ops.append(qml.PauliZ(i))
    for i in range(n):
        for j in range(i + 1, n):
            if J[i, j] != 0:
                coeffs.append(float(J[i, j])); ops.append(qml.PauliZ(i) @ qml.PauliZ(j))
    Hc = qml.Hamiltonian(coeffs, ops)

    @qml.qnode(dev)
    def cost(params):
        for w in range(n):
            qml.Hadamard(wires=w)
        for layer in range(p):
            qml.templates.ApproxTimeEvolution(Hc, params[layer], 1)
            for w in range(n):
                qml.RX(2 * params[p + layer], wires=w)
        return qml.expval(Hc)

    t0 = time.perf_counter()
    params = pnp.array([0.3] * p + [0.4] * p, requires_grad=True)
    opt = qml.AdamOptimizer(0.1)
    for _ in range(steps):
        params = opt.step(cost, params)

    @qml.qnode(dev)
    def get_probs(params):
        for w in range(n):
            qml.Hadamard(wires=w)
        for layer in range(p):
            qml.templates.ApproxTimeEvolution(Hc, params[layer], 1)
            for w in range(n):
                qml.RX(2 * params[p + layer], wires=w)
        return qml.probs(wires=range(n))

    pr = np.array(get_probs(params))
    ms = (time.perf_counter() - t0) * 1e3
    top = np.argsort(-pr)[:64]
    best_sel, best = [], -1.0
    for s in top:
        sel = repair_to_feasible(inst, decode_bits(bits[int(s)], sq))
        sc = coverage_score(inst, sel)
        if sc > best:
            best, best_sel = sc, sel
    return SolverResult(f"qaoa_p{p}_pennylane", best_sel, float(best), ms,
                        note=f"PennyLane QAOA depth-{p}, {n} qubits, default.qubit sim"), \
           {"qubits": n, "depth_p": p}