knowledge_base.py

"""
knowledge_base.py
Single source of truth for the Quantum Algorithm Recommendation Engine (QARE).

Two consumers:
  1) generate_dataset.py  -> uses recommend() as the GROUND-TRUTH labeler.
  2) evaluation/baseline.py -> uses recommend() as the RULE-BASED baseline.

Facts are grounded in standard complexity-theory / NISQ-era hardware knowledge
(Nielsen & Chuang; Qiskit textbook; PennyLane demos; Preskill 'NISQ' 2018;
Shor 1994; Grover 1996; Farhi QAOA 2014; Peruzzo/McClean VQE 2014; Harrow-
Hassidim-Lloyd 2009). No copyrighted text is reproduced; only structured facts.
"""
from __future__ import annotations
from dataclasses import dataclass, field
from typing import Callable
import math

# --------------------------------------------------------------------------- #
# Problem taxonomy
# --------------------------------------------------------------------------- #
PROBLEM_TYPES = [
    "integer_factoring",
    "discrete_log",
    "unstructured_search",
    "combinatorial_optimization",   # MaxCut, TSP-like, portfolio
    "ground_state_energy",          # quantum chemistry / materials
    "eigenvalue_estimation",        # spectra, phase
    "linear_system",                # Ax=b
    "sampling",                     # boson/thermal/prob sampling
    "classification",               # supervised ML
    "graph_connectivity",           # search on graphs
    "simulation_dynamics",          # Hamiltonian time evolution
]

HARDWARE_TYPES = [
    "superconducting",   # IBM, Google
    "trapped_ion",       # IonQ, Quantinuum
    "neutral_atom",      # QuEra, Pasqal
    "photonic",          # Xanadu, PsiQuantum
    "annealer",          # D-Wave
    "simulator",         # statevector / classical sim
    "fault_tolerant",    # hypothetical FT device w/ logical qubits
]

NOISE_LEVELS = ["none", "low", "medium", "high"]  # none == FT / ideal sim


# --------------------------------------------------------------------------- #
# Algorithm records
# --------------------------------------------------------------------------- #
@dataclass
class Algo:
    name: str
    category: str                       # e.g. "fault-tolerant", "variational", "annealing", "classical"
    solves: list[str]                   # PROBLEM_TYPES it targets
    nisq_friendly: bool                 # runs meaningfully on noisy near-term HW
    needs_fault_tolerance: bool
    hybrid: bool                        # classical/quantum loop
    # qubit requirement as a function of problem size n (returns int)
    qubits_fn: Callable[[int], int]
    # depth class as a function of n (returns rough gate depth); big => needs FT
    depth_fn: Callable[[int], int]
    hardware_fit: list[str]
    advantages: list[str]
    limitations: list[str]
    references: list[str] = field(default_factory=list)


def _poly(a, b=0, c=0):
    return lambda n: int(a * n * n + b * n + c) + 1


