
import numpy as np

# ─────────────────────────────────────────────────────────────────
#  PKTron Auto Backend Optimizer
#
#  Automatically picks the best simulation backend based on:
#    - Number of qubits
#    - Number of gates
#    - Circuit depth (how many layers of gates)
#    - Whether GPU (CuPy) is available
#
#  Rules (simple version):
#    ≤ 15 qubits  → Statevector  (fast, exact, stores full state)
#    16-29 qubits → MPS          (memory efficient, tensor network)
#    ≥ 30 qubits  → GPU          (parallel, if CuPy installed)
#                               → MPS fallback if no GPU
#
#  You can also override the thresholds yourself.
# ─────────────────────────────────────────────────────────────────

# Default thresholds — change these if you want
_STATEVECTOR_MAX_QUBITS = 15
_MPS_MAX_QUBITS         = 29

def _has_gpu() -> bool:
    """Check if CuPy (GPU support) is installed and a GPU is available."""
    try:
        import cupy as cp
        cp.array([1])   # actually try to use it
        return True
    except Exception:
        return False

def _circuit_depth(gate_history: list) -> int:
    """
    Estimate circuit depth = number of gate layers.
    Simple approximation: count gates since we don't
    have a full DAG — good enough for backend selection.
    """
    return len(gate_history)

def _memory_estimate_mb(n_qubits: int) -> float:
    """
    Estimate statevector memory in MB.
    statevector has 2^n complex128 amplitudes = 16 bytes each.
    """
    return (2**n_qubits * 16) / (1024**2)


def select_backend(circuit, verbose: bool = True) -> str:
    """
    Automatically select the best PKTron backend for a circuit.

    Returns one of:
        "statevector"  — exact simulation, best for small circuits
        "mps"          — tensor network, best for medium circuits
        "gpu"          — GPU acceleration, best for large circuits

    Args:
        circuit: any PKTron QuantumCircuit
        verbose: if True, prints the decision and reasoning

    Example:
        from pktron import generate_circuit
        from pktron.auto_backend import select_backend, run_auto

        qc = generate_circuit("bell state")
        backend = select_backend(qc)
        # → "statevector"
    """
    n       = circuit.num_qubits
    history = getattr(circuit, 'gate_history', [])
    n_gates = len(history)
    depth   = _circuit_depth(history)
    mem_mb  = _memory_estimate_mb(n)
    gpu_ok  = _has_gpu()

    # ── Decision logic ────────────────────────────────────────────
    if n <= _STATEVECTOR_MAX_QUBITS:
        choice = "statevector"
        reason = (f"{n} qubits → statevector is fast and exact. "
                  f"Memory needed: {mem_mb:.2f} MB")

    elif n <= _MPS_MAX_QUBITS:
        choice = "mps"
        reason = (f"{n} qubits → too large for statevector "
                  f"({mem_mb:.0f} MB). MPS saves memory via "
                  f"tensor compression.")

    else:
        if gpu_ok:
            choice = "gpu"
            reason = (f"{n} qubits → very large circuit. "
                      f"GPU detected → using GPU acceleration.")
        else:
            choice = "mps"
            reason = (f"{n} qubits → very large circuit. "
                      f"No GPU found → falling back to MPS.")

    if verbose:
        _print_decision(n, n_gates, depth, mem_mb, gpu_ok, choice, reason)

    return choice


def _print_decision(n, n_gates, depth, mem_mb, gpu_ok, choice, reason):
    """Print a nicely formatted backend selection report."""
    print("=" * 58)
    print("  PKTron Auto Backend Optimizer")
    print("=" * 58)
    print(f"  Circuit info:")
    print(f"    Qubits      : {n}")
    print(f"    Gates       : {n_gates}")
    print(f"    Depth       : {depth}")
    print(f"    Memory (SV) : {mem_mb:.3f} MB")
    print(f"    GPU avail.  : {'✅ Yes' if gpu_ok else '❌ No'}")
    print()
    print(f"  ➤  Selected backend : {choice.upper()}")
    print(f"  ➤  Reason           : {reason}")
    print("=" * 58)


