
import numpy as np
import time

# ─────────────────────────────────────────────────────────────────
#  PKTron Documentation & Benchmark Suite
#
#  1. generate_sphinx_docs() — writes .rst files for Sphinx
#  2. run_full_benchmarks()  — times all backends vs reference
#  3. print_benchmark_table()— pretty-prints results
# ─────────────────────────────────────────────────────────────────

SPHINX_INDEX = """
PKTron Quantum Computing Framework
====================================

.. toctree::
   :maxdepth: 2
   :caption: Contents:

   installation
   quickstart
   backends
   noise
   error_mitigation
   tensor_network
   fault_tolerant
   gpu
   qiskit_compat
   api_reference
   benchmarks
   changelog

Introduction
------------
PKTron is a pure-Python quantum computing simulation framework designed
for research, education, and prototyping. It supports statevector,
density matrix, MPS tensor-network, stabilizer, and GPU-accelerated
backends with a unified interface.

Key Features
~~~~~~~~~~~~
- Statevector simulation up to ~28 qubits (CPU) / ~32 qubits (GPU)
- MPS / tensor-network simulation up to ~100+ qubits (low entanglement)
- Full noise modelling: depolarizing, amplitude damping, phase damping,
  bit-flip, phase-flip, thermal relaxation, readout errors
- Error mitigation: ZNE, PEC, Readout Error Mitigation
- Fault-tolerant primitives: repetition code, Steane code, surface code
- Differentiable circuits via PyTorch autograd (parameter-shift rule)
- Qiskit-compatible interface
"""

SPHINX_QUICKSTART = """
Quickstart
==========

Installation
------------

.. code-block:: bash

    pip install numpy matplotlib quimb scipy joblib

Running your first circuit
--------------------------

.. code-block:: python

    from pktron import QuantumCircuit, AerSimulator

    # Create a 2-qubit Bell state
    qc = QuantumCircuit(2)
    qc.h(0)
    qc.cnot(0, 1)

    # Simulate with 1024 shots
    sim = AerSimulator()
    result = sim.run(qc, shots=1024)
    print(result.get_counts())
    # Expected: {'00': ~512, '11': ~512}

Using noise
-----------

.. code-block:: python

    from pktron.noise_v2 import NoiseModelV2, DepolarizingError

    noise = NoiseModelV2()
    noise.add_gate_error(DepolarizingError(0.01), "H")
    noise.add_readout_error(p0given1=0.02, p1given0=0.03, qubit=0)

    from pktron.density_matrix_hp import HPDensityMatrixSimulator
    sim = HPDensityMatrixSimulator()
    result = sim.run(qc, shots=1024, noise_model=noise)
    print(result.purity())

Error Mitigation
----------------

.. code-block:: python

    from pktron.error_mitigation_v2 import ZNE
    zne = ZNE(scale_factors=(1, 2, 3))
    result = zne.run(executor_fn, circuit, observable_fn)
    print(result['zne_value'])

Surface Code (Fault Tolerance)
-------------------------------

.. code-block:: python

    from pktron.fault_tolerant import SurfaceCode, LogicalQubit
    lq = LogicalQubit(distance=3)
    lq.initialize(0)
    lq.apply_noise(p=0.005)
    print(lq.measure())   # should be 0
    print(lq.fidelity())  # close to 1.0
"""

SPHINX_API = """
API Reference
=============

QuantumCircuit
--------------
.. autoclass:: pktron.quantum.QuantumCircuit
   :members:
   :undoc-members:

AerSimulator
------------
.. autoclass:: pktron.simulator.AerSimulator
   :members:

HPDensityMatrixSimulator
------------------------
.. autoclass:: pktron.density_matrix_hp.HPDensityMatrixSimulator
   :members:

ProductionMPS
-------------
.. autoclass:: pktron.tensor_network_hp.ProductionMPS
   :members:

AutoSimulator
-------------
.. autoclass:: pktron.tensor_network_hp.AutoSimulator
   :members:

SurfaceCode
-----------
.. autoclass:: pktron.fault_tolerant.SurfaceCode
   :members:

LogicalQubit
------------
.. autoclass:: pktron.fault_tolerant.LogicalQubit
   :members:

GPUBatchedSampler
-----------------
.. autoclass:: pktron.gpu_primitives.GPUBatchedSampler
   :members:

ZNE
---
.. autoclass:: pktron.error_mitigation_v2.ZNE
   :members:

REM
---
.. autoclass:: pktron.error_mitigation_v2.REM
   :members:
"""


