"""
Real Zero-Noise Extrapolation (ZNE) error mitigation, wired to the same
engine and real noise channels (dense_evolution.NoiseModel) the rest of
the Composer uses. Noise is scaled by running the real channel at
noise_p, 2*noise_p, 3*noise_p and Richardson-extrapolating the measured
Pauli expectation back to zero noise via dense_evolution's own
zero_noise_extrapolation -- not a fabricated mitigated curve.
"""

from dataclasses import dataclass
from typing import Optional

import numpy as np

import dense_evolution as de

__all__ = ['MitigationResult', 'run_zne_mitigation', 'DensityMatrixZNEResult', 'run_density_matrix_zne']

_DEFAULT_NOISE_FACTORS = (1.0, 2.0, 3.0)
# polynomial_extrapolate's own docstring: with exactly degree+1 points the
# fit is the unique interpolating polynomial, IDENTICAL to
# richardson_extrapolate -- the extra 2 points here (4x, 5x) are what
# actually changes anything, trading a bit of interpolation bias for
# averaging down statistical noise across more measured scales (verified
# in that docstring against a real 5-seed sweep, not asserted here).
_POLYNOMIAL_NOISE_FACTORS = (1.0, 2.0, 3.0, 4.0, 5.0)
_POLYNOMIAL_DEGREE = 2


@dataclass
class MitigationResult:
    n_qubits: int
    pauli_string: str
    ideal_expectation: float
    noise_factors: list
    noisy_expectations: list
    zne_extrapolated: float
    extrapolation_method: str = "richardson"


def run_zne_mitigation(
    qasm_text: str,
    pauli_string: str,
    noise_model: str,
    noise_p: float,
    seed: Optional[int] = None,
    noise_factors=None,
    n_trials: int = 200,
    extrapolation_method: str = "richardson",
) -> MitigationResult:
    """Real ZNE: <P> is measured at the real ideal state and at the real
    channel applied at noise_p * each factor, then extrapolated to zero
    noise -- either Richardson (dense_evolution.zero_noise_extrapolation,
    the exact interpolating polynomial through 3 points, the default) or
    a degree-2 least-squares polynomial fit through 5 points
    (dense_evolution.polynomial_extrapolate) -- caller's choice, not
    silently picked: noise_factors defaults to the 3- or 5-point set that
    matches whichever method was requested, unless the caller overrides
    it explicitly.

    NoiseModel.apply_to_sv is a *stochastic single-shot* Kraus draw (one
    random outcome per call, not the channel's averaged/ensemble
    behavior) -- feeding a single draw straight into ZNE gives a
    discontinuous, meaningless curve (verified: bitflip/depolarizing
    jumped 1.0 -> 0.0 -> 0.0 across noise scales instead of decaying
    smoothly). <P> under a Kraus channel is Tr(rho P) = mean over the
    channel's trajectories, so each scale here averages n_trials
    independent stochastic draws -- the actual expectation value ZNE is
    defined against, not one random sample of it.

    pauli_string uses dense_evolution's own qubit-0-is-position-0
    convention (same as pauli_expectation), independent of Qiskit's
    little-endian display convention used elsewhere on this page --
    this function never touches a Qiskit-ordered array.
    """
    if extrapolation_method not in ("richardson", "polynomial"):
        raise ValueError(
            f"unknown extrapolation_method {extrapolation_method!r}, must be 'richardson' or 'polynomial'"
        )
    if noise_factors is None:
        noise_factors = (
            _POLYNOMIAL_NOISE_FACTORS if extrapolation_method == "polynomial" else _DEFAULT_NOISE_FACTORS
        )

    # Parsed with dense_evolution's own QASMParser, never Qiskit's --
    # QuantumCircuit.from_qasm_str itself segfaults on macOS (see
    # dashboard_core/engine.py's module docstring for the full story);
    # this function never needs a Qiskit circuit object at all, only the
    # ideal statevector, so there's no reason to build one here.
    parsed = de.QASMParser().parse(qasm_text)
    n_qubits = parsed.n_qubits
    if len(pauli_string) != n_qubits:
        raise ValueError(f"pauli_string length {len(pauli_string)} != n_qubits {n_qubits}")
    if noise_model not in de.NoiseModel.MODELS:
        raise ValueError(f"unknown noise model {noise_model!r}, must be one of {de.NoiseModel.MODELS}")

    # Same real anti-OOM guard as dashboard_core.engine.run_circuit_from_qasm
    # -- this kernel now runs on whatever machine a Composer visitor has,
    # not just a dev laptop, and this function was the one real gap that
    # never checked before allocating (engine.py's own functions always did).
    required_mb = (2 ** n_qubits) * 16 / 1e6
    de.chunk.SafeMemoryGuard().check_allocation(required_mb, context=f"{n_qubits}-qubit statevector")

    sim = de.DenseSVSimulator(n_qubits, use_float32=False)
    sim.run_circuit(parsed.to_tuples())
    sv_ideal = np.asarray(sim.sv)
    ideal_expectation = float(np.real(de.pauli_expectation(sv_ideal, pauli_string)))

    rng = np.random.default_rng(seed)
    noisy_expectations = []
    for factor in noise_factors:
        scaled_p = min(noise_p * factor, 1.0)
        trial_values = np.empty(n_trials)
        for i in range(n_trials):
            sv_noisy = de.NoiseModel.apply_to_sv(
                sv_ideal.copy(), n_qubits, noise_model, scaled_p, rng=rng,
            )
            trial_values[i] = np.real(de.pauli_expectation(sv_noisy, pauli_string))
        noisy_expectations.append(float(trial_values.mean()))

    if extrapolation_method == "polynomial":
        zne_value = float(de.polynomial_extrapolate(noisy_expectations, list(noise_factors), degree=_POLYNOMIAL_DEGREE))
    else:
        zne_value = float(de.zero_noise_extrapolation(noisy_expectations, list(noise_factors)))

    return MitigationResult(
        n_qubits=n_qubits,
        pauli_string=pauli_string,
        ideal_expectation=ideal_expectation,
        noise_factors=list(noise_factors),
        noisy_expectations=noisy_expectations,
        zne_extrapolated=zne_value,
        extrapolation_method=extrapolation_method,
    )


