#!/usr/bin/python3
# -*- coding: utf-8 -*-
#
# Copyright (c) 2022 Baidu, Inc. All Rights Reserved.
#
# 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
#
#     http://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.

"""
Utility functions used in the ``qcompute_qep.mitigation`` module.
"""

from matplotlib.figure import Figure
from typing import List
import matplotlib.pyplot as plt
from matplotlib import rc, pylab
import numpy as np


def plot_zne_sequences(data_ideal: List[float],
                       data_miti: List[float],
                       data_noisy: np.ndarray,
                       scale_factors: List[float],
                       title: str = None,
                       fname: str = None) -> Figure:
    r"""Visualize the dataset generated by the Zero-Noise Extrapolation method.

    Assume the number of data points is :math:`n` and the length of @scale_factors is :math:`k`, then

    + `@data_ideal` is a :math:`n\times 1` vector of ideal values;
    + `@data_miti` is a :math:`n\times 1` vector of mitigated values; and
    + `@data_noisy` is a :math:`n\times k` matrix of noisy values, each row records the set of rescaled values.

    :param data_ideal: np.ndarray, the set of ideal values
    :param data_miti: np.ndarray, the set of mitigated values
    :param data_noisy: np.ndarray, the set of noisy values
    :param scale_factors: List[float], the set of rescaling factors
    :param title: str, default to None, the title of the figure
    :param fname: str, default to None, the file name of the figure that will be saved
    :return: Figure, the created figure instance
    """
    n = len(data_ideal)
    x_labels = list(range(1, n + 1))

    # Visualize data
    fig = plt.figure(figsize=(8, 5))
    ax = plt.gca()

    # Plot ideal data
    plt.plot(x_labels, data_ideal, ':', color='black', alpha=0.8, linewidth=1.5, label='Ideal')

    # Plot noisy data
    # Setup the colormap, which should guarantee that each data type has its own representative color
    cm = pylab.get_cmap('Dark2')  # Dark2, Accent, Paired
    markers = ['^', 'v', '<', '>']

    for i, lam in enumerate(scale_factors):
        plt.plot(x_labels, data_noisy[i], '-.', marker=markers[i % len(markers)],
                 linewidth=1,
                 markersize=8,
                 alpha=0.8,
                 color=cm(1.*(i % 8)/8), label=r'Noisy: $\lambda={}$'.format(lam))

    # Plot mitigated data. The color is fixed to red, the marker is fixed to star.
    # We plot the mitigated data at last so that they will not be covered by the noisy data points.
    plt.plot(x_labels, data_miti, '-*', color='red', linewidth=1, markersize=12, label='ZNE')

    plt.xticks(x_labels)
    plt.legend()

    plt.title(title)
    plt.xlabel(r'Size of the Clifford Sequence $L$', fontsize=14)
    plt.ylabel(r'Expectation Value', fontsize=14)

    # Save the figure
    fname = fname if fname is not None else title
    plt.savefig(fname, dpi=400)
    return fig
