"""Circuit Analyzer -- extract features from a Qiskit QuantumCircuit.

This module is the foundation of the EMRG pipeline. It inspects a
QuantumCircuit and returns a :class:`CircuitFeatures` dataclass that
downstream modules (heuristics, codegen) consume.

Typical usage::

    from qiskit import QuantumCircuit
    from emrg.analyzer import analyze_circuit

    qc = QuantumCircuit(2, 2)
    qc.h(0)
    qc.cx(0, 1)
    qc.measure_all()

    features = analyze_circuit(qc)
    print(features)
"""

from __future__ import annotations

import math
import warnings
from dataclasses import dataclass, field
from types import MappingProxyType
from typing import TYPE_CHECKING, Any

if TYPE_CHECKING:
    from collections.abc import Mapping

    from qiskit import QuantumCircuit

__all__ = [
    "CircuitFeatures",
    "analyze_circuit",
    "MULTI_QUBIT_GATE_NAMES",
    "CLIFFORD_GATE_NAMES",
    "DEFAULT_MULTI_QUBIT_ERROR_RATE",
    "DEFAULT_SINGLE_QUBIT_ERROR_RATE",
    "NOISE_THRESHOLD_LOW",
    "NOISE_THRESHOLD_MODERATE",
    "PEC_OVERHEAD_NOISE_MULTIPLIER",
    "CDR_NON_CLIFFORD_FRACTION_THRESHOLD",
    "CDR_MIN_DEPTH",
    "CDR_MAX_DEPTH",
]

# ---------------------------------------------------------------------------
# Constants
# ---------------------------------------------------------------------------

#: Gate names that Qiskit uses for two-qubit and multi-qubit operations.
#: Includes 3-qubit gates (ccx, ccz, cswap) that also benefit from
#: multi-qubit error rate estimates.
MULTI_QUBIT_GATE_NAMES: frozenset[str] = frozenset(
    {
        "cx",
        "cz",
        "cy",
        "ch",
        "crx",
        "cry",
        "crz",
        "cp",
        "swap",
        "iswap",
        "ecr",
        "rzx",
        "rzz",
        "rxx",
        "ryy",
        "csx",
        "cu",
        "cu1",
        "cu3",
        "dcx",
        "ccx",  # 3-qubit
        "ccz",  # 3-qubit
        "cswap",  # 3-qubit
    }
)

#: Default proxy for average multi-qubit gate error rate on NISQ hardware.
#: IBM Eagle r3 ~ 0.005-0.01; we use a conservative midpoint.
DEFAULT_MULTI_QUBIT_ERROR_RATE: float = 0.01

#: Default proxy for average single-qubit gate error rate.
DEFAULT_SINGLE_QUBIT_ERROR_RATE: float = 0.001

#: Noise factor below this is classified as "low".
NOISE_THRESHOLD_LOW: float = 0.05

#: Noise factor below this (but >= NOISE_THRESHOLD_LOW) is "moderate";
#: at or above this threshold is "high".
NOISE_THRESHOLD_MODERATE: float = 0.20

#: Multiplier used in the PEC overhead formula:
#: ``exp(PEC_OVERHEAD_NOISE_MULTIPLIER * noise_factor * multi_qubit_gate_count)``.
PEC_OVERHEAD_NOISE_MULTIPLIER: float = 2.0

#: Gate names that are always Clifford (regardless of parameters).
#: The Clifford group is generated by {H, S, CNOT}.
CLIFFORD_GATE_NAMES: frozenset[str] = frozenset(
    {
        "h", "x", "y", "z",
        "s", "sdg",
        "sx", "sxdg",
        "cx", "cz", "cy",
        "swap", "iswap",
        "id", "i",
        "ecr",
        "dcx",
    }
)

#: Rotation gate names that are Clifford only at specific angles
#: (multiples of pi/2).
_ROTATION_GATE_NAMES: frozenset[str] = frozenset(
    {
        "r",
        "rz", "ry", "rx",
        "p", "u1",
        "rxx", "ryy", "rzz", "rzx",
        "crx", "cry", "crz", "cp",
        "u2", "u3", "u",
        "cu", "cu1", "cu3",
    }
)

#: Minimum non-Clifford gate fraction for CDR to be considered.
CDR_NON_CLIFFORD_FRACTION_THRESHOLD: float = 0.2

#: Minimum circuit depth for CDR consideration.
CDR_MIN_DEPTH: int = 10

#: Maximum circuit depth for CDR consideration.
CDR_MAX_DEPTH: int = 40


