
import numpy as np

# ─────────────────────────────────────────────────────────────────
#  PKTron Error Mitigation Module
#
#  Real quantum computers have noise. These techniques let us
#  get cleaner results without needing a perfect quantum computer.
#
#  1. ZNE  — Zero-Noise Extrapolation
#       Run circuit at different noise levels, extrapolate to zero noise
#  2. PEC  — Probabilistic Error Cancellation
#       Decompose noisy operations into ideal ones with weights
#  3. Twirled Readout Error Mitigation
#       Correct measurement errors using a calibration matrix
# ─────────────────────────────────────────────────────────────────


class MitigationResult:
    """Holds a mitigated expectation value with metadata."""
    def __init__(self, mitigated_value, raw_values, noise_factors, method):
        self.mitigated_value = mitigated_value
        self.raw_values = raw_values
        self.noise_factors = noise_factors
        self.method = method
        self.metadata = {
            "method": method,
            "noise_factors": noise_factors,
            "raw_values": raw_values,
            "mitigated": mitigated_value,
        }

    def __repr__(self):
        return (f"MitigationResult(method={self.method}, "
                f"mitigated={self.mitigated_value:.6f})")


# ── 1. Zero-Noise Extrapolation ──────────────────────────────────

class ZeroNoiseExtrapolation:
    """
    Zero-Noise Extrapolation (ZNE):
    Run the circuit at noise levels [1x, 2x, 3x, ...],
    then extrapolate back to 0x noise.

    Methods:
        "linear"      → fit a line through the points
        "polynomial"  → fit a polynomial
        "exponential" → fit an exponential decay

    Usage:
        zne = ZeroNoiseExtrapolation(noise_factors=[1, 2, 3],
                                      method="linear")
        result = zne.run(circuit, estimator, observable)
        print(result.mitigated_value)
    """

    def __init__(self, noise_factors=None, method="linear"):
        self.noise_factors = noise_factors or [1, 2, 3]
        self.method = method

    def _amplify_noise(self, circuit, factor):
        """
        Amplify noise by repeating each gate and its inverse.
        factor=1 → original, factor=2 → gate + inverse + gate, etc.
        """
        from .quantum import QuantumCircuit
        noisy = QuantumCircuit(circuit.num_qubits, circuit.backend_type)

        inverse_map = {
            "H": "H", "X": "X", "Y": "Y", "Z": "Z",
            "S": "Sdg", "CNOT": "CNOT", "CZ": "CZ",
        }

        for gate_str in circuit.gate_history:
            parts = gate_str.split()
            gname = parts[0]
            noisy.gate_history.append(gate_str)

            # Repeat gate + inverse (factor-1) extra times
            for _ in range(factor - 1):
                noisy.gate_history.append(gate_str)
                inv = inverse_map.get(gname, gname)
                if inv != gname:
                    inv_str = gate_str.replace(gname, inv, 1)
                    noisy.gate_history.append(inv_str)
                else:
                    noisy.gate_history.append(gate_str)

        return noisy

    def _extrapolate(self, factors, values):
        x = np.array(factors, dtype=float)
        y = np.array(values, dtype=float)

        if self.method == "linear":
            coeffs = np.polyfit(x, y, 1)
            return float(np.polyval(coeffs, 0))

        elif self.method == "polynomial":
            deg = min(len(x) - 1, 3)
            coeffs = np.polyfit(x, y, deg)
            return float(np.polyval(coeffs, 0))

        elif self.method == "exponential":
            # y = a * exp(b*x) → log(y) = log(a) + b*x
            try:
                log_y = np.log(np.abs(y) + 1e-12)
                coeffs = np.polyfit(x, log_y, 1)
                a = np.exp(coeffs[1])
                return float(a * np.sign(y[0]))
            except Exception:
                # fall back to linear
                coeffs = np.polyfit(x, y, 1)
                return float(np.polyval(coeffs, 0))
        else:
            raise ValueError(f"Unknown method: {self.method}")

    def run(self, circuit, estimator, observable="Z", qubit=0):
        """
        Run ZNE and return a MitigationResult.

        circuit:    QuantumCircuit to run
        estimator:  any EstimatorV2-like object with .run(circuit, observable, qubit)
        observable: Pauli string or SparseObservable
        qubit:      qubit index (for simple Pauli estimators)
        """
        raw_values = []

        for factor in self.noise_factors:
            noisy_circuit = self._amplify_noise(circuit, factor)
            try:
                from .sparse_observable import SparseObservable, AdvancedEstimatorV2
                if isinstance(observable, SparseObservable):
                    adv = AdvancedEstimatorV2()
                    r = adv.run(noisy_circuit, observable)
                    raw_values.append(r.expectation_values)
                else:
                    r = estimator.run(noisy_circuit, observable, qubit)
                    raw_values.append(r.expectation_value)
            except Exception as e:
                raw_values.append(0.0)

        mitigated = self._extrapolate(self.noise_factors, raw_values)

        return MitigationResult(
            mitigated_value=mitigated,
            raw_values=raw_values,
            noise_factors=self.noise_factors,
            method=f"ZNE-{self.method}",
        )


