"""Matplotlib visualisation of error budgets."""

from __future__ import annotations

import colorsys
import warnings
from collections.abc import Sequence
from pathlib import Path
from typing import TYPE_CHECKING, Any

import matplotlib.pyplot as plt
import numpy as np
from adjustText import adjust_text
from matplotlib.colors import ColorConverter, to_rgb

if TYPE_CHECKING:
    import numpy.typing as npt
    from matplotlib.axes import Axes
    from matplotlib.figure import Figure
    from matplotlib.text import Text

_RIVERLANE_PLOT_COLOURS: list[str] = [
    "#006F62",
    "#FF6A00",
    "#4B5BFF",
    "#E01500",
    "#003B40",
    "#D835A2",
    "#00A3A3",
    "#9B3EC6",
    "#54170E",
    "#FFD00B",
    "#0AD2F2",
]
_RIVERLANE_DARK_TEXT_COLOUR = "#0A1600"
_RIVERLANE_WHITE_TEXT_COLOUR = "#E5F0EF"
_RIVERLANE_LINE_COLOUR = "#003B40"
_RIVERLANE_NEUTRAL_GREY = "#DBDBDB"


def _scale_lightness(rgb: tuple[float, float, float] | str, scale_l: float):
    # From https://stackoverflow.com/a/60562502
    if isinstance(rgb, str):
        rgb = ColorConverter.to_rgb(rgb)
    # convert rgb to hls
    # Ignore E741 because "l" here is not really ambiguous and this function is small.
    h, l, s = colorsys.rgb_to_hls(*rgb)  # noqa: E741
    # manipulate h, l, s values and return as rgb
    return colorsys.hls_to_rgb(h, min(1, l * scale_l), s=s)


def _colour_brightness(colour: str) -> float:
    r, g, b = to_rgb(colour)
    return 0.2125 * r + 0.7152 * g + 0.0722 * b


def _text_size_in_data_coordinates(
    fig: Figure, ax: Axes, text: str
) -> tuple[float, float]:
    """Get the size of ``text`` when drawn on ``fig``.

    Adapted From https://stackoverflow.com/a/36959454.
    """
    get_renderer = getattr(fig.canvas, "get_renderer")
    assert get_renderer is not None
    renderer = get_renderer()
    text_object = ax.text(0.5, 0.5, text)
    bb = text_object.get_window_extent(renderer=renderer).transformed(
        ax.transData.inverted()
    )
    text_object.remove()
    return (bb.width, bb.height)


def _draw_line_from_text_to_position(
    fig: Figure, ax: Axes, text: Text, x: float, y: float, colour: str, t: float = 0.2
) -> None:
    """Draw a line between text and (x, y) but with the end at (x, y) slightly shortened."""
    # First, get the coordinates at which the line will intersect with the bounding-box
    # of the text.
    get_renderer = getattr(fig.canvas, "get_renderer")
    assert get_renderer is not None
    renderer = get_renderer()
    bl, tr = ax.transData.inverted().transform(text.get_window_extent(renderer))
    cx, cy = (tr + bl) / 2
    slope = (cy - y) / (cx - x)
    offset = cy - slope * cx
    x_intersect, y_intersect = (bl[1] - offset) / slope, bl[1]

    x, y = (1 - t) * x + t * x_intersect, (1 - t) * y + t * y_intersect
    ax.plot([x, x_intersect], [y, y_intersect], color=colour)