# ---------------------------------------------------------------------------
# Data model
# ---------------------------------------------------------------------------


def _freeze_dict(value: dict[str, Any]) -> MappingProxyType[str, Any]:
    """Wrap a dict in a read-only ``MappingProxyType``."""
    return MappingProxyType(value)


@dataclass(frozen=True)
class CircuitFeatures:
    """Immutable snapshot of circuit properties relevant to error mitigation.

    All fields are derived from Qiskit introspection -- no simulation needed.

    Attributes:
        num_qubits: Number of quantum bits in the circuit.
        depth: Circuit depth (longest path of dependent operations).
        gate_counts: Read-only mapping of gate name to count
            (e.g. ``{'cx': 4, 'h': 3}``).
        total_gate_count: Total number of gate operations.
        multi_qubit_gate_count: Number of multi-qubit gates (2-qubit and
            3-qubit, e.g. CX, CCX, CSWAP).
        single_qubit_gate_count: Number of single-qubit gates.
        num_parameters: Number of unbound parameters (0 for non-variational).
        has_measurements: Whether the circuit contains measurement operations.
        estimated_noise_factor: Simple heuristic noise proxy
            (``multi_qubit_gate_count * multi_qubit_error_rate
              + single_qubit_gate_count * single_qubit_error_rate``).
        noise_category: Human-readable category: ``"low"``, ``"moderate"``,
            or ``"high"``.
        noise_model_available: Whether a noise model is available for PEC.
            Defaults to ``False``; set to ``True`` when a depolarizing (or
            other) noise model can be supplied so the heuristic engine may
            consider Probabilistic Error Cancellation.
        pec_overhead_estimate: Approximate PEC sampling overhead computed as
            ``exp(2 * noise_factor * multi_qubit_gate_count)``.  Values much
            larger than ~1000 make PEC impractical.
        layer_heterogeneity: Ratio of the maximum to minimum multi-qubit gate
            count across circuit layers (with +1 denominator smoothing).
            Higher values indicate more varied layer structure, which favours
            layerwise noise scaling over uniform folding.  Returns 0.0 when
            the circuit has fewer than two layers containing multi-qubit gates.
        non_clifford_count: Number of non-Clifford gates in the circuit.
            Non-Clifford gates include T, Tdg, and any rotation gate whose
            angle is not a multiple of pi/2.
        non_clifford_fraction: Fraction of gates that are non-Clifford
            (``non_clifford_count / total_gate_count``).  0.0 means all
            Clifford, 1.0 means all non-Clifford.
    """

    num_qubits: int
    depth: int
    gate_counts: Mapping[str, int] = field(default_factory=dict)
    total_gate_count: int = 0
    multi_qubit_gate_count: int = 0
    single_qubit_gate_count: int = 0
    num_parameters: int = 0
    has_measurements: bool = False
    estimated_noise_factor: float = 0.0
    noise_category: str = "low"
    noise_model_available: bool = False
    pec_overhead_estimate: float = 1.0
    layer_heterogeneity: float = 0.0
    non_clifford_count: int = 0
    non_clifford_fraction: float = 0.0


# ---------------------------------------------------------------------------
# Validation helpers
# ---------------------------------------------------------------------------


def _validate_circuit(qc: QuantumCircuit) -> list[str]:
    """Validate the circuit and return a list of warning messages (may be empty).

    Checks:
    - Circuit must be a ``QuantumCircuit`` instance.
    - Circuit must contain at least one gate operation.
    - Circuit should have measurements for meaningful mitigation.

    Raises:
        TypeError: If *qc* is not a ``QuantumCircuit``.
        ValueError: If the circuit has zero gate operations.
    """
    from qiskit import QuantumCircuit as _QuantumCircuit  # lazy runtime import

    if not isinstance(qc, _QuantumCircuit):
        raise TypeError(f"Expected a qiskit.QuantumCircuit, got {type(qc).__name__}.")

    ops = qc.count_ops()
    # Filter out measurement/barrier/delay -- only count real gates
    gate_ops = {
        name: count
        for name, count in ops.items()
        if name not in ("measure", "barrier", "delay", "reset")
    }

    if not gate_ops:
        raise ValueError(
            "Circuit has no gate operations. "
            "Provide a circuit with at least one quantum gate."
        )

    warns: list[str] = []
    if "measure" not in ops:
        warns.append(
            "Circuit has no measurements. Error mitigation requires "
            "measurement results -- consider adding qc.measure_all()."
        )
    if qc.num_parameters > 0:
        warns.append(
            f"Circuit has {qc.num_parameters} unbound parameter(s). "
            "Parameters must be bound before execution "
            "(e.g. qc = qc.assign_parameters({...})). "
            "EMRG will analyze the circuit structure, but the generated "
            "code cannot run until all parameters are resolved."
        )
    return warns