ALGOS: dict[str, Algo] = {
    "Shor": Algo(
        "Shor's Algorithm", "fault-tolerant",
        ["integer_factoring", "discrete_log"],
        nisq_friendly=False, needs_fault_tolerance=True, hybrid=False,
        qubits_fn=lambda n: 2 * n + 3,           # ~2n logical qubits for n-bit number
        depth_fn=_poly(1, 0, 0),                 # O(n^2 log n) ~ very deep
        hardware_fit=["fault_tolerant", "simulator"],
        advantages=["Exponential speedup over GNFS for factoring/DLP"],
        limitations=["Requires many low-error logical qubits", "Impractical on NISQ"],
        references=["Shor 1994", "Nielsen & Chuang ch.5"],
    ),
    "Grover": Algo(
        "Grover's Algorithm", "fault-tolerant",
        ["unstructured_search", "graph_connectivity"],
        nisq_friendly=False, needs_fault_tolerance=True, hybrid=False,
        qubits_fn=lambda n: n + 1,
        depth_fn=lambda n: int(math.pi / 4 * math.sqrt(2 ** min(n, 30))) + 1,
        hardware_fit=["fault_tolerant", "trapped_ion", "simulator"],
        advantages=["Quadratic speedup for unstructured search"],
        limitations=["Deep iterate; quadratic gain erased by NISQ noise",
                     "Rarely practical below fault tolerance"],
        references=["Grover 1996"],
    ),
    "QAOA": Algo(
        "QAOA", "variational",
        ["combinatorial_optimization", "graph_connectivity"],
        nisq_friendly=True, needs_fault_tolerance=False, hybrid=True,
        qubits_fn=lambda n: n,                   # 1 qubit per binary var
        depth_fn=lambda n: 2 * n,                # p-layer, shallow-ish
        hardware_fit=["superconducting", "trapped_ion", "neutral_atom", "simulator"],
        advantages=["Shallow tunable depth (p layers)", "NISQ-compatible"],
        limitations=["Barren plateaus at depth", "No proven speedup",
                     "Often matched by classical heuristics"],
        references=["Farhi et al. 2014"],
    ),
    "VQE": Algo(
        "VQE", "variational",
        ["ground_state_energy", "eigenvalue_estimation", "simulation_dynamics"],
        nisq_friendly=True, needs_fault_tolerance=False, hybrid=True,
        qubits_fn=lambda n: n,                   # ~1 qubit per spin-orbital
        depth_fn=lambda n: 4 * n,
        hardware_fit=["superconducting", "trapped_ion", "neutral_atom", "simulator"],
        advantages=["Leading NISQ chemistry method", "Shallow ansatz options"],
        limitations=["Barren plateaus", "Measurement overhead",
                     "Optimizer can stall"],
        references=["Peruzzo et al. 2014", "McClean et al. 2016"],
    ),
    "QPE": Algo(
        "Quantum Phase Estimation", "fault-tolerant",
        ["eigenvalue_estimation", "ground_state_energy", "simulation_dynamics"],
        nisq_friendly=False, needs_fault_tolerance=True, hybrid=False,
        qubits_fn=lambda n: n + 8,               # system + ancilla precision qubits
        depth_fn=_poly(0, 8, 0),
        hardware_fit=["fault_tolerant", "simulator"],
        advantages=["High-precision eigenvalues", "Backbone of many FT algos"],
        limitations=["Deep controlled-U", "Needs fault tolerance"],
        references=["Kitaev 1995", "Nielsen & Chuang ch.5"],
    ),
    "HHL": Algo(
        "HHL", "fault-tolerant",
        ["linear_system"],
        nisq_friendly=False, needs_fault_tolerance=True, hybrid=False,
        qubits_fn=lambda n: int(math.log2(max(n, 2))) + 10,
        depth_fn=_poly(0, 12, 0),
        hardware_fit=["fault_tolerant", "simulator"],
        advantages=["Exponential speedup for sparse, well-conditioned Ax=b (with caveats)"],
        limitations=["Strong assumptions (condition number, state prep, readout)",
                     "Impractical on NISQ", "Speedup often not end-to-end"],
        references=["Harrow, Hassidim, Lloyd 2009", "Aaronson 2015 (caveats)"],
    ),
    "QuantumWalk": Algo(
        "Quantum Walk Search", "fault-tolerant",
        ["graph_connectivity", "unstructured_search"],
        nisq_friendly=False, needs_fault_tolerance=True, hybrid=False,
        qubits_fn=lambda n: 2 * int(math.log2(max(n, 2))) + 2,
        depth_fn=_poly(0, 6, 0),
        hardware_fit=["fault_tolerant", "simulator"],
        advantages=["Speedups for element distinctness / graph search"],
        limitations=["Deep circuits", "Needs fault tolerance"],
        references=["Ambainis 2003", "Childs 2009"],
    ),
    "QuantumAnnealing": Algo(
        "Quantum Annealing", "annealing",
        ["combinatorial_optimization"],
        nisq_friendly=True, needs_fault_tolerance=False, hybrid=True,
        qubits_fn=lambda n: n,
        depth_fn=lambda n: 1,                    # analog, no gate depth
        hardware_fit=["annealer"],
        advantages=["Native QUBO/Ising solving", "Thousands of physical qubits available"],
        limitations=["Restricted to QUBO", "Embedding overhead",
                     "No proven asymptotic speedup"],
        references=["Kadowaki & Nishimori 1998", "D-Wave docs"],
    ),
    "QSVM": Algo(
        "Quantum Kernel / VQC (QML)", "variational",
        ["classification"],
        nisq_friendly=True, needs_fault_tolerance=False, hybrid=True,
        qubits_fn=lambda n: n,                   # ~1 qubit per feature
        depth_fn=lambda n: 3 * n,
        hardware_fit=["superconducting", "trapped_ion", "simulator"],
        advantages=["Access to high-dim feature maps"],
        limitations=["No general advantage shown", "Kernel concentration",
                     "Classical ML usually competitive"],
        references=["Havlicek et al. 2019", "Schuld & Killoran 2019"],
    ),
    "GaussianBosonSampling": Algo(
        "Gaussian Boson Sampling", "sampling",
        ["sampling"],
        nisq_friendly=True, needs_fault_tolerance=False, hybrid=False,
        qubits_fn=lambda n: n,                   # modes
        depth_fn=lambda n: n,
        hardware_fit=["photonic"],
        advantages=["Demonstrated sampling advantage on photonic HW"],
        limitations=["Narrow applicability", "Not general-purpose compute"],
        references=["Hamilton et al. 2017", "Zhong et al. 2020"],
    ),
    "Trotter": Algo(
        "Trotterized Hamiltonian Simulation", "digital-simulation",
        ["simulation_dynamics", "ground_state_energy"],
        nisq_friendly=True, needs_fault_tolerance=False, hybrid=False,
        qubits_fn=lambda n: n,
        depth_fn=lambda n: 6 * n,
        hardware_fit=["superconducting", "trapped_ion", "neutral_atom", "simulator"],
        advantages=["Direct simulation of local Hamiltonians", "Tunable accuracy via steps"],
        limitations=["Depth grows with time & accuracy", "Trotter error"],
        references=["Lloyd 1996", "Childs et al. 2018"],
    ),
    "SurfaceCode": Algo(
        "Surface-Code Error Correction", "error-correction",
        [],  # not a solver; recommended in EC scenarios
        nisq_friendly=False, needs_fault_tolerance=True, hybrid=False,
        qubits_fn=lambda n: 1000 * n,            # ~physical per logical, illustrative
        depth_fn=lambda n: n,
        hardware_fit=["superconducting", "neutral_atom"],
        advantages=["High threshold (~1%)", "2D nearest-neighbor layout"],
        limitations=["Large physical-qubit overhead"],
        references=["Fowler et al. 2012"],
    ),
    # Classical fallbacks (the correct answer when quantum isn't practical)
    "Classical": Algo(
        "Classical algorithm", "classical",
        PROBLEM_TYPES,
        nisq_friendly=True, needs_fault_tolerance=False, hybrid=False,
        qubits_fn=lambda n: 0,
        depth_fn=lambda n: 0,
        hardware_fit=["simulator"],
        advantages=["Mature, reliable, no quantum hardware needed"],
        limitations=["No quantum speedup"],
        references=["Cormen et al. (CLRS)", "Gurobi/CPLEX docs"],
    ),
}

