# Copyright 2021 The Cirq Developers
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     https://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
"""Tooling for cross entropy benchmark experiments."""

from __future__ import annotations

from collections.abc import Iterable, Sequence
from dataclasses import dataclass
from typing import Any

import cirq
import matplotlib as mpl
import matplotlib.pyplot as plt
import numpy as np
import numpy.typing as npt
import pandas as pd
import scipy
import seaborn as sns
import tqdm.auto
import tqdm.contrib.itertools

from supermarq.qcvv.base_experiment import QCVVExperiment, QCVVResults, Sample


@dataclass
class XEBResults(QCVVResults):
    """Results from an XEB experiment."""

    _circuit_fidelities: pd.DataFrame | None = None
    _cycle_fidelity_estimate: float | None = None
    """Estimated cycle fidelity."""
    _cycle_fidelity_estimate_std: float | None = None
    """Standard deviation for the cycle fidelity estimate."""

    @property
    def cycle_fidelity_estimate(self) -> float:
        """Estimated cycle fidelity."""
        if self._cycle_fidelity_estimate is None:
            raise self._not_analyzed
        return self._cycle_fidelity_estimate

    @property
    def cycle_fidelity_estimate_std(self) -> float:
        """Standard deviation for the cycle fidelity estimate."""
        if self._cycle_fidelity_estimate_std is None:
            raise self._not_analyzed
        return self._cycle_fidelity_estimate_std

    def _analytical_data(self) -> pd.DataFrame:
        """Create a copy of `self.data` with analytical probabilities in place of experimental data.

        Returns:
            The `pd.DataFrame` containing analytical probabilities.
        """
        assert self.data is not None

        indices = [f"{i:0>{self.num_qubits}b}" for i in range(2**self.num_qubits)]

        analytic_probabilities: list[npt.NDArray[np.float64]] = []
        for _, row in self.data.iterrows():
            sample = self.experiment[row.uuid]
            subcircuit = sample.circuit.copy()
            subcircuit.clear_operations_touching(
                subcircuit.all_qubits().difference(self.qubits), range(len(subcircuit))
            )
            analytic_final_state = cirq.final_state_vector(
                subcircuit,
                qubit_order=self.qubits,
                ignore_terminal_measurements=True,
                dtype=np.complex128,
            )
            analytic_probabilities.append(np.abs(analytic_final_state) ** 2)

        analytical_data = self.data.drop(columns=indices, errors="ignore")
        analytical_data[indices] = analytic_probabilities
        return analytical_data

    def _analyze(self) -> None:
        """Analyse the results and calculate the estimated circuit fidelity.

        Args:
            plot_results (optional): Whether to generate the data plots. Defaults to True.

        Raises:
            RuntimeError: If there is no data stored.

        Returns:
           The final results from the experiment.
        """
        if self.data is None:
            raise RuntimeError("No data stored. Cannot perform analysis.")

        analytical_data = self._analytical_data()

        indices = [f"{i:0>{self.num_qubits}b}" for i in range(2**self.num_qubits)]
        self.data["sum_p(x)p^(x)"] = (self.data[indices] * analytical_data[indices]).sum(axis=1)
        self.data["sum_p(x)p(x)"] = (analytical_data[indices] ** 2).sum(axis=1)

        # Fit a linear model for each cycle depth to estimate the circuit fidelity
        records = []
        for depth in set(self.data.cycle_depth):
            df = self.data[self.data.cycle_depth == depth]
            fit = scipy.stats.linregress(
                x=df["sum_p(x)p(x)"] - 1 / 2**self.num_qubits,
                y=df["sum_p(x)p^(x)"] - 1 / 2**self.num_qubits,
            )

            records.append(
                {
                    "cycle_depth": depth,
                    "circuit_fidelity_estimate": fit.slope,
                    "circuit_fidelity_estimate_std": fit.stderr,
                }
            )

        self._circuit_fidelities = pd.DataFrame(records)

        self._circuit_fidelities["log_fidelity_estimate"] = np.log(
            self._circuit_fidelities.circuit_fidelity_estimate
        )

        # Fit a linear model to the depth ~ log(fidelity) to approximate the cycle fidelity
        cycle_fit = scipy.stats.linregress(
            x=self._circuit_fidelities.cycle_depth,
            y=np.log(self._circuit_fidelities.circuit_fidelity_estimate),
        )
        self._cycle_fidelity_estimate = np.exp(cycle_fit.slope)
        self._cycle_fidelity_estimate_std = self.cycle_fidelity_estimate * cycle_fit.stderr
        self._zero_cycle_fidelity_estimate = np.exp(cycle_fit.intercept)

    def plot_results(
        self,
        filename: str | None = None,
    ) -> plt.Figure:
        """Plot the experiment data and the corresponding fits.

        Args:
            filename: Optional argument providing a filename to save the plots to. Defaults to None,
                indicating not to save the plot.

        Returns:
            A single matplotlib figure containing both the linear fit per cycle depth and the decay
            with cycle depth.

        Raises:
            RuntimeError: If there is no data stored.
        """
        if self.data is None:
            raise RuntimeError("No data stored. Cannot plot results.")

        if self._circuit_fidelities is None:
            raise RuntimeError(
                "No stored dataframe of circuit fidelities. Something has gone wrong."
            )

        fig, axs = plt.subplots(1, 2, figsize=(10, 4.8))
        colours = sns.color_palette("dark:r", n_colors=len(self.data.cycle_depth.unique()))
        for depth, color in zip(sorted(self.data.cycle_depth.unique()), colours):
            sns.regplot(
                data=self.data[self.data.cycle_depth == depth],
                x="sum_p(x)p(x)",
                y="sum_p(x)p^(x)",
                ci=None,
                ax=axs[0],
                color=color,
                label=depth,
            )
        axs[0].legend(title="Cycle depth", bbox_to_anchor=(1.0, 0.5), loc="center left")
        axs[0].set_xlabel(r"$\sum p(x)^2$", fontsize=15)
        axs[0].set_ylabel(r"$\sum p(x) \hat{p}(x)$", fontsize=15)
        axs[0].set_title(r"Linear fit per cycle depth", fontsize=15, wrap=True)

        sns.regplot(
            data=self._circuit_fidelities,
            x="cycle_depth",
            y="circuit_fidelity_estimate",
            ax=axs[1],
            fit_reg=False,
            color="tab:red",
        )
        axs[1].set_xlabel(r"Cycle depth", fontsize=15)
        axs[1].set_ylabel(r"Circuit fidelity", fontsize=15)
        axs[1].set_title(r"Exponential decay of circuit fidelity", fontsize=15, wrap=True)

        # Add fit line
        x = np.linspace(
            self._circuit_fidelities.cycle_depth.min(), self._circuit_fidelities.cycle_depth.max()
        )
        y = self.cycle_fidelity_estimate**x * self._zero_cycle_fidelity_estimate
        y_p = (
            self.cycle_fidelity_estimate + 1.96 * self.cycle_fidelity_estimate_std
        ) ** x * self._zero_cycle_fidelity_estimate
        y_m = (
            self.cycle_fidelity_estimate - 1.96 * self.cycle_fidelity_estimate_std
        ) ** x * self._zero_cycle_fidelity_estimate
        axs[1].plot(x, y, color="tab:red", linewidth=2)
        axs[1].fill_between(x, y_m, y_p, alpha=0.25, color="tab:red", label="95% CI")
        axs[1].legend()

        fig.tight_layout()

        if filename is not None:
            fig.savefig(filename, bbox_inches="tight")

        return fig

    def print_results(self) -> None:
        """Prints the key results data."""
        print(  # noqa: T201
            f"Estimated cycle fidelity: {self.cycle_fidelity_estimate:.5} "
            f"+/- {self.cycle_fidelity_estimate_std:.5}"
        )

    def plot_speckle(self, filename: str | None = None) -> plt.Figure:
        """Creates the speckle plot of the XEB data. See Fig. S18 of
        https://arxiv.org/abs/1910.11333 for an explanation of this plot.

        Args:
            filename: Optional argument providing a filename to save the plots to. Defaults to None,
                indicating not to save the plot.

        Returns:
            A matplotlib figure with the speckle plot.

        Raises:
            RuntimeError: If there is no data stored.
        """
        if self.data is None:
            raise RuntimeError("No data stored. Cannot plot results.")

        indices = [f"{i:0>{self.num_qubits}b}" for i in range(2**self.num_qubits)]

        # Reformat dataframe
        df = pd.melt(
            self.data,
            value_vars=indices,
            id_vars=["cycle_depth", "circuit_realization"],
            var_name="bitstring",
        )

        # Create the axes needed
        fig, axs = plt.subplot_mosaic(
            [[f"P({idx})", "cbar", ".", "Decay"] for idx in indices],
            width_ratios=[1, 0.05, 0.05, 1],
            figsize=(12, 4.8),
        )
        fig.subplots_adjust(hspace=0)

        # Plot the heatmaps
        for k, bitstring in enumerate(indices):
            ax = axs[f"P({bitstring})"]

            data = df[df["bitstring"] == bitstring].pivot(
                index="circuit_realization", columns="cycle_depth", values="value"
            )
            cmap = mpl.colormaps["rocket"]
            # norm = mpl.colors.Normalize(0, 1)  # or vmin, vmax # noqa: ERA001
            sns.heatmap(data, vmin=0, vmax=1, ax=ax, cbar_ax=axs["cbar"], cmap=cmap)
            ax.set_ylabel("")
            ax.set_xlabel("")
            ax.set_yticks([])
            ax.set_yticklabels([])
            plt.text(
                0.99,
                0.90,
                f"P({bitstring})",
                ha="right",
                va="top",
                transform=ax.transAxes,
                color="white",
            )
            if k == 0:
                ax.set_title("Speckle plots")
            else:
                ax.axhline(y=0, linewidth=1.5, color="white", linestyle="--")

            if k == 2**self.num_qubits - 1:
                ax.set_xlabel("Cycle depth")

        # Format colour bar
        axs["cbar"].set_ylabel("Probability")
        axs["cbar"].yaxis.set_label_position("right")
        axs["cbar"].yaxis.tick_left()

        # Plot the decay of purity

        # Calculate the average std of the probability distributions.
        purity_data = (
            df.groupby(by=["cycle_depth"])
            .std(numeric_only=True)
            .reset_index()
            .rename(columns={"value": "sqrt_speckle_purity"})
            .drop(columns=["circuit_realization"])
        )
        # Rescale the purity estimate according to Porter-Thomas distribution
        dim = np.prod([q.dimension for q in self.qubits])
        purity_data["sqrt_speckle_purity"] = purity_data["sqrt_speckle_purity"] * np.sqrt(
            dim**2 * (dim + 1) / (dim - 1)
        )

        # Plot decay
        sns.regplot(
            data=purity_data,
            x="cycle_depth",
            y="sqrt_speckle_purity",
            logx=True,
            ax=axs["Decay"],
        )
        # purity_plot.tight_layout() # noqa: ERA001
        axs["Decay"].set_xlabel(r"Cycle depth")
        axs["Decay"].set_ylabel(r"$\sqrt{\mathrm{Purity}}$")
        axs["Decay"].set_title(r"Purity Decay")
        # Estimate decay coefficient
        purity_fit = scipy.stats.linregress(
            x=purity_data["cycle_depth"],
            y=np.log(purity_data["sqrt_speckle_purity"]),
        )
        # Add label with coefficient estimate
        axs["Decay"].text(
            0.95,
            0.95,
            (
                f"Decay coefficient: {np.exp(purity_fit.slope):4f} "
                f"+/- {np.exp(purity_fit.slope) * purity_fit.stderr:4f}"
            ),
            ha="right",
            va="center",
            transform=axs["Decay"].transAxes,
        )

        if filename is not None:
            fig.savefig(filename, bbox_inches="tight")

        return fig