# ---------------------------------------------------------------------------
# Noise estimation
# ---------------------------------------------------------------------------


def _classify_noise(noise_factor: float) -> str:
    """Map a numeric noise factor to a human-readable category.

    Thresholds are deliberately conservative and tunable.
    """
    if noise_factor < NOISE_THRESHOLD_LOW:
        return "low"
    elif noise_factor < NOISE_THRESHOLD_MODERATE:
        return "moderate"
    else:
        return "high"


def _estimate_noise(
    multi_qubit_count: int,
    single_qubit_count: int,
    *,
    multi_qubit_error_rate: float = DEFAULT_MULTI_QUBIT_ERROR_RATE,
    single_qubit_error_rate: float = DEFAULT_SINGLE_QUBIT_ERROR_RATE,
) -> float:
    """Return a simple additive noise proxy.

    This is *not* a rigorous noise model -- it is a fast heuristic used to
    guide mitigation strategy selection.
    """
    return (
        multi_qubit_count * multi_qubit_error_rate
        + single_qubit_count * single_qubit_error_rate
    )


def _estimate_pec_overhead(
    noise_factor: float,
    multi_qubit_gate_count: int,
) -> float:
    """Return the approximate PEC sampling overhead.

    The formula ``exp(2 * noise_factor * multi_qubit_gate_count)`` captures
    the exponential cost that makes PEC impractical for large or very noisy
    circuits.  The exponent is capped at 700 to avoid ``OverflowError`` on
    circuits with hundreds of multi-qubit gates.
    """
    exponent = PEC_OVERHEAD_NOISE_MULTIPLIER * noise_factor * multi_qubit_gate_count
    # math.exp overflows around ~709; cap to prevent OverflowError.
    if exponent > 700:
        return math.inf
    return math.exp(exponent)


# ---------------------------------------------------------------------------
# Layer heterogeneity
# ---------------------------------------------------------------------------


def _compute_layer_heterogeneity(qc: QuantumCircuit) -> float:
    """Compute how unevenly multi-qubit gates are distributed across layers.

    Returns ``max(counts) / (min(counts) + 1)`` over circuit layers that
    contain at least one multi-qubit gate.  Returns 0.0 when fewer than
    two such layers exist (no variation is possible).
    """
    from qiskit.converters import circuit_to_dag

    dag = circuit_to_dag(qc)
    mq_counts: list[int] = []
    for layer_dict in dag.layers():
        layer_dag = layer_dict["graph"]
        mq = sum(
            1
            for node in layer_dag.op_nodes()
            if node.op.name in MULTI_QUBIT_GATE_NAMES
        )
        if mq > 0:
            mq_counts.append(mq)

    if len(mq_counts) < 2:
        return 0.0
    return max(mq_counts) / (min(mq_counts) + 1)


# ---------------------------------------------------------------------------
# Non-Clifford gate counting
# ---------------------------------------------------------------------------


def _is_clifford_angle(angle: float) -> bool:
    """Return ``True`` if *angle* is a multiple of pi/2 (within tolerance).

    Clifford rotations are those at multiples of pi/2:
    0, pi/2, pi, 3*pi/2, 2*pi, etc.
    """
    # Normalise to units of pi/2; check if close to an integer.
    units = angle / (math.pi / 2)
    return abs(units - round(units)) < 1e-8


def _count_non_clifford_gates(qc: QuantumCircuit) -> int:
    """Count gates that are outside the Clifford group.

    Classification:
    - Gates in ``CLIFFORD_GATE_NAMES`` are always Clifford.
    - T, Tdg are always non-Clifford.
    - Rotation gates (rz, ry, rx, p, u1, u2, u3, u, crx, cry, crz, cp,
      rxx, ryy, rzz, rzx, cu, cu1, cu3) are Clifford only when all their
      parameters are multiples of pi/2.
    - Any other gate not in the known sets is counted as non-Clifford.
    """
    from qiskit.converters import circuit_to_dag

    dag = circuit_to_dag(qc)
    non_gate_names = {"measure", "barrier", "delay", "reset"}
    count = 0

    for node in dag.op_nodes():
        name = node.op.name
        if name in non_gate_names:
            continue

        # Always-Clifford gates.
        if name in CLIFFORD_GATE_NAMES:
            continue

        # Always non-Clifford gates.
        if name in ("t", "tdg"):
            count += 1
            continue

        # Rotation gates: Clifford iff all angles are multiples of pi/2.
        if name in _ROTATION_GATE_NAMES:
            params = node.op.params
            if params and all(
                isinstance(p, (int, float)) and _is_clifford_angle(float(p))
                for p in params
            ):
                continue
            count += 1
            continue

        # Unknown gate -- conservatively count as non-Clifford.
        count += 1

    return count


