
import numpy as np
import gc

# ──────────────────────────────────────────────
#  Built-in device profiles
# ──────────────────────────────────────────────
DEVICE_PROFILES = {
    "ibm_5q": {
        "num_qubits": 5,
        "connectivity": [(0,1),(1,2),(2,3),(3,4)],   # linear chain
        "gate_error": 0.001,        # single-qubit gate error rate
        "cx_error":  0.01,          # two-qubit (CX) gate error rate
        "t1": 100e-6,               # relaxation time (seconds)
        "t2":  80e-6,               # dephasing time (seconds)
        "gate_time": 50e-9,         # gate duration (seconds)
    },
    "noisy_20q": {
        "num_qubits": 20,
        "connectivity": [(i, i+1) for i in range(19)],  # linear chain
        "gate_error": 0.005,
        "cx_error":  0.02,
        "t1": 60e-6,
        "t2": 40e-6,
        "gate_time": 100e-9,
    },
}


class HardwareEmulator:
    """
    Simulates a real quantum device with:
      - qubit connectivity constraints
      - gate error rates
      - decoherence (T1 / T2)

    Usage:
        emu = HardwareEmulator("ibm_5q")
        result = emu.run(circuit, shots=1024)
        emu.summary()
    """

    def __init__(self, profile_name: str):
        if profile_name not in DEVICE_PROFILES:
            raise ValueError(
                f"Unknown profile '{profile_name}'. "
                f"Available: {list(DEVICE_PROFILES.keys())}"
            )
        self.profile_name = profile_name
        self.profile      = DEVICE_PROFILES[profile_name]
        self.num_qubits   = self.profile["num_qubits"]
        self.connectivity = set(self.profile["connectivity"])
        self.gate_error   = self.profile["gate_error"]
        self.cx_error     = self.profile["cx_error"]
        self.t1           = self.profile["t1"]
        self.t2           = self.profile["t2"]
        self.gate_time    = self.profile["gate_time"]

    # ── connectivity check ───────────────────────────────
    def is_connected(self, q1: int, q2: int) -> bool:
        """Return True if q1-q2 edge exists in device topology."""
        return (q1, q2) in self.connectivity or (q2, q1) in self.connectivity

    # ── decoherence probability ──────────────────────────
    def decoherence_prob(self, gate_depth: int) -> float:
        """
        Rough probability of a decoherence error after `gate_depth` gates.
        Uses T1 relaxation: p = 1 - exp(-depth * gate_time / T1)
        Capped at 0.5 so results stay physically meaningful.
        """
        t_total = gate_depth * self.gate_time
        p = 1.0 - np.exp(-t_total / self.t1)
        return min(p, 0.5)

    # ── run emulation ────────────────────────────────────
    def run(self, circuit, shots: int = 1024) -> dict:
        """
        Run a circuit on the emulated device.

        Returns a dict of bitstring -> count with realistic noise applied.
        Memory-safe: works shot-by-shot, never builds full statevector for large n.
        """
        n      = circuit.num_qubits
        depth  = len(getattr(circuit, "gate_history", []))

        # Warn if circuit is wider than the device
        if n > self.num_qubits:
            print(f"  ⚠️  Circuit has {n} qubits but device only has "
                  f"{self.num_qubits}. Results may be unreliable.")

        # Get ideal probabilities (use existing simulator if available)
        try:
            sv = circuit.get_statevector()
            probs = np.abs(sv) ** 2
            del sv
            gc.collect()
        except Exception:
            # Fallback: uniform distribution (circuit not simulated)
            dim   = 2 ** min(n, self.num_qubits)
            probs = np.ones(dim) / dim

        dim = len(probs)

        # ── apply gate errors ────────────────────────────
        noise_per_gate = self.gate_error * depth
        noise_per_gate = min(noise_per_gate, 0.3)   # cap at 30 %

        # Depolarise: mix ideal probs with uniform
        uniform = np.ones(dim) / dim
        noisy   = (1 - noise_per_gate) * probs + noise_per_gate * uniform

        # ── apply decoherence ────────────────────────────
        d_prob = self.decoherence_prob(depth)
        noisy  = (1 - d_prob) * noisy + d_prob * uniform

        # Renormalise (floating-point safety)
        noisy = np.abs(noisy)
        noisy /= noisy.sum()

        # ── sample shots ─────────────────────────────────
        indices  = np.random.choice(dim, size=shots, p=noisy)
        fmt      = f"{{:0{min(n, self.num_qubits)}b}}"
        counts   = {}
        for idx in indices:
            key = fmt.format(idx)
            counts[key] = counts.get(key, 0) + 1

        # Clean up large arrays immediately
        del probs, noisy, uniform, indices
        gc.collect()

        return counts

    # ── pretty summary ───────────────────────────────────
    def summary(self):
        p = self.profile
        print("=" * 52)
        print(f"  Hardware Emulator — {self.profile_name}")
        print("=" * 52)
        print(f"  Qubits         : {p['num_qubits']}")
        print(f"  Connectivity   : {len(self.connectivity)} edges")
        print(f"  Gate error     : {p['gate_error']*100:.2f}%")
        print(f"  CX error       : {p['cx_error']*100:.2f}%")
        print(f"  T1 (relaxation): {p['t1']*1e6:.1f} µs")
        print(f"  T2 (dephasing) : {p['t2']*1e6:.1f} µs")
        print(f"  Gate time      : {p['gate_time']*1e9:.1f} ns")
        print("=" * 52)

    def __repr__(self):
        return f"HardwareEmulator(profile='{self.profile_name}', qubits={self.num_qubits})"


# ── convenience factory ──────────────────────────────────
def get_emulator(profile_name: str) -> HardwareEmulator:
    """Quick factory: emu = get_emulator('ibm_5q')"""
    return HardwareEmulator(profile_name)
