# This code is part of Qiskit.
#
# (C) Copyright IBM 2020.
#
# 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.

# pylint: disable=invalid-name

"""Matplotlib plotter API."""

from typing import Optional

import matplotlib
import matplotlib.pyplot as plt
import numpy as np
from matplotlib.patches import Rectangle

from qiskit.visualization.exceptions import VisualizationError
from qiskit.visualization.pulse_v2 import core, drawings, types
from qiskit.visualization.pulse_v2.plotters.base_plotter import BasePlotter
from qiskit.visualization.utils import matplotlib_close_if_inline


class Mpl2DPlotter(BasePlotter):
    """Matplotlib API for pulse drawer.

    This plotter places canvas charts along y axis of 2D canvas with vertical offset.
    Each chart is map to X-Y axis of the canvas.
    """

    def __init__(self, canvas: core.DrawerCanvas, axis: Optional[plt.Axes] = None):
        """Create new plotter.

        Args:
            canvas: Configured drawer canvas object. Canvas object should be updated
                with `.update` method before set to the plotter API.
            axis: Matplotlib axis object. When `axis` is provided, the plotter updates
                given axis instead of creating and returning new matplotlib figure.
        """
        super().__init__(canvas=canvas)

        # calculate height of all charts
        canvas_height = 0
        for chart in self.canvas.charts:
            if not chart.is_active and not self.canvas.formatter["control.show_empty_channel"]:
                continue
            canvas_height += chart.vmax - chart.vmin
        # set min canvas_height size
        canvas_height = max(canvas_height, 0.1)

        if axis is None:
            fig_h = canvas_height * self.canvas.formatter["general.fig_chart_height"]
            fig_w = self.canvas.formatter["general.fig_width"]

            self.figure = plt.figure(figsize=(fig_w, fig_h))
            self.ax = self.figure.add_subplot(1, 1, 1)
        else:
            self.figure = axis.figure
            self.ax = axis

        self.initialize_canvas()

    def initialize_canvas(self):
        """Format appearance of matplotlib canvas."""
        self.ax.set_facecolor(self.canvas.formatter["color.background"])

        # axis labels
        self.ax.set_yticklabels([])
        self.ax.yaxis.set_tick_params(left=False)

    def draw(self):
        """Output drawings stored in canvas object."""
        # axis configuration
        axis_config = self.canvas.layout["time_axis_map"](
            time_window=self.canvas.time_range,
            axis_breaks=self.canvas.time_breaks,
            dt=self.canvas.device.dt,
        )

        current_y = 0
        margin_y = self.canvas.formatter["margin.between_channel"]
        for chart in self.canvas.charts:
            if not chart.is_active and not self.canvas.formatter["control.show_empty_channel"]:
                continue
            current_y -= chart.vmax
            for _, data in chart.collections:
                # calculate scaling factor
                if not data.ignore_scaling:
                    # product of channel-wise scaling and chart level scaling
                    scale = max(self.canvas.chan_scales.get(chan, 1.0) for chan in data.channels)
                    scale *= chart.scale
                else:
                    scale = 1.0

                x = data.xvals
                y = scale * data.yvals + current_y

                if isinstance(data, drawings.LineData):
                    # line object
                    if data.fill:
                        self.ax.fill_between(x, y1=y, y2=current_y * np.ones_like(y), **data.styles)
                    else:
                        self.ax.plot(x, y, **data.styles)
                elif isinstance(data, drawings.TextData):
                    # text object
                    text = rf"${data.latex}$" if data.latex else data.text
                    # replace dynamic text
                    text = text.replace(types.DynamicString.SCALE, f"{chart.scale:.1f}")
                    self.ax.text(x=x[0], y=y[0], s=text, **data.styles)
                elif isinstance(data, drawings.BoxData):
                    xy = x[0], y[0]
                    box = Rectangle(
                        xy, width=x[1] - x[0], height=y[1] - y[0], fill=True, **data.styles
                    )
                    self.ax.add_patch(box)
                else:
                    VisualizationError(
                        "Data {name} is not supported "
                        "by {plotter}".format(name=data, plotter=self.__class__.__name__)
                    )
            # axis break
            for pos in axis_config.axis_break_pos:
                self.ax.text(
                    x=pos,
                    y=current_y,
                    s="//",
                    ha="center",
                    va="center",
                    zorder=self.canvas.formatter["layer.axis_label"],
                    fontsize=self.canvas.formatter["text_size.axis_break_symbol"],
                    rotation=180,
                )

            # shift chart position
            current_y += chart.vmin - margin_y

        # remove the last margin
        current_y += margin_y

        y_max = self.canvas.formatter["margin.top"]
        y_min = current_y - self.canvas.formatter["margin.bottom"]

        # plot axis break line
        for pos in axis_config.axis_break_pos:
            self.ax.plot(
                [pos, pos],
                [y_min, y_max],
                zorder=self.canvas.formatter["layer.fill_waveform"] + 1,
                linewidth=self.canvas.formatter["line_width.axis_break"],
                color=self.canvas.formatter["color.background"],
            )

        # label
        self.ax.set_xticks(list(axis_config.axis_map.keys()))
        self.ax.set_xticklabels(
            list(axis_config.axis_map.values()),
            fontsize=self.canvas.formatter["text_size.axis_label"],
        )
        self.ax.set_xlabel(
            axis_config.label, fontsize=self.canvas.formatter["text_size.axis_label"]
        )

        # boundary
        if axis_config.window == (0, 0):
            self.ax.set_xlim(0, 1)
        else:
            self.ax.set_xlim(*axis_config.window)
        self.ax.set_ylim(y_min, y_max)

        # title
        if self.canvas.fig_title:
            self.ax.text(
                x=axis_config.window[0],
                y=y_max,
                s=self.canvas.fig_title,
                ha="left",
                va="bottom",
                zorder=self.canvas.formatter["layer.fig_title"],
                color=self.canvas.formatter["color.fig_title"],
                size=self.canvas.formatter["text_size.fig_title"],
            )

    def get_image(self, interactive: bool = False) -> matplotlib.pyplot.Figure:
        """Get image data to return.

        Args:
            interactive: When set `True` show the circuit in a new window.
                This depends on the matplotlib backend being used supporting this.

        Returns:
            Matplotlib figure data.
        """
        matplotlib_close_if_inline(self.figure)

        if self.figure and interactive:
            self.figure.show()

        return self.figure