# specific classical method names by problem (for nicer reasoning text)
CLASSICAL_METHOD = {
    "integer_factoring": "General Number Field Sieve (GNFS)",
    "discrete_log": "index calculus / Pollard's rho",
    "unstructured_search": "linear scan / hashing",
    "combinatorial_optimization": "simulated annealing / Gurobi (branch-and-bound)",
    "ground_state_energy": "coupled cluster (CCSD(T)) / DMRG",
    "eigenvalue_estimation": "Lanczos / dense LAPACK eigensolver",
    "linear_system": "conjugate gradient / sparse LU",
    "sampling": "MCMC (Metropolis-Hastings)",
    "classification": "gradient-boosted trees / SVM / neural nets",
    "graph_connectivity": "BFS/DFS / union-find",
    "simulation_dynamics": "tensor networks / classical ODE integrators",
}


# --------------------------------------------------------------------------- #
# Problem instance
# --------------------------------------------------------------------------- #
@dataclass
class Problem:
    problem_type: str
    size: int                 # n: bits / variables / orbitals / features / nodes(log)
    available_qubits: int
    noise: str                # NOISE_LEVELS
    max_depth: int
    hardware: str             # HARDWARE_TYPES
    desired_accuracy: float   # 0..1 (target solution quality / precision)


# --------------------------------------------------------------------------- #
# Core recommender (ground truth + baseline)
# --------------------------------------------------------------------------- #
# Which algorithms are candidates for each problem type, in preference order
CANDIDATES = {
    "integer_factoring": ["Shor", "Classical"],
    "discrete_log": ["Shor", "Classical"],
    "unstructured_search": ["Grover", "QuantumWalk", "Classical"],
    "combinatorial_optimization": ["QAOA", "QuantumAnnealing", "Classical"],
    "ground_state_energy": ["VQE", "QPE", "Trotter", "Classical"],
    "eigenvalue_estimation": ["QPE", "VQE", "Classical"],
    "linear_system": ["HHL", "Classical"],
    "sampling": ["GaussianBosonSampling", "Classical"],
    "classification": ["QSVM", "Classical"],
    "graph_connectivity": ["QuantumWalk", "Grover", "QAOA", "Classical"],
    "simulation_dynamics": ["Trotter", "VQE", "QPE", "Classical"],
}