class XEB(QCVVExperiment[XEBResults]):
    r"""Cross-entropy benchmarking (XEB) experiment.

    The XEB experiment can be used to estimate the combined fidelity of a repeating cycle of gates.
    Cycles are made up of randomly selected single qubit gates and a constant layer of interest.
    This is illustrated as follows:

    For each randomly generated circuit, with a given number of cycle, we compare the
    simulated state probabilities, :math:`p(x)` with those achieved by running the circuit
    on a given target, :math:`\hat{p}(x)`. The fidelity of a circuit containing :math:`d`
    cycles, :math:`f_d` can then be estimated as

    .. math::

        \sum_{x \in \{0, 1\}^n} p(x) \hat{p}(x) - \frac{1}{2^n} =
        f_d \left(\sum_{x \in \{0, 1\}^n} p(x)^2 -  \frac{1}{2^n}\right)

    We can therefore fit a linear model to estimate the value of :math:`f_d`. We the estimate
    the fidelity of the cycle, :math:`f_{\mathrm{cycle}}` as

    .. math::

        f_d = A(f_{cycle})^d

    Thus fitting another linear model to :math:`\log(f_d) \sim d` provides us with an estimate
    of the cycle fidelity.

    For more details see: https://www.nature.com/articles/s41586-019-1666-5
    """

    def __init__(
        self,
        num_circuits: int,
        cycle_depths: Iterable[int],
        interleaved_layer: cirq.Gate | cirq.OP_TREE | None = cirq.CZ,
        single_qubit_gate_set: Sequence[cirq.Gate] | None = None,
        *,
        random_seed: int | np.random.Generator | None = None,
        _samples: list[Sample] | None = None,
        **kwargs: str,
    ) -> None:
        """Initializes a cross-entropy benchmarking experiment.

        Args:
            num_circuits: Number of circuits to sample.
            cycle_depths: The cycle depths to sample.
            interleaved_layer: The gate or operation(s) to interleave between the single qubit
                gates. If None then no gates are interleaved.
            single_qubit_gate_set: Optional list of single qubit gates to randomly sample from when
                generating random circuits. If not provided defaults to phased X gates with 1/4 pi
                intervals.
            random_seed: An optional seed to use for randomization.
            kwargs: Any additional supported string keyword args.
        """
        if interleaved_layer is None:
            qubits: Sequence[cirq.Qid] = cirq.LineQubit.range(2)

        elif isinstance(interleaved_layer, cirq.Gate):
            qubits = (
                cirq.LineQubit.range(cirq.num_qubits(interleaved_layer))
                if all(d == 2 for d in cirq.qid_shape(interleaved_layer))
                else cirq.LineQid.for_gate(interleaved_layer)
            )
            interleaved_layer = interleaved_layer.on(*qubits)

        elif isinstance(interleaved_layer, (cirq.Operation, cirq.Moment)):
            qubits = sorted(interleaved_layer.qubits)

        else:
            interleaved_layer = cirq.Circuit(interleaved_layer)
            qubits = sorted(interleaved_layer.all_qubits())

        self.interleaved_layer: cirq.OP_TREE | None = interleaved_layer
        """The layer to interleave."""

        self.single_qubit_gate_set: list[cirq.Gate]
        """The single qubit gates to randomly sample from"""

        if single_qubit_gate_set is None:
            self.single_qubit_gate_set = [
                cirq.PhasedXPowGate(exponent=0.5, phase_exponent=phase_exponent)
                for phase_exponent in [0.0, 0.25, 0.5]
            ]
        else:
            self.single_qubit_gate_set = list(single_qubit_gate_set)

        super().__init__(
            qubits=qubits,
            num_circuits=num_circuits,
            cycle_depths=cycle_depths,
            random_seed=random_seed,
            results_cls=XEBResults,
            _samples=_samples,
            **kwargs,
        )

    def independent_qubit_groups(self) -> list[tuple[cirq.Qid, ...]]:
        """Get all independent subsets of qubits in this experiment.

        Returns:
            A list of disjoint tuples of `cirq.Qid` objects, each of which can be analyzed as an
            independent XEB experiment.
        """
        if self.interleaved_layer:
            qubit_sets = cirq.Circuit(self.interleaved_layer).get_independent_qubit_sets()
            return sorted((tuple(sorted(qubit_set)) for qubit_set in qubit_sets), key=min)

        return [self.qubits]

    ###################
    # Private Methods #
    ###################
    def _build_circuits(
        self,
        num_circuits: int,
        cycle_depths: Iterable[int],
    ) -> Sequence[Sample]:
        """Build a list of random circuits to perform the XEB experiment with.

        Args:
            num_circuits: Number of circuits to generate.
            cycle_depths: An iterable of the different numbers of cycles to include in each circuit.

        Returns:
            The list of experiment samples.
        """
        random_circuits = []
        for k, depth in tqdm.contrib.itertools.product(
            range(num_circuits), cycle_depths, desc="Building circuits"
        ):
            # Choose single-qubit gates, avoiding repeats on the same qubit in sequential layers
            num_choices = len(self.single_qubit_gate_set)
            block_repeats = num_choices > 1
            chosen_gate_indices = np.append(
                self._rng.integers(0, num_choices, size=(1, self.num_qubits)),
                self._rng.integers(block_repeats, num_choices, size=(depth, self.num_qubits)),
                axis=0,
            )
            chosen_gate_indices = np.add.accumulate(chosen_gate_indices, axis=0) % num_choices

            circuit = cirq.Circuit(
                self.single_qubit_gate_set[gate_index].on(qubit)
                for gate_indices_in_layer in chosen_gate_indices
                for gate_index, qubit in zip(gate_indices_in_layer, self.qubits)
            )

            circuit = self._interleave_layer(circuit, self.interleaved_layer)

            random_circuits.append(
                Sample(
                    circuit=circuit + cirq.measure(*self.qubits),
                    data={
                        "circuit_depth": len(circuit),
                        "cycle_depth": depth,
                        "interleaved_layer": str(self.interleaved_layer),
                    },
                    circuit_realization=k,
                )
            )

        return random_circuits

    def _json_dict_(self) -> dict[str, Any]:
        """Converts the experiment to a json-able dictionary that can be used to recreate the
        experiment object. Note that the state of the random number generator is not stored.

        Returns:
            Json-able dictionary of the experiment data.
        """
        json_dict = super()._json_dict_()
        json_dict["interleaved_layer"] = self.interleaved_layer
        json_dict["single_qubit_gate_set"] = self.single_qubit_gate_set
        del json_dict["qubits"]

        return json_dict