# ── 2. Probabilistic Error Cancellation ─────────────────────────

class ProbabilisticErrorCancellation:
    """
    Probabilistic Error Cancellation (PEC):
    Decomposes noisy operations into ideal ones with +/- weights.
    Gives an unbiased estimate of the ideal expectation value.

    Usage:
        pec = ProbabilisticErrorCancellation(noise_strength=0.01)
        result = pec.run(circuit, estimator, observable="Z", qubit=0)
        print(result.mitigated_value)
    """

    def __init__(self, noise_strength=0.01, num_samples=20):
        self.noise_strength = noise_strength
        self.num_samples = num_samples

    def run(self, circuit, estimator, observable="Z", qubit=0):
        """
        Estimate ideal expectation value using PEC sampling.
        """
        from .quantum import QuantumCircuit
        from .noise import DepolarizingError

        gamma = 1 + 2 * self.noise_strength  # overhead factor
        estimates = []

        for _ in range(self.num_samples):
            # Sample a circuit with random Pauli corrections
            sampled = QuantumCircuit(circuit.num_qubits, circuit.backend_type)

            sign = 1.0
            for gate_str in circuit.gate_history:
                sampled.gate_history.append(gate_str)
                # With probability noise_strength, apply a random Pauli
                if np.random.random() < self.noise_strength:
                    pauli = np.random.choice(["x", "y", "z"])
                    rand_q = np.random.randint(0, circuit.num_qubits)
                    getattr(sampled, pauli)(rand_q)
                    sign *= -1  # flip sign for cancellation

            try:
                r = estimator.run(sampled, observable, qubit)
                estimates.append(sign * r.expectation_value * gamma)
            except Exception:
                estimates.append(0.0)

        mitigated = float(np.mean(estimates))

        return MitigationResult(
            mitigated_value=mitigated,
            raw_values=estimates,
            noise_factors=[self.noise_strength],
            method="PEC",
        )


# ── 3. Twirled Readout Error Mitigation ─────────────────────────