def _feasible(algo: Algo, p: Problem) -> tuple[bool, list[str]]:
    """Return (feasible, reasons_it_fails)."""
    fails = []
    req_q = algo.qubits_fn(p.size)
    req_d = algo.depth_fn(p.size)

    if algo.name == "Classical algorithm":
        return True, []

    if req_q > p.available_qubits:
        fails.append(f"needs ~{req_q} qubits but only {p.available_qubits} available")

    # hardware compatibility
    if p.hardware not in algo.hardware_fit and p.hardware != "simulator":
        fails.append(f"not suited to {p.hardware} hardware")

    # fault tolerance vs noise
    if algo.needs_fault_tolerance and p.noise in ("low", "medium", "high") \
            and p.hardware not in ("fault_tolerant", "simulator"):
        fails.append("requires fault tolerance; current noise is prohibitive")

    # depth budget (skip for annealer analog / simulator ideal)
    if p.hardware not in ("annealer", "simulator") and req_d > p.max_depth:
        fails.append(f"needs depth ~{req_d} but budget is {p.max_depth}")

    # NISQ + high noise kills non-nisq-friendly algos
    if not algo.nisq_friendly and p.noise == "high":
        fails.append("high noise erases the theoretical advantage")

    return (len(fails) == 0), fails


def _confidence(algo: Algo, p: Problem, feasible: bool, n_fails: int) -> float:
    if not feasible:
        return round(max(0.15, 0.4 - 0.1 * n_fails), 2)
    base = 0.9 if algo.category in ("fault-tolerant", "annealing", "digital-simulation") else 0.75
    if algo.hybrid and p.noise in ("low", "none"):
        base += 0.05
    if p.hardware == "simulator":
        base += 0.05
    # accuracy pressure: variational methods lose confidence at very high accuracy demand
    if algo.category == "variational" and p.desired_accuracy > 0.95:
        base -= 0.15
    return round(min(0.98, base), 2)


def recommend(p: Problem) -> dict:
    """Ground-truth recommendation for a Problem. Returns the QARE output schema."""
    cands = CANDIDATES.get(p.problem_type, ["Classical"])
    scored = []
    for key in cands:
        algo = ALGOS[key]
        feas, fails = _feasible(algo, p)
        conf = _confidence(algo, p, feas, len(fails))
        scored.append((key, algo, feas, fails, conf))

    # Prefer a feasible quantum method with highest confidence; else fall to Classical.
    feasible_quantum = [s for s in scored if s[2] and s[0] != "Classical"]
    if feasible_quantum:
        feasible_quantum.sort(key=lambda s: -s[4])
        primary_key, primary, _, _, conf = feasible_quantum[0]
        quantum_practical = True
    else:
        primary_key, primary = "Classical", ALGOS["Classical"]
        conf = 0.9
        quantum_practical = False

    # Build ranked alternatives (exclude primary), keep order by confidence then list order
    alts = []
    for key, algo, feas, fails, c in scored:
        if key == primary_key:
            continue
        label = algo.name
        note = "feasible" if feas else "; ".join(fails)
        alts.append({"algorithm": label, "feasible": feas, "confidence": c, "note": note})

    # Reasoning
    if quantum_practical:
        why = (f"{primary.name} targets {p.problem_type.replace('_',' ')} and fits the "
               f"constraints: ~{primary.qubits_fn(p.size)} qubits (<= {p.available_qubits}), "
               f"depth within budget, and tolerates the stated {p.noise} noise on "
               f"{p.hardware} hardware. " + "; ".join(primary.advantages) + ".")
        limitations = primary.limitations
        hw_req = (f"~{primary.qubits_fn(p.size)} qubits, depth ~{primary.depth_fn(p.size)}, "
                  f"{'fault tolerance required' if primary.needs_fault_tolerance else 'NISQ-compatible'}")
        refs = primary.references
    else:
        method = CLASSICAL_METHOD.get(p.problem_type, "a classical solver")
        blockers = []
        for key, algo, feas, fails, c in scored:
            if key != "Classical" and fails:
                blockers.append(f"{ALGOS[key].name} ({fails[0]})")
        why = (f"No quantum method is practical here: " + "; ".join(blockers[:3]) +
               f". Use {method} on classical hardware until larger, lower-noise or "
               f"fault-tolerant devices are available.")
        limitations = ["No quantum speedup at this problem scale / hardware maturity"]
        hw_req = "classical CPU/GPU; 0 qubits"
        refs = ALGOS["Classical"].references

    return {
        "primary_algorithm": primary.name,
        "confidence": conf,
        "quantum_practical": quantum_practical,
        "reasoning": why,
        "alternatives": alts,
        "hardware_requirements": hw_req,
        "advantages": primary.advantages,
        "limitations": limitations,
        "references": refs,
    }


if __name__ == "__main__":
    import json
    demo = Problem("integer_factoring", size=1024, available_qubits=50,
                   noise="high", max_depth=100, hardware="superconducting",
                   desired_accuracy=0.99)
    print(json.dumps(recommend(demo), indent=2))