def run_auto(circuit, shots: int = 1024, verbose: bool = True):
    """
    Run a circuit using the automatically selected backend.

    Returns a result object or dict depending on which backend
    was chosen.

    Args:
        circuit: PKTron QuantumCircuit
        shots:   number of measurement shots
        verbose: print backend selection info

    Example:
        result = run_auto(my_circuit, shots=2048)
    """
    backend_name = select_backend(circuit, verbose=verbose)

    if backend_name == "statevector":
        return _run_statevector(circuit, shots)
    elif backend_name == "mps":
        return _run_mps(circuit, shots)
    elif backend_name == "gpu":
        return _run_gpu(circuit, shots)


def _run_statevector(circuit, shots):
    """Run using PKTron's built-in AerSimulator (statevector)."""
    try:
        from pktron import AerSimulator
        sim    = AerSimulator(method="statevector")
        result = sim.run(circuit, shots=shots)
        print(f"[AutoBackend] ✅ Statevector done. shots={shots}")
        return result
    except Exception as e:
        print(f"[AutoBackend] ⚠️  Statevector failed: {e}")
        # Return raw statevector as fallback
        sv    = circuit.get_state()
        probs = np.abs(sv)**2
        n     = circuit.num_qubits
        return {
            "backend":     "statevector_fallback",
            "statevector": sv,
            "probs":       probs,
            "labels":      [f'|{format(i,f"0{n}b")}>'
                            for i in range(2**n)],
        }


def _run_mps(circuit, shots):
    """Run using PKTron's MPS (tensor network) simulator."""
    try:
        from pktron.tensor_network_hp import MPSSimulatorHP
        sim = MPSSimulatorHP(circuit.num_qubits)
        sv  = sim.get_statevector()
        probs = np.abs(sv)**2
        n = circuit.num_qubits
        print(f"[AutoBackend] ✅ MPS done.")
        return {
            "backend": "mps",
            "statevector": sv,
            "probs":  probs,
            "labels": [f'|{format(i,f"0{n}b")}>'
                       for i in range(2**n)],
        }
    except Exception:
        # Graceful fallback to statevector
        print("[AutoBackend] ⚠️  MPS unavailable → falling back to statevector")
        return _run_statevector(circuit, shots)


def _run_gpu(circuit, shots):
    """Run using PKTron's GPU simulator."""
    try:
        from pktron.gpu_primitives import GPUStatevector
        sim = GPUStatevector(circuit.num_qubits)
        sv  = sim.get_statevector()
        probs = np.abs(sv)**2
        n = circuit.num_qubits
        print(f"[AutoBackend] ✅ GPU done.")
        return {
            "backend": "gpu",
            "statevector": sv,
            "probs":  probs,
            "labels": [f'|{format(i,f"0{n}b")}>'
                       for i in range(2**n)],
        }
    except Exception:
        # Graceful fallback to MPS then statevector
        print("[AutoBackend] ⚠️  GPU unavailable → falling back to MPS")
        return _run_mps(circuit, shots)


def backend_report(circuits: list, names: list = None) -> None:
    """
    Print a comparison table of backend choices for
    a list of circuits. Useful for planning which backend
    to use before running big experiments.

    Example:
        backend_report([qc_small, qc_medium, qc_large],
                       names=["Bell", "20q", "40q"])
    """
    if names is None:
        names = [f"Circuit {i+1}" for i in range(len(circuits))]

    print("=" * 70)
    print("  PKTron Backend Report")
    print("=" * 70)
    print(f"  {'Name':<14} {'Qubits':>6} {'Gates':>6} "
          f"{'Mem(MB)':>8} {'Backend':>12}")
    print(f"  {'-'*56}")

    for name, circuit in zip(names, circuits):
        n       = circuit.num_qubits
        n_gates = len(getattr(circuit, 'gate_history', []))
        mem_mb  = _memory_estimate_mb(n)
        backend = select_backend(circuit, verbose=False)
        print(f"  {name:<14} {n:>6} {n_gates:>6} "
              f"{mem_mb:>8.2f} {backend:>12}")

    print("=" * 70)
