
import numpy as np
import time

# ─────────────────────────────────────────────────────────────────
#  PKTron Benchmark Suite
#
#  Measures speed and accuracy of PKTron backends.
#  Compares against reference values.
#
#  Usage:
#      from pktron import run_benchmarks
#      run_benchmarks()
# ─────────────────────────────────────────────────────────────────

def _make_bell_circuit(n=2):
    from .quantum import QuantumCircuit
    qc = QuantumCircuit(n)
    qc.h(0)
    for i in range(n-1):
        qc.cnot(i, i+1)
    return qc

def _make_qft_circuit(n=4):
    from .quantum import QuantumCircuit
    from .algorithms import qft
    qc = QuantumCircuit(n)
    qft(qc)
    return qc

def _make_random_circuit(n=6, depth=10):
    from .quantum import QuantumCircuit
    import numpy as np
    qc = QuantumCircuit(n)
    gates = ["h","x","z","s"]
    for _ in range(depth):
        g = np.random.choice(gates)
        q = np.random.randint(0, n)
        getattr(qc, g)(q)
        if n > 1:
            c = np.random.randint(0, n)
            t = (c + 1) % n
            qc.cnot(c, t)
    return qc

def _time_backend(backend, circuit, shots=512, runs=3):
    times = []
    for _ in range(runs):
        t0 = time.time()
        backend.run(circuit, shots=shots)
        times.append(time.time() - t0)
    return float(np.mean(times))

def run_benchmarks(shots=512, verbose=True):
    """
    Run the full PKTron benchmark suite.
    Returns a dict of results.
    """
    from .simulator import AerSimulator
    from .backend import StatevectorSimulator
    from .performance import FullGPUStatevector, AutoBackend, DecisionDiagramBackend

    results = {}

    if verbose:
        print("=" * 55)
        print("  PKTron Benchmark Suite")
        print("=" * 55)

    # ── Bell circuit ──────────────────────────────────────────
    bell = _make_bell_circuit(2)
    backends = {
        "StatevectorSimulator": StatevectorSimulator(),
        "AerSimulator(sv)":     AerSimulator(method="statevector"),
        "GPUStatevector":       FullGPUStatevector(),
        "DecisionDiagram":      DecisionDiagramBackend(),
        "AutoBackend":          AutoBackend(),
    }

    if verbose:
        print("\n📌 Bell Circuit (2 qubits):")
    for name, backend in backends.items():
        try:
            t = _time_backend(backend, bell, shots=shots)
            results[f"bell_{name}"] = t
            if verbose:
                print(f"   {name:30s} → {t*1000:.2f} ms")
        except Exception as e:
            if verbose:
                print(f"   {name:30s} → ERROR: {e}")

    # ── Random circuit ────────────────────────────────────────
    rand = _make_random_circuit(4, 8)
    if verbose:
        print("\n📌 Random Circuit (4 qubits, depth 8):")
    for name, backend in backends.items():
        try:
            t = _time_backend(backend, rand, shots=shots)
            results[f"random_{name}"] = t
            if verbose:
                print(f"   {name:30s} → {t*1000:.2f} ms")
        except Exception as e:
            if verbose:
                print(f"   {name:30s} → ERROR: {e}")

    # ── Accuracy check ────────────────────────────────────────
    if verbose:
        print("\n📌 Accuracy Check (Bell state |00> + |11>):")
    bell2 = _make_bell_circuit(2)
    sv_backend = StatevectorSimulator()
    r = sv_backend.run(bell2, shots=4096)
    counts = r.get_counts()
    p00 = counts.get("00", 0) / 4096
    p11 = counts.get("11", 0) / 4096
    results["accuracy_p00"] = p00
    results["accuracy_p11"] = p11
    if verbose:
        print(f"   P(|00>) = {p00:.4f}  (expected ~0.5)")
        print(f"   P(|11>) = {p11:.4f}  (expected ~0.5)")
        error = abs(p00 - 0.5) + abs(p11 - 0.5)
        print(f"   Total error = {error:.4f}  (< 0.05 is good)")

    if verbose:
        print("\n✅ Benchmarks complete!")
        print("=" * 55)

    return results
