
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

class BackendV2:
    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, noise_model=None):
        raise NotImplementedError("Subclasses must implement run()")


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

    def run(self, circuit, shots=1024, noise_model=None):
        counts = {}
        for _ in range(shots):
            active = noise_model.apply(circuit) if noise_model is not None else circuit
            outcome, _ = active.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, noise_model=None):
        active = noise_model.apply(circuit) if noise_model is not None else circuit
        affected = set()
        for gate in active.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, noise_model=None):
        active = noise_model.apply(circuit) if noise_model is not None else circuit
        n = active.num_qubits
        counts = {}
        for _ in range(shots):
            outcome, _ = active.measure()
            key = format(outcome, f"0{n}b")
            counts[key] = counts.get(key, 0) + 1
        return Result(counts, circuit)


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