@dataclass
class DensityMatrixZNEResult:
    n_qubits: int
    noise_factors: list
    fidelity_raw: float          # Uhlmann fidelity: ideal vs raw noisy (base scale) rho
    fidelity_corrected: float    # Uhlmann fidelity: ideal vs ZNE-corrected rho


def run_density_matrix_zne(
    qasm_text: str,
    noise_model: str,
    noise_p: float,
    seed: Optional[int] = None,
    noise_factors=_DEFAULT_NOISE_FACTORS,
    n_trials: int = 200,
) -> DensityMatrixZNEResult:
    """Density-matrix ZNE (dense_evolution.zne_density_matrix): builds a
    real Monte-Carlo density-matrix estimate at each noise scale (the
    mean of n_trials |psi_k><psi_k| projectors from independent
    NoiseModel draws -- a real ensemble reconstruction, not a single
    trajectory), extrapolates to zero noise, and projects onto the
    nearest physical (PSD, trace-1) density matrix internally.

    Graded (never fed back into the extrapolation) against the true
    ideal density matrix via dense_evolution.uhlmann_fidelity, so the
    reported improvement is an honest measurement of whether the
    correction actually helped -- matches the pattern in
    zne_density_matrix's own docstring (experiments/matrix_healing_zne.py:
    raw ~0.865, corrected ~0.947 on a 2-qubit Bell state).
    """
    # Same QASMParser-only parsing as run_zne_mitigation above -- see its
    # comment for why a Qiskit circuit object is never built here.
    parsed = de.QASMParser().parse(qasm_text)
    n_qubits = parsed.n_qubits
    if noise_model not in de.NoiseModel.MODELS:
        raise ValueError(f"unknown noise model {noise_model!r}, must be one of {de.NoiseModel.MODELS}")

    # Density matrices are dim x dim (dim = 2**n_qubits), not just dim --
    # this holds rho_ideal, one rho per noise factor (rhos_at_scales, kept
    # alive simultaneously so zne_density_matrix can extrapolate across all
    # of them at once) and rho_corrected: (len(noise_factors) + 2) separate
    # dim*dim complex128 arrays, quadratically worse than a plain
    # statevector at the same qubit count. Same real anti-OOM guard as
    # dashboard_core.engine.run_circuit_from_qasm, sized for what this
    # function actually allocates -- this was the one real gap that never
    # checked before allocating.
    dim = 2 ** n_qubits
    required_mb = dim * dim * 16 / 1e6 * (len(noise_factors) + 2)
    de.chunk.SafeMemoryGuard().check_allocation(required_mb, context=f"{n_qubits}-qubit density matrix ZNE")

    sim = de.DenseSVSimulator(n_qubits, use_float32=False)
    sim.run_circuit(parsed.to_tuples())
    sv_ideal = np.asarray(sim.sv)
    rho_ideal = np.outer(sv_ideal, sv_ideal.conj())

    rng = np.random.default_rng(seed)
    rhos_at_scales = []
    for factor in noise_factors:
        scaled_p = min(noise_p * factor, 1.0)
        rho_acc = np.zeros((dim, dim), dtype=complex)
        for _ in range(n_trials):
            sv_noisy = de.NoiseModel.apply_to_sv(
                sv_ideal.copy(), n_qubits, noise_model, scaled_p, rng=rng,
            )
            rho_acc += np.outer(sv_noisy, sv_noisy.conj())
        rho_acc /= n_trials
        rhos_at_scales.append(rho_acc)

    rho_corrected = np.asarray(de.zne_density_matrix(rhos_at_scales, list(noise_factors)))

    fidelity_raw = float(de.uhlmann_fidelity(rho_ideal, rhos_at_scales[0]))
    fidelity_corrected = float(de.uhlmann_fidelity(rho_ideal, rho_corrected))

    return DensityMatrixZNEResult(
        n_qubits=n_qubits,
        noise_factors=list(noise_factors),
        fidelity_raw=fidelity_raw,
        fidelity_corrected=fidelity_corrected,
    )