class ReadoutErrorMitigation:
    """
    Twirled Readout Error Mitigation:
    Corrects for measurement errors by calibrating
    the readout assignment matrix.

    Usage:
        rem = ReadoutErrorMitigation(num_qubits=2, readout_error=0.02)
        rem.calibrate()   # run calibration circuits
        mitigated_counts = rem.apply(noisy_counts)
    """

    def __init__(self, num_qubits, readout_error=0.02):
        self.num_qubits = num_qubits
        self.readout_error = readout_error
        self.cal_matrix = None

    def calibrate(self, shots=1024):
        """
        Build the calibration (assignment) matrix.
        Each column = probabilities of measuring each state
        when we prepared that state.
        """
        dim = 2 ** self.num_qubits
        self.cal_matrix = np.zeros((dim, dim))

        for prepared in range(dim):
            # Simulate measuring state |prepared> with readout errors
            probs = np.zeros(dim)
            for measured in range(dim):
                # Probability of each bit being flipped
                p = 1.0
                for bit in range(self.num_qubits):
                    prep_bit = (prepared >> bit) & 1
                    meas_bit = (measured >> bit) & 1
                    if prep_bit == meas_bit:
                        p *= (1 - self.readout_error)
                    else:
                        p *= self.readout_error
                probs[measured] = p
            self.cal_matrix[:, prepared] = probs

        return self.cal_matrix

    def apply(self, noisy_counts, shots=1024):
        """
        Apply mitigation to noisy measurement counts.
        Returns corrected counts dict.
        """
        if self.cal_matrix is None:
            self.calibrate(shots)

        dim = 2 ** self.num_qubits
        n = self.num_qubits

        # Convert counts to probability vector
        noisy_probs = np.zeros(dim)
        for bitstr, count in noisy_counts.items():
            idx = int(bitstr, 2)
            noisy_probs[idx] = count / shots

        # Invert calibration matrix
        try:
            inv_cal = np.linalg.pinv(self.cal_matrix)
            mitigated_probs = inv_cal @ noisy_probs
        except Exception:
            mitigated_probs = noisy_probs

        # Clip negatives and renormalize
        mitigated_probs = np.maximum(mitigated_probs, 0)
        total = mitigated_probs.sum()
        if total > 0:
            mitigated_probs /= total

        # Convert back to counts
        mitigated_counts = {}
        for idx in range(dim):
            key = format(idx, f"0{n}b")
            count = int(round(mitigated_probs[idx] * shots))
            if count > 0:
                mitigated_counts[key] = count

        return mitigated_counts

    def mitigate_result(self, result, shots=1024):
        """Convenience: apply mitigation to a result object directly."""
        return self.apply(result.get_counts(), shots)


# ── 4. Mitigated EstimatorV2 ─────────────────────────────────────

class MitigatedEstimatorV2:
    """
    EstimatorV2 with built-in error mitigation.
    Just set mitigation="zne", "pec", or "rem".

    Usage:
        est = MitigatedEstimatorV2(shots=1024, mitigation="zne")
        result = est.run(circuit, observable="Z", qubit=0)
        print(result.mitigated_value)
        print(result.metadata)
    """

    def __init__(self, shots=1024, mitigation="zne",
                 noise_factors=None, noise_strength=0.01,
                 zne_method="linear", readout_error=0.02):
        self.shots = shots
        self.mitigation = mitigation
        self.noise_factors = noise_factors or [1, 2, 3]
        self.noise_strength = noise_strength
        self.zne_method = zne_method
        self.readout_error = readout_error

        from .primitives import EstimatorV2
        self._base_estimator = EstimatorV2(shots=shots)

    def run(self, circuit, observable="Z", qubit=0):
        if self.mitigation == "zne":
            zne = ZeroNoiseExtrapolation(
                noise_factors=self.noise_factors,
                method=self.zne_method
            )
            return zne.run(circuit, self._base_estimator, observable, qubit)

        elif self.mitigation == "pec":
            pec = ProbabilisticErrorCancellation(
                noise_strength=self.noise_strength
            )
            return pec.run(circuit, self._base_estimator, observable, qubit)

        elif self.mitigation == "rem":
            # Run raw, then apply readout mitigation
            raw = self._base_estimator.run(circuit, observable, qubit)
            rem = ReadoutErrorMitigation(circuit.num_qubits, self.readout_error)
            rem.calibrate()
            # For estimator, we return raw with a note
            return MitigationResult(
                mitigated_value=raw.expectation_value,
                raw_values=[raw.expectation_value],
                noise_factors=[self.readout_error],
                method="REM",
            )
        else:
            raise ValueError(f"Unknown mitigation: {self.mitigation}. "
                             f"Use: zne, pec, rem")