# ---------------------------------------------------------------------------
# Public API
# ---------------------------------------------------------------------------


def analyze_circuit(
    qc: QuantumCircuit,
    *,
    multi_qubit_error_rate: float = DEFAULT_MULTI_QUBIT_ERROR_RATE,
    single_qubit_error_rate: float = DEFAULT_SINGLE_QUBIT_ERROR_RATE,
    noise_model_available: bool = False,
) -> CircuitFeatures:
    """Analyze a Qiskit ``QuantumCircuit`` and extract mitigation-relevant features.

    Parameters
    ----------
    qc:
        The circuit to analyze.
    multi_qubit_error_rate:
        Proxy error rate for multi-qubit gates (default 0.01).
    single_qubit_error_rate:
        Proxy error rate for single-qubit gates (default 0.001).
    noise_model_available:
        Whether a noise model is available for techniques like PEC
        (default ``False``).  Set to ``True`` when a depolarizing or
        tomographic noise model can be supplied.

    Returns
    -------
    CircuitFeatures
        Immutable dataclass with extracted metrics.

    Raises
    ------
    TypeError
        If *qc* is not a ``QuantumCircuit``.
    ValueError
        If the circuit has zero gate operations.

    Examples
    --------
    >>> from qiskit import QuantumCircuit
    >>> qc = QuantumCircuit(2, 2); _ = qc.h(0); _ = qc.cx(0, 1)
    >>> qc.measure([0, 1], [0, 1])
    >>> features = analyze_circuit(qc)
    >>> features.num_qubits
    2
    >>> features.multi_qubit_gate_count
    1
    """
    # --- validate -----------------------------------------------------------
    warn_msgs = _validate_circuit(qc)
    for msg in warn_msgs:
        warnings.warn(msg, UserWarning, stacklevel=2)

    # --- extract raw ops ----------------------------------------------------
    ops = dict(qc.count_ops())  # OrderedDict -> dict

    # Separate gate ops from non-gate ops (measure, barrier, etc.)
    non_gate_names = {"measure", "barrier", "delay", "reset"}
    gate_counts: dict[str, int] = {
        name: count for name, count in ops.items() if name not in non_gate_names
    }

    total_gate_count = sum(gate_counts.values())
    multi_qubit_gate_count = sum(
        count for name, count in gate_counts.items() if name in MULTI_QUBIT_GATE_NAMES
    )
    single_qubit_gate_count = total_gate_count - multi_qubit_gate_count

    # --- noise estimate -----------------------------------------------------
    noise_factor = _estimate_noise(
        multi_qubit_gate_count,
        single_qubit_gate_count,
        multi_qubit_error_rate=multi_qubit_error_rate,
        single_qubit_error_rate=single_qubit_error_rate,
    )

    pec_overhead = _estimate_pec_overhead(noise_factor, multi_qubit_gate_count)
    layer_het = _compute_layer_heterogeneity(qc)
    non_clifford = _count_non_clifford_gates(qc)
    nc_fraction = non_clifford / total_gate_count if total_gate_count > 0 else 0.0

    return CircuitFeatures(
        num_qubits=qc.num_qubits,
        depth=qc.depth(),
        gate_counts=_freeze_dict(gate_counts),
        total_gate_count=total_gate_count,
        multi_qubit_gate_count=multi_qubit_gate_count,
        single_qubit_gate_count=single_qubit_gate_count,
        num_parameters=qc.num_parameters,
        has_measurements="measure" in ops,
        estimated_noise_factor=round(noise_factor, 6),
        noise_category=_classify_noise(noise_factor),
        noise_model_available=noise_model_available,
        pec_overhead_estimate=round(pec_overhead, 6),
        layer_heterogeneity=round(layer_het, 4),
        non_clifford_count=non_clifford,
        non_clifford_fraction=round(nc_fraction, 6),
    )