def generate_sphinx_docs(output_dir="pktron_docs"):
    """Write Sphinx .rst documentation files to output_dir."""
    import os
    os.makedirs(output_dir, exist_ok=True)

    files = {
        "index.rst": SPHINX_INDEX,
        "quickstart.rst": SPHINX_QUICKSTART,
        "api_reference.rst": SPHINX_API,
    }
    for fname, content in files.items():
        path = os.path.join(output_dir, fname)
        with open(path, "w") as f:
            f.write(content.strip())

    # Write conf.py
    with open(os.path.join(output_dir, "conf.py"), "w") as f:
        f.write("""
project = "PKTron"
author = "PKTron Contributors"
release = "2.2"
extensions = ["sphinx.ext.autodoc","sphinx.ext.napoleon","sphinx.ext.viewcode"]
html_theme = "sphinx_rtd_theme"
""".strip())

    print(f"  Sphinx docs written to: {output_dir}/")
    return output_dir


# ── Benchmark Suite ───────────────────────────────────────────────

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

def _make_random(n=6, depth=8):
    from pktron.quantum import QuantumCircuit
    qc = QuantumCircuit(n)
    for _ in range(depth):
        q = np.random.randint(n)
        np.random.choice([qc.h, qc.x, qc.z])(q)
        if n > 1:
            c = np.random.randint(n); t = (c+1)%n
            qc.cnot(c, t)
    return qc

def _time_it(fn, runs=3):
    times = []
    for _ in range(runs):
        t0 = time.perf_counter()
        fn()
        times.append(time.perf_counter() - t0)
    return float(np.mean(times)), float(np.std(times))

def run_full_benchmarks(shots=512, verbose=True):
    """
    Run comprehensive benchmarks across all PKTron backends.
    Returns a dict of results suitable for display or export.
    """
    from pktron.simulator import AerSimulator
    from pktron.density_matrix_hp import HPDensityMatrixSimulator
    from pktron.tensor_network_hp import ProductionMPS
    from pktron.gpu_primitives import GPUBatchedSampler
    from pktron.fault_tolerant import SurfaceCode, RepetitionCode

    results = []

    configs = [
        ("Bell(2q)",   _make_bell(2)),
        ("Bell(4q)",   _make_bell(4)),
        ("Random(6q)", _make_random(6, 8)),
        ("Random(10q)",_make_random(10, 6)),
    ]

    backends = [
        ("AerSimulator",         lambda qc: AerSimulator().run(qc, shots=shots)),
        ("HPDensityMatrix",      lambda qc: HPDensityMatrixSimulator().run(qc, shots=shots)),
        ("ProductionMPS",        lambda qc: ProductionMPS(max_bond_dim=32).run(qc, shots=shots)),
        ("GPUBatchedSampler",    lambda qc: GPUBatchedSampler(shots=shots).run(qc)),
    ]

    if verbose:
        print("\n" + "="*70)
        print(f"  PKTron v2.2 Benchmark Suite  |  shots={shots}")
        print("="*70)
        print(f"  {'Circuit':<14} {'Backend':<22} {'Mean(ms)':>10} {'Std(ms)':>9}")
        print("-"*70)

    for cname, qc in configs:
        for bname, bfn in backends:
            try:
                mean_t, std_t = _time_it(lambda: bfn(qc))
                row = {"circuit": cname, "backend": bname,
                       "mean_ms": round(mean_t*1000,2),
                       "std_ms":  round(std_t*1000,2)}
                results.append(row)
                if verbose:
                    print(f"  {cname:<14} {bname:<22} {mean_t*1000:>9.2f}ms "
                          f"{std_t*1000:>8.2f}ms")
            except Exception as ex:
                if verbose:
                    print(f"  {cname:<14} {bname:<22}   ERROR: {ex}")

    # Fault-tolerant benchmarks
    if verbose:
        print("\n  Fault-Tolerant Error Rates:")
        print(f"  {'Code':<22} {'p_physical':<14} {'p_logical':>10}")
        print("-"*50)

    for code, name in [(RepetitionCode(), "RepetitionCode(3)"),
                        (SurfaceCode(3),   "SurfaceCode(d=3)")]:
        for p in [0.001, 0.01, 0.05]:
            try:
                p_l = code.logical_error_rate(p, n_trials=1000)
                if verbose:
                    print(f"  {name:<22} {p:<14.3f} {p_l:>10.4f}")
                results.append({"circuit": name, "backend": "FaultTolerant",
                                 "p_physical": p, "p_logical": p_l})
            except Exception as ex:
                if verbose:
                    print(f"  {name}: ERROR {ex}")

    if verbose:
        print("="*70)
        print("  Benchmark complete.\n")

    return results
