
import numpy as np
from itertools import product as iproduct

class ZNE:
    """Zero Noise Extrapolation — run at multiple noise levels, fit to zero."""
    def __init__(self, scale_factors=(1,2,3)):
        self.scale_factors = scale_factors

    def run(self, executor_fn, circuit, observable_fn):
        values = []
        for s in self.scale_factors:
            noisy_circuit = self._fold_gates(circuit, s)
            result = executor_fn(noisy_circuit)
            values.append(observable_fn(result))
        coeffs = np.polyfit(self.scale_factors, values, deg=len(self.scale_factors)-1)
        zne_value = float(np.polyval(coeffs, 0))
        return {"zne_value": zne_value, "raw_values": values, "scale_factors": list(self.scale_factors)}

    def _fold_gates(self, circuit, scale):
        from copy import deepcopy
        new = deepcopy(circuit)
        original = list(circuit._ops)
        extra = []
        for _ in range(scale - 1):
            extra += [(g, q, p) for g, q, p in reversed(original)]
            extra += original
        new._ops = original + extra
        return new

class PEC:
    """Probabilistic Error Cancellation (simplified quasi-probability sampling)."""
    def __init__(self, noise_strength=0.01):
        self.noise_strength = noise_strength

    def run(self, executor_fn, circuit, observable_fn, num_samples=50):
        estimates = []
        for _ in range(num_samples):
            sign = np.random.choice([-1, 1], p=[self.noise_strength/2, 1-self.noise_strength/2])
            result = executor_fn(circuit)
            estimates.append(sign * observable_fn(result))
        return {"pec_value": float(np.mean(estimates)),
                "std": float(np.std(estimates)),
                "num_samples": num_samples}

class REM:
    """Readout Error Mitigation via calibration matrix inversion."""
    def __init__(self, num_qubits):
        self.num_qubits = num_qubits
        self.cal_matrix = np.eye(2**num_qubits)

    def calibrate(self, executor_fn, circuit_factory):
        n = self.num_qubits
        dim = 2**n
        for i in range(dim):
            bitstring = format(i, f"0{n}b")
            qc = circuit_factory(bitstring)
            result = executor_fn(qc)
            counts = result if isinstance(result, dict) else result.get_counts()
            total = sum(counts.values())
            col = np.zeros(dim)
            for k, v in counts.items():
                col[int(k, 2)] = v / total
            self.cal_matrix[:, i] = col

    def apply(self, counts):
        n = self.num_qubits
        dim = 2**n
        total = sum(counts.values())
        measured = np.zeros(dim)
        for k, v in counts.items():
            measured[int(k, 2)] = v / total
        try:
            corrected = np.linalg.solve(self.cal_matrix, measured)
            corrected = np.clip(corrected, 0, None)
            corrected /= corrected.sum()
        except np.linalg.LinAlgError:
            corrected = measured
        return {format(i, f"0{n}b"): float(v) for i, v in enumerate(corrected) if v > 1e-6}