def plot_error_budget(
    contributions: npt.NDArray[np.floating] | Sequence[float],
    contribution_descriptions: Sequence[str],
    contributions_stddev: (
        npt.NDArray[np.floating] | Sequence[float | None] | None
    ) = None,
    fig: Figure | None = None,
    ax: Axes | None = None,
    save_to_filename: Path | None = None,
    excess_noise_description: str = "excess",
) -> tuple[Figure, Axes]:
    """Plot the provided error budget.

    Args:
        contributions (npt.NDArray[numpy.floating] | Sequence[float]):
            error-budget as an array of numbers.
        contribution_descriptions (Sequence[str]): textual description of each error
            mechanism involved in ``contributions``. Will be plotted on the error-budget
            figure.
        contributions_stddev (npt.NDArray[numpy.floating] | Sequence[float | None] | None):
            standard deviation of the provided values in ``contributions``. Default to
            ``None``, which does not draw the error-bars. If an array is provided,
            ``None`` values mean "no error-bars for this value" (this is often the case
            for the excess noise).
        fig (Figure | None): the figure to draw the error budget on. If not provided,
            a new figure and ax will be generated by calling ``plt.subplots()``.
        ax (Axes | None): the ax to draw the error budget on. If not provided,
            a new figure and ax will be generated by calling ``plt.subplots()``.
        save_to_filename (Path | None): if provided, the drawn figure will be saved to
            this path before returning.
        excess_noise_description (str): description of the budget that is associated
            with excess noise that has not been modelised in the error-budgeting
            process. If it matches (case-insensitive) one or more of the entries in
            ``contribution_descriptions``, this entry will be draw with grey hatches.

    Returns:
        Tuple[Figure, Axes]:
            The figure and ax that have been used to draw the resulting error-budget
            plot.

    Examples:
        Calculating per-round logical error probability and its standard deviation
        given number of fails, and number of shots for several rounds::

            import matplotlib.pyplot as plt

            fig, ax = plot_error_budget(
                contributions=[1.9, 0.3, 0.9, 0.6, 0.15, 0.4, 0.65],
                contribution_descriptions=[
                    "CZ",
                    "CZ stray int.",
                    "Data idle",
                    "Meas.",
                    "Leakage",
                    "1Q",
                    "Excess",
                ],
                excess_noise_description="Excess",
            )
            plt.show()
    """
    # Declare dimensions here. Every dimension is defined as a fraction of bar_height to
    # have scaling independent plots.
    bar_height: float = 0.5
    y_axis_offset_from_zero: float = bar_height / 10
    bar_center: float = y_axis_offset_from_zero + bar_height / 2
    bar_top: float = y_axis_offset_from_zero + bar_height
    non_fitting_texts_ys: list[float] = [bar_top + i * bar_height / 3 for i in (1, 2)]
    ymax = non_fitting_texts_ys[-1] + bar_height / 2

    # In case the user did not provide a figure or an axe, get it ourselves.
    if fig is None or ax is None:
        fig, ax = plt.subplots()

    # Plot a stacked horizontal bar plot. Do not include labels that do not fit in their
    # respective bar, but keep them to include them later.
    non_fitting_texts: list[str] = []
    nft_centres: list[float] = []
    nft_bar_colours: list[str] = []
    nft_bar_centres: list[float] = []
    twidths: list[float] = [
        _text_size_in_data_coordinates(fig, ax, d)[0] for d in contribution_descriptions
    ]

    offset: float = 0
    for i in range(len(contributions)):
        colour = _RIVERLANE_PLOT_COLOURS[i % len(_RIVERLANE_PLOT_COLOURS)]
        description = contribution_descriptions[i]
        contribution = contributions[i]
        twidth = twidths[i]
        # We need to handle the special case of the excess noise. Doing that with
        # kwargs.
        kwargs: dict[str, Any] = {
            "color": colour,
            "edgecolor": colour,
            "ecolor": _scale_lightness(colour, 0.8),
        }
        if description.lower() == excess_noise_description.lower():
            # Note that the RGB values of the colour below matter: we want it to be
            # white because the background will be white, and text colour will be picked
            # according to the background colour brightness.
            kwargs["color"] = "#FFFFFF00"
            kwargs["edgecolor"] = _RIVERLANE_NEUTRAL_GREY
            kwargs["hatch"] = "///"
            kwargs["ecolor"] = _scale_lightness(_RIVERLANE_NEUTRAL_GREY, 0.8)
        stddev = contributions_stddev[i] if contributions_stddev is not None else None
        xerr_bounds = (
            ((min(stddev, abs(contribution)),), (stddev,))
            if stddev is not None
            else None
        )
        bar = ax.barh(
            bar_center,
            contribution,
            height=bar_height,
            left=offset,
            xerr=xerr_bounds,
            capsize=2,
            **kwargs,
        )
        text_colour = (
            _RIVERLANE_DARK_TEXT_COLOUR
            if _colour_brightness(kwargs["color"]) > 0.5
            else _RIVERLANE_WHITE_TEXT_COLOUR
        )
        if twidth < 0.8 * contribution:
            ax.bar_label(
                bar,
                labels=[description],
                label_type="center",
                color=text_colour,
            )
        else:
            non_fitting_texts.append(description)
            nft_centres.append(offset + contribution / 2)
            nft_bar_colours.append(kwargs["color"])
            nft_bar_centres.append(offset + contribution / 2)
        offset += contribution

    # Ensure that axes are fixed before adding texts and automatically place them.
    # Fix the aspect ratio for a better looking plot.
    ax.set_aspect((offset / ymax) * (1 / 5))
    # Fixing the Y-axis bounds
    ax.set_ybound(0, ymax)
    # Add labels that did not fit into the bar.
    texts = [
        ax.text(
            center,
            non_fitting_texts_ys[i % len(non_fitting_texts_ys)],
            text,
            horizontalalignment="center",
            color=colour,
        )
        for i, (center, text, colour) in enumerate(
            zip(nft_centres, non_fitting_texts, nft_bar_colours)
        )
    ]

    if texts:
        ret = adjust_text(
            texts,
            target_x=nft_bar_centres,
            target_y=[bar_top for _ in range(len(nft_bar_centres))],
            avoid_self=False,
            only_move=dict.fromkeys(("explode", "pull", "text", "static"), "x"),
            ax=ax,
        )
        if ret is None:
            warnings.warn(
                "Call to adjust_text failed. Labels will likely not be well positioned."
            )
        else:
            texts, _ = ret
    # Draw lines
    for target, text, colour in zip(nft_bar_centres, texts, nft_bar_colours):
        _draw_line_from_text_to_position(fig, ax, text, target, bar_top, colour, t=0)

    # Removing the ticks on the Y-axis as these have no meaning here.
    ax.set_yticks([])
    # Remove part of the black frame because it is not useful.
    ax.spines["top"].set_visible(False)
    ax.spines["right"].set_visible(False)
    # Add axis labels
    ax.set_xlabel("Error budget, 1/Λ")

    if save_to_filename is not None:
        fig.savefig(save_to_filename)

    return fig, ax
