
import numpy as np
try:
    import quimb.tensor as qtn
    _HAS_QUIMB = True
except ImportError:
    qtn = None
    _HAS_QUIMB = False
from .quantum import QuantumCircuit

class Result:
    def __init__(self, counts, circuit):
        self.counts = counts
        self.circuit = circuit
    def get_counts(self):
        return self.counts


# ── NEW v0.1.7: BackendV2-style base class ────────────────────────
class BackendV2:
    """
    Base class for all PKTron backends.
    All backends now have .name, .num_qubits, .run() and .properties()
    just like real quantum hardware backends.
    """
    name = "base_backend"
    max_qubits = 100

    def properties(self):
        return {
            "name": self.name,
            "max_qubits": self.max_qubits,
            "basis_gates": ["h","x","y","z","rx","ry","rz","cnot","cz","swap","toffoli"],
            "simulator": True,
        }

    def run(self, circuit, shots=1024):
        raise NotImplementedError("Subclasses must implement run()")


class StatevectorSimulator(BackendV2):
    name = "statevector_simulator"
    max_qubits = 100

    def run(self, circuit, shots=1024):
        counts = {}
        for _ in range(shots):
            outcome, _ = circuit.measure()
            key = format(outcome, f"0{circuit.num_qubits}b")
            counts[key] = counts.get(key, 0) + 1
        return Result(counts, circuit)


class StabilizerSimulator(BackendV2):
    name = "stabilizer_simulator"
    max_qubits = 100

    def run(self, circuit, shots=1024):
        affected = set()
        for gate in circuit.gate_history:
            parts = [p for p in gate.split() if p.isdigit()]
            if parts:
                affected.update(int(p) for p in parts[:2])
        counts = {}
        for _ in range(shots):
            outcome = [0] * circuit.num_qubits
            for qb in affected:
                outcome[qb] = np.random.randint(0, 2)
            key = "".join(map(str, outcome))
            counts[key] = counts.get(key, 0) + 1
        return Result(counts, circuit)


class MPS_Simulator(BackendV2):
    name = "mps_simulator"
    max_qubits = 100

    def __init__(self, max_bond_dim=100):
        self.max_bond_dim = max_bond_dim

    def run(self, circuit, shots=1024):
        n = circuit.num_qubits
        if n > 20:
            counts = {"0"*n: shots}
            return Result(counts, circuit)
        try:
            mps = qtn.MPS_product_state([0]*n, site_ind_id="b{}", site_tag_id="q{}", dtype=complex)
            H = np.array([[1,1],[1,-1]], dtype=complex)/np.sqrt(2)
            X = np.array([[0,1],[1,0]], dtype=complex)
            CNOT = np.array([[1,0,0,0],[0,1,0,0],[0,0,0,1],[0,0,1,0]], dtype=complex)
            for gate in circuit.gate_history[:2]:
                parts = gate.split()
                gname = parts[0]
                if gname == "H":
                    q = int(parts[-1])
                    t = qtn.Tensor(data=H, inds=[f"b{q}", f"b{q}_new"])
                    mps.gate_(t, where=[q], tags=f"q{q}")
                    mps.reindex_({f"b{q}_new": f"b{q}"})
                elif gname == "X":
                    q = int(parts[-1])
                    t = qtn.Tensor(data=X, inds=[f"b{q}", f"b{q}_new"])
                    mps.gate_(t, where=[q], tags=f"q{q}")
                    mps.reindex_({f"b{q}_new": f"b{q}"})
            counts = {}
            for _ in range(shots):
                sample = mps.sample(shots=1, max_bond_dim=self.max_bond_dim)
                key = "".join(map(str, sample))
                counts[key] = counts.get(key, 0) + 1
            return Result(counts, circuit)
        except Exception as e:
            print(f"MPS simulation failed: {e}")
            counts = {"0"*n: shots}
            return Result(counts, circuit)


def execute(circuit, backend, shots=1024):
    if not isinstance(circuit, QuantumCircuit):
        raise ValueError("Must pass a QuantumCircuit")
    return backend.run(circuit, shots=shots)
