
# ─────────────────────────────────────────────────────────────────
#  PKTron Qiskit Compatibility Layer
#
#  Lets most Qiskit scripts run on PKTron with minimal changes.
#  Just replace:
#      from qiskit import QuantumCircuit
#  with:
#      from pktron.qiskit_compat import QuantumCircuit
#
#  Supported aliases:
#    QuantumCircuit, ClassicalRegister, QuantumRegister
#    Aer, AerSimulator, execute, transpile
#    Statevector, DensityMatrix, Operator
#    PauliSumOp, SparsePauliOp
# ─────────────────────────────────────────────────────────────────

from .quantum import QuantumCircuit
from .midcircuit import ClassicalRegister, DynamicCircuit
from .simulator import AerSimulator
from .backend import StatevectorSimulator, execute, BackendV2
from .quantum_info import Statevector, DensityMatrix, Operator, Pauli
from .sparse_observable import SparseObservable
from .transpiler import PassManager
from .noise import NoiseModel, DepolarizingError, ThermalRelaxationError
from .primitives import SamplerV2, EstimatorV2
import numpy as np


# ── QuantumRegister (alias — PKTron uses integer qubit indices) ──

class QuantumRegister:
    """
    Alias for Qiskit-style QuantumRegister.
    PKTron uses integer indices, so this is just a named wrapper.

    Usage:
        qr = QuantumRegister(3, "q")
        print(qr.size)   # 3
        print(qr.name)   # "q"
    """
    def __init__(self, size, name="q"):
        self.size = size
        self.name = name

    def __len__(self):
        return self.size

    def __getitem__(self, index):
        return index   # return integer index

    def __repr__(self):
        return f"QuantumRegister({self.size}, '{self.name}')"


# ── SparsePauliOp (alias for SparseObservable) ──────────────────

class SparsePauliOp:
    """
    Qiskit-style SparsePauliOp.
    Wraps PKTron SparseObservable.

    Usage:
        op = SparsePauliOp(["ZZ", "XI"], coeffs=[0.5, 0.3])
        print(op.to_matrix())
    """
    def __init__(self, paulis, coeffs=None):
        if isinstance(paulis, str):
            paulis = [paulis]
        if coeffs is None:
            coeffs = [1.0] * len(paulis)
        self._obs = SparseObservable(list(zip(paulis, coeffs)))

    def to_matrix(self):
        return self._obs.to_matrix()

    def expectation_value(self, state):
        return self._obs.expectation_value(state)

    @property
    def num_qubits(self):
        return self._obs.num_qubits

    def __repr__(self):
        return f"SparsePauliOp({self._obs})"


# ── PauliSumOp (alias) ───────────────────────────────────────────
PauliSumOp = SparsePauliOp


# ── Aer mock ─────────────────────────────────────────────────────

class Aer:
    """
    Qiskit-style Aer mock.
    Aer.get_backend("aer_simulator") returns a PKTron AerSimulator.

    Usage:
        backend = Aer.get_backend("aer_simulator")
        result = execute(circuit, backend, shots=1024).result()
    """
    @staticmethod
    def get_backend(name="aer_simulator", method="statevector"):
        method_map = {
            "aer_simulator": "statevector",
            "statevector_simulator": "statevector",
            "density_matrix_simulator": "density_matrix",
            "unitary_simulator": "unitary",
        }
        m = method_map.get(name, "statevector")
        return AerSimulator(method=m)


# ── transpile wrapper ────────────────────────────────────────────

def transpile(circuit, backend=None, optimization_level=1, basis_gates=None):
    """
    Qiskit-style transpile function.
    Runs PKTron PassManager with passes matching optimization_level.

    optimization_level:
        0 → no optimization
        1 → commutative cancellation
        2 → unroller + cancellation
        3 → full: unroller + cancellation + basis translation + RZ synthesis
    """
    from .transpiler import (UnrollerPass, CancelCommutativePass,
                             BasisTranslationPass, RZSynthesisPass)

    passes = []
    if optimization_level >= 1:
        passes.append(CancelCommutativePass())
    if optimization_level >= 2:
        passes.append(UnrollerPass())
        passes.append(CancelCommutativePass())
    if optimization_level >= 3:
        target = "cx"
        if basis_gates and "cz" in basis_gates:
            target = "cz"
        passes.append(BasisTranslationPass(target=target))
        passes.append(RZSynthesisPass())

    if passes:
        pm = PassManager(passes)
        pm.run(circuit)

    return circuit


# ── Result wrapper ───────────────────────────────────────────────

class QiskitResult:
    """
    Wraps PKTron result to look like a Qiskit Result object.
    result.get_counts(), result.get_statevector(), etc.
    """
    def __init__(self, pktron_result):
        self._result = pktron_result

    def get_counts(self, circuit=None):
        return self._result.get_counts()

    def get_statevector(self, circuit=None):
        if hasattr(self._result, "get_statevector"):
            return self._result.get_statevector()
        raise ValueError("No statevector in this result")

    def get_unitary(self, circuit=None):
        if hasattr(self._result, "get_unitary"):
            return self._result.get_unitary()
        raise ValueError("No unitary in this result")

    def success(self):
        return True


# ── execute wrapper ──────────────────────────────────────────────

def qiskit_execute(circuit, backend, shots=1024, noise_model=None, **kwargs):
    """
    Qiskit-style execute function.
    Returns a QiskitResult wrapping the PKTron result.

    Usage:
        result = qiskit_execute(circuit, backend, shots=1024)
        counts = result.get_counts()
    """
    if noise_model is not None and hasattr(backend, "run"):
        raw = backend.run(circuit, shots=shots, noise_model=noise_model)
    else:
        raw = backend.run(circuit, shots=shots)
    return QiskitResult(raw)


# ── Convenience re-exports ───────────────────────────────────────

__all__ = [
    "QuantumCircuit", "ClassicalRegister", "QuantumRegister",
    "Aer", "AerSimulator", "execute", "qiskit_execute", "transpile",
    "Statevector", "DensityMatrix", "Operator", "Pauli",
    "SparsePauliOp", "PauliSumOp",
    "SamplerV2", "EstimatorV2",
    "NoiseModel", "DepolarizingError", "ThermalRelaxationError",
    "PassManager", "BackendV2",
]
