# This code is part of Qiskit.
#
# (C) Copyright IBM 2024.
#
# This code is licensed under the Apache License, Version 2.0. You may
# obtain a copy of this license in the LICENSE.txt file in the root directory
# of this source tree or at http://www.apache.org/licenses/LICENSE-2.0.
#
# Any modifications or derivative works of this code must retain this
# copyright notice, and modified files need to carry a notice indicating
# that they have been altered from the originals.

"""A class to help understand the expected performance of estimator jobs."""

from __future__ import annotations
from typing import Optional, Sequence

from qiskit.exceptions import QiskitError
from qiskit.transpiler.passmanager import PassManager
from qiskit.primitives.containers import EstimatorPubLike
from qiskit.primitives.containers.estimator_pub import EstimatorPub
from qiskit.providers import BackendV2 as Backend

from qiskit_ibm_runtime.debug_tools.neat_results import NeatPubResult, NeatResult
from qiskit_ibm_runtime.transpiler.passes.cliffordization import ConvertISAToClifford


try:
    from qiskit_aer.noise import NoiseModel
    from qiskit_aer.primitives.estimator_v2 import EstimatorV2 as AerEstimator

    HAS_QISKIT_AER = True
except ImportError:
    HAS_QISKIT_AER = False


class Neat:
    r"""A class to help understand the expected performance of estimator jobs.

    The "Noisy Estimator Analyzer Tool" (or "NEAT") is a convenience tool that users of the
    :class:`~.Estimator` primitive can employ to analyze and predict the performance of
    their queries. Its simulate method uses ``qiskit-aer`` to simulate the estimation task
    classically efficiently, either in ideal conditions or in the presence of noise. The
    simulations' results can be compared with other simulation results or with primitive results
    results to draw custom figures of merit.

    .. code::python

        # Initialize a Neat object
        analyzer = Neat(backend)

        # Map arbitrary PUBs to Clifford PUBs
        cliff_pubs = analyzer.to_clifford(pubs)

        # Calculate the expectation values in the absence of noise
        r_ideal = analyzer.ideal_sim(cliff_pubs)

        # Calculate the expectation values in the presence of noise
        r_noisy = analyzer.noisy_sim(cliff_pubs)

        # Calculate the expectation values for a different noise model
        analyzer.noise_model = another_noise_model
        another_r_noisy = analyzer.noisy_sim(cliff_pubs)

        # Run the Clifford PUBs on a QPU
        r_qpu = estimator.run(cliff_pubs)

        # Calculate useful figures of merit using mathematical operators, for example the relative
        # difference between experimental and noisy results, ...
        rel_diff = abs(r_noisy[0] - r_qpu[0]) / r_noisy[0]

        # ... the signal-to-noise ratio between experimental and ideal results, ...
        ratio = r_qpu[0] / r_ideal[0]

        # ... or the absolute difference between results obtained with different noise models
        abs_diff = abs(r_noisy[0] - another_r_noisy[0])

    Args:
        backend: A backend.
        noise_model: A noise model for the operations of the given backend. If ``None``, it
            defaults to the noise model generated by :meth:`NoiseModel.from_backend`.
    """

    def __init__(self, backend: Backend, noise_model: Optional[NoiseModel] = None) -> None:
        if not HAS_QISKIT_AER:
            raise ValueError(
                "Cannot initialize object of type 'Neat' since 'qiskit-aer' is not installed. "
                "Install 'qiskit-aer' and try again."
            )

        self._backend = backend
        self.noise_model = (
            noise_model
            if noise_model is not None
            else NoiseModel.from_backend(backend, thermal_relaxation=False)
        )

    @property
    def noise_model(self) -> NoiseModel:
        r"""
        The noise model used by this analyzer tool for the noisy simulations.
        """
        return self._noise_model

    @noise_model.setter
    def noise_model(self, value: NoiseModel) -> NoiseModel:
        """Sets a new noise model.

        Args:
            value: A new noise model.
        """
        self._noise_model = value

    def backend(self) -> Backend:
        r"""
        The backend used by this analyzer tool.
        """
        return self._backend

    def _simulate(
        self,
        pubs: Sequence[EstimatorPubLike],
        with_noise: bool,
        cliffordize: bool,
        seed_simulator: Optional[int],
        precision: float = 0,
    ) -> NeatResult:
        r"""
        Perform a noisy or noiseless simulation of the estimator task specified by ``pubs``.

        Args:
            pubs: The PUBs specifying the estimation task of interest.
            with_noise: Whether to perform an ideal, noiseless simulation (``False``) or a noisy
                simulation (``True``).
            cliffordize: Whether or not to automatically apply the
                :class:`.~ConvertISAToClifford` transpiler pass to the given ``pubs`` before
                performing the simulations.
            seed_simulator: A seed for the simulator.
            precision: The target precision for the estimates of each expectation value in the
                returned results.

        Returns:
            The results of the simulation.
        """
        if cliffordize:
            coerced_pubs = self.to_clifford(pubs)
        else:
            coerced_pubs = [EstimatorPub.coerce(p) for p in pubs]

        backend_options = {
            "method": "stabilizer",
            "noise_model": self.noise_model if with_noise else None,
            "seed_simulator": seed_simulator,
        }
        estimator = AerEstimator(
            options={"backend_options": backend_options, "default_precision": precision}
        )

        aer_job = estimator.run(coerced_pubs)
        try:
            aer_result = aer_job.result()
        except QiskitError as err:
            if "invalid parameters" in str(err):
                raise ValueError(
                    "Couldn't run the simulation, likely because the given PUBs contain one or "
                    "more non-Clifford instructions. To fix, try setting ``cliffordize`` to "
                    "``True``."
                ) from err
            raise err

        pub_results = [NeatPubResult(r.data.evs) for r in aer_result]
        return NeatResult(pub_results)

    def ideal_sim(
        self,
        pubs: Sequence[EstimatorPubLike],
        cliffordize: bool = False,
        seed_simulator: Optional[int] = None,
        precision: float = 0,
    ) -> NeatResult:
        r"""
        Perform an ideal, noiseless simulation of the estimator task specified by ``pubs``.

        This function uses ``qiskit-aer``'s ``Estimator`` class to simulate the estimation task
        classically.

        .. note::
            To ensure scalability, every circuit in ``pubs`` is required to be a Clifford circuit,
            so that it can be simulated efficiently regardless of its size. For estimation tasks
            that involve non-Clifford circuits, the recommended workflow consists of mapping
            the non-Clifford circuits to the nearest Clifford circuits using the
            :class:`.~ConvertISAToClifford` transpiler pass, or equivalently, to use the Neat's
            :meth:`to_clifford` convenience method. Alternatively, setting ``cliffordize`` to
            ``True`` ensures that the :meth:`to_clifford` method is applied automatically to the
            given ``pubs`` prior to the simulation.

        Args:
            pubs: The PUBs specifying the estimation task of interest.
            cliffordize: Whether or not to automatically apply the
                :class:`.~ConvertISAToClifford` transpiler pass to the given ``pubs`` before
                performing the simulations.
            seed_simulator: A seed for the simulator.
            precision: The target precision for the estimates of each expectation value in the
                returned results.

        Returns:
            The results of the simulation.
        """
        return self._simulate(pubs, False, cliffordize, seed_simulator, precision)

    def noisy_sim(
        self,
        pubs: Sequence[EstimatorPubLike],
        cliffordize: bool = False,
        seed_simulator: Optional[int] = None,
        precision: float = 0,
    ) -> NeatResult:
        r"""
        Perform a noisy simulation of the estimator task specified by ``pubs``.

        This function uses ``qiskit-aer``'s ``Estimator`` class to simulate the estimation task
        classically.

        .. note::
            To ensure scalability, every circuit in ``pubs`` is required to be a Clifford circuit,
            so that it can be simulated efficiently regardless of its size. For estimation tasks
            that involve non-Clifford circuits, the recommended workflow consists of mapping
            the non-Clifford circuits to the nearest Clifford circuits using the
            :class:`.~ConvertISAToClifford` transpiler pass, or equivalently, to use the Neat's
            :meth:`to_clifford` convenience method. Alternatively, setting ``cliffordize`` to
            ``True`` ensures that the :meth:`to_clifford` method is applied automatically to the
            given ``pubs`` prior to the simulation.

        Args:
            pubs: The PUBs specifying the estimation task of interest.
            cliffordize: Whether or not to automatically apply the
                :class:`.~ConvertISAToClifford` transpiler pass to the given ``pubs`` before
                performing the simulations.
            seed_simulator: A seed for the simulator.
            precision: The target precision for the estimates of each expectation value in the
                returned results.

        Returns:
            The results of the simulation.
        """
        return self._simulate(pubs, True, cliffordize, seed_simulator, precision)

    def to_clifford(self, pubs: Sequence[EstimatorPubLike]) -> list[EstimatorPub]:
        r"""
        Return the cliffordized version of the given ``pubs``.

        This convenience method runs the :class:`.~ConvertISAToClifford` transpiler pass on the
        PUBs' circuits.

        Args:
            pubs: The PUBs to turn into Clifford PUBs.

        Returns:
            The Clifford PUBs.
        """
        coerced_pubs = []
        for pub in pubs:
            coerced_pub = EstimatorPub.coerce(pub)
            coerced_pubs.append(
                EstimatorPub(
                    PassManager([ConvertISAToClifford()]).run(coerced_pub.circuit),
                    coerced_pub.observables,
                    coerced_pub.parameter_values,
                    coerced_pub.precision,
                    False,
                )
            )

        return coerced_pubs

    def __repr__(self) -> str:
        return f'Neat(backend="{self.backend().name}")'