"""Matplotlib FigureCanvas widget for manual scoring."""

from collections.abc import Callable

import matplotlib.axes
import matplotlib.ticker as mticker
import numpy as np
from matplotlib.backend_bases import MouseButton
from matplotlib.backends.backend_qtagg import FigureCanvasQTAgg as FigureCanvas
from matplotlib.figure import Figure
from matplotlib.gridspec import GridSpec
from matplotlib.patches import Rectangle
from matplotlib.widgets import RectangleSelector
from PySide6 import QtWidgets

from accusleepy.brain_state_set import BrainStateSet

# upper limit of spectrogram y-axis, in Hz
SPEC_UPPER_F = 30
# interval of spectrogram y-axis ticks, in Hz
SPEC_Y_TICK_INTERVAL = 10

# margins around subplots in the figure
SUBPLOT_TOP_MARGIN = 0.98
SUBPLOT_BOTTOM_MARGIN = 0.02
SUBPLOT_LEFT_MARGIN = 0.07
SUBPLOT_RIGHT_MARGIN = 0.95

# maximum number of x-axis ticks to show on the lower plot
MAX_LOWER_X_TICK_N = 7


class MplWidget(QtWidgets.QWidget):
    """Widget that displays a matplotlib FigureCanvas"""

    def __init__(self, parent=None):
        QtWidgets.QWidget.__init__(self, parent)
        # set up the canvas and store a reference to its axes
        self.canvas = FigureCanvas(Figure())
        self.canvas.axes = None

        # set the widget layout and remove the margins
        vertical_layout = QtWidgets.QVBoxLayout()
        vertical_layout.addWidget(self.canvas)
        vertical_layout.setContentsMargins(0, 0, 0, 0)
        self.setLayout(vertical_layout)

        # given during the setup process
        self.epoch_length = None

        # upper plot references
        self.upper_marker = None
        self.label_img_ref = None
        self.spec_ref = None
        self.roi = None
        self.editing_patch = None
        self.roi_patch = None

        # lower plot references
        self.eeg_line = None
        self.emg_line = None
        self.top_marker = None
        self.bottom_marker = None
        self.eeg_epoch_plot = None
        self.emg_epoch_plot = None

    def setup_upper_figure(
        self,
        n_epochs: int,
        label_img: np.array,
        confidence_scores: np.array,
        confidence_img: np.array,
        spec: np.array,
        f: np.array,
        emg: np.array,
        epochs_to_show: int,
        label_display_options: np.array,
        brain_state_set: BrainStateSet,
        roi_function: Callable,
    ):
        """Initialize upper FigureCanvas for the manual scoring GUI

        :param n_epochs: number of epochs in the recording
        :param label_img: brain state labels, as an image
        :param confidence_scores: confidence scores
        :param confidence_img: confidence scores, as an image
        :param spec: EEG spectrogram image
        :param f: EEG spectrogram frequency axis
        :param emg: EMG RMS per epoch
        :param epochs_to_show: number of epochs to show in the lower plot
        :param label_display_options: valid brain state y-axis locations
        :param brain_state_set: set of brain states options
        :param roi_function: callback for ROI selection
        """
        # references to parts of the epoch marker
        self.upper_marker = list()

        # subplot layout
        height_ratios = [2, 8, 2, 12, 13]
        gs1 = GridSpec(5, 1, hspace=0.02, height_ratios=height_ratios)
        gs2 = GridSpec(5, 1, hspace=0.4, height_ratios=height_ratios)
        axes = list()
        axes.append(self.canvas.figure.add_subplot(gs1[0]))
        axes.append(self.canvas.figure.add_subplot(gs1[1]))
        axes.append(self.canvas.figure.add_subplot(gs1[2]))
        axes.append(self.canvas.figure.add_subplot(gs1[3]))
        axes.append(self.canvas.figure.add_subplot(gs2[4]))

        # subplots have different axes limits
        for i in [0, 1, 2, 4]:
            axes[i].set_xlim((-0.5, n_epochs - 0.5))
        axes[3].set_xlim(0, n_epochs)

        # confidence score subplot
        if confidence_scores is None:
            axes[0].set_visible(False)
        else:
            axes[0].set_ylim([-0.5, 0.5])
            axes[0].set_xticks([])
            axes[0].set_yticks([0])
            axes[0].set_yticklabels(["Conf."])
            axes[0].tick_params(axis="y", color="white")
            axes[0].imshow(
                confidence_img, aspect="auto", origin="lower", interpolation="None"
            )
            confidence_x = (
                np.repeat(list(range(len(confidence_scores) + 1)), 2)[1:-1] - 0.5
            )
            confidence_y = np.repeat(confidence_scores, 2) - 0.5
            axes[0].plot(confidence_x, confidence_y, "k", linewidth=0.5)
            for side in ["left", "right", "bottom", "top"]:
                axes[0].spines[side].set_visible(False)

        # brain state subplot
        axes[1].set_ylim(
            [-0.5, np.max(label_display_options) - np.min(label_display_options) + 0.5]
        )
        axes[1].set_xticks([])
        axes[1].set_yticks(
            label_display_options - np.min(label_display_options),
        )
        axes[1].set_yticklabels([b.name for b in brain_state_set.brain_states])
        ax2 = axes[1].secondary_yaxis("right")
        ax2.set_yticks(
            label_display_options - np.min(label_display_options),
        )
        ax2.set_yticklabels([b.digit for b in brain_state_set.brain_states])
        self.label_img_ref = axes[1].imshow(
            label_img, aspect="auto", origin="lower", interpolation="None"
        )
        # add patch to dim the display when creating an ROI
        self.editing_patch = axes[1].add_patch(
            Rectangle(
                xy=(-0.5, -0.5),
                width=n_epochs,
                height=np.max(label_display_options)
                - np.min(label_display_options)
                + 1,
                color="white",
                edgecolor=None,
                alpha=0.4,
                visible=False,
            )
        )
        # add the ROI selection widget, but disable it until it's needed
        self.roi = RectangleSelector(
            ax=axes[1],
            onselect=roi_function,
            interactive=False,
            button=MouseButton(1),
        )
        self.roi.set_active(False)
        # keep a reference to the ROI patch so we can change its color later
        # index 0 is the "editing_patch" created earlier
        self.roi_patch = [c for c in axes[1].get_children() if type(c) is Rectangle][1]

        # epoch marker subplot
        axes[2].set_ylim((0, 1))
        axes[2].axis("off")
        self.upper_marker.append(
            axes[2].plot([-0.5, epochs_to_show - 0.5], [0.5, 0.5], "r")[0]
        )
        self.upper_marker.append(axes[2].plot([0], [0.5], "rD")[0])

        # EEG spectrogram subplot
        # select subset of frequencies to show
        f = f[f <= SPEC_UPPER_F]
        spec = spec[0 : len(f), :]
        axes[3].set_ylabel("EEG", rotation="horizontal", ha="right")
        axes[3].set_yticks(
            np.linspace(
                0,
                len(f),
                1 + round(SPEC_UPPER_F / SPEC_Y_TICK_INTERVAL),
            ),
        )
        axes[3].set_yticklabels(
            np.arange(0, SPEC_UPPER_F + SPEC_Y_TICK_INTERVAL, SPEC_Y_TICK_INTERVAL)
        )
        axes[3].tick_params(axis="both", which="major", labelsize=8)
        axes[3].xaxis.set_major_formatter(mticker.FuncFormatter(self.time_formatter))
        self.spec_ref = axes[3].imshow(
            spec,
            vmin=np.percentile(spec, 2),
            vmax=np.percentile(spec, 98),
            aspect="auto",
            origin="lower",
            interpolation="None",
            extent=(
                0,
                n_epochs,
                -0.5,
                len(f) + 0.5,
            ),
        )

        # EMG subplot
        axes[4].set_xticks([])
        axes[4].set_yticks([])
        axes[4].set_ylabel("EMG", rotation="horizontal", ha="right")
        axes[4].plot(
            emg,
            "k",
            linewidth=0.5,
        )

        self.canvas.figure.subplots_adjust(
            left=SUBPLOT_LEFT_MARGIN,
            right=SUBPLOT_RIGHT_MARGIN,
            top=SUBPLOT_TOP_MARGIN,
            bottom=SUBPLOT_BOTTOM_MARGIN,
        )

        self.canvas.axes = axes

    def setup_lower_figure(
        self,
        label_img: np.array,
        sampling_rate: int | float,
        epochs_to_show: int,
        brain_state_set: BrainStateSet,
        label_display_options: np.array,
    ):
        """Initialize lower FigureCanvas for the manual scoring GUI

        :param label_img: brain state labels, as an image
        :param sampling_rate: EEG/EMG sampling rate, in Hz
        :param epochs_to_show: number of epochs to show in the lower plot
        :param brain_state_set: set of brain states options
        :param label_display_options: valid brain state y-axis locations
        """
        # number of samples in one epoch
        # we can expect this to be very close to an integer
        samples_per_epoch = round(sampling_rate * self.epoch_length)
        # number of EEG/EMG samples to plot
        samples_shown = round(samples_per_epoch * epochs_to_show)

        # references to parts of the epoch markers
        self.top_marker = list()
        self.bottom_marker = list()
        # epoch marker display parameters
        marker_dy = 0.25
        marker_y_offset_top = 0.02
        marker_y_offset_bottom = 0.01

        # subplot layout
        gs1 = GridSpec(3, 1, hspace=0)
        gs2 = GridSpec(3, 1, hspace=0.5)
        axes = list()
        axes.append(self.canvas.figure.add_subplot(gs1[0]))
        axes.append(self.canvas.figure.add_subplot(gs1[1]))
        axes.append(self.canvas.figure.add_subplot(gs2[2]))

        # EEG subplot
        axes[0].set_xticks([])
        axes[0].set_yticks([])
        axes[0].set_xlim((0, samples_shown))
        axes[0].set_ylim((-1, 1))
        axes[0].set_ylabel("EEG", rotation="horizontal", ha="right")
        self.eeg_line = axes[0].plot(
            np.zeros(samples_shown),
            "k",
            linewidth=0.5,
        )[0]
        self.eeg_epoch_plot = plot_eeg_emg_epochs(
            ax=axes[0],
            samples_per_epoch=samples_per_epoch,
            samples_shown=samples_shown,
            epochs_to_show=epochs_to_show,
        )
        self.eeg_epoch_plot.set(visible=False)

        # top epoch marker
        marker_x = [
            [0, 0],
            [0, samples_per_epoch],
            [samples_per_epoch, samples_per_epoch],
        ]
        marker_y = np.array(
            [
                [1 - marker_dy, 1],
                [1, 1],
                [1 - marker_dy, 1],
            ]
        )
        for x, y in zip(marker_x, marker_y, strict=True):
            self.top_marker.append(axes[0].plot(x, y - marker_y_offset_top, "r")[0])

        # EMG subplot
        axes[1].set_xticks(
            resample_x_ticks(
                np.arange(
                    0,
                    samples_shown,
                    samples_per_epoch,
                )
            )
        )
        axes[1].tick_params(axis="x", which="major", labelsize=8)
        axes[1].set_yticks([])
        axes[1].set_xlim((0, samples_shown))
        axes[1].set_ylim((-1, 1))
        axes[1].set_ylabel("EMG", rotation="horizontal", ha="right")

        self.emg_epoch_plot = plot_eeg_emg_epochs(
            ax=axes[1],
            samples_per_epoch=samples_per_epoch,
            samples_shown=samples_shown,
            epochs_to_show=epochs_to_show,
        )
        self.emg_epoch_plot.set(visible=False)

        self.emg_line = axes[1].plot(
            np.zeros(samples_shown),
            "k",
            linewidth=0.5,
        )[0]

        for x, y in zip(marker_x, marker_y, strict=True):
            self.bottom_marker.append(
                axes[1].plot(x, -1 * (y - marker_y_offset_bottom), "r")[0]
            )

        # brain state subplot
        axes[2].set_xticks([])
        axes[2].set_yticks(
            label_display_options - np.min(label_display_options),
        )
        axes[2].set_yticklabels([b.name for b in brain_state_set.brain_states])
        ax2 = axes[2].secondary_yaxis("right")
        ax2.set_yticks(
            label_display_options - np.min(label_display_options),
        )
        ax2.set_yticklabels([b.digit for b in brain_state_set.brain_states])
        axes[2].set_xlim((-0.5, epochs_to_show - 0.5))
        axes[2].set_ylim(
            [-0.5, np.max(label_display_options) - np.min(label_display_options) + 0.5]
        )
        self.label_img_ref = axes[2].imshow(
            label_img[:, 0:epochs_to_show, :],
            aspect="auto",
            origin="lower",
            interpolation="None",
        )

        self.canvas.figure.subplots_adjust(
            left=SUBPLOT_LEFT_MARGIN,
            right=SUBPLOT_RIGHT_MARGIN,
            top=SUBPLOT_TOP_MARGIN,
            bottom=SUBPLOT_BOTTOM_MARGIN,
        )

        self.canvas.axes = axes

    def time_formatter(self, x, pos):
        x = x * self.epoch_length
        return f"{int(x // 3600):02d}:{int(x // 60) % 60:02d}:{x % 60:05.2f}"


def resample_x_ticks(x_ticks: np.array) -> np.array:
    """Choose a subset of x_ticks to display

    The x-axis can get crowded if there are too many timestamps shown.
    This function finds a subset of evenly spaced x-axis ticks that
    includes the one at the beginning of the central epoch.

    :param x_ticks: full set of x_ticks
    :return: smaller subset of x_ticks
    """
    if len(x_ticks) <= MAX_LOWER_X_TICK_N:
        return x_ticks

    # number of ticks to the left of the central epoch
    # this will always be an integer
    nl = round((len(x_ticks) - 1) / 2)

    # search for even tick spacings that include the central epoch
    # if necessary, skip the leftmost tick
    for offset in [0, 1]:
        if (nl - offset) % 3 == 0:
            return x_ticks[offset :: round((nl - offset) / 3)]
        elif (nl - offset) % 2 == 0:
            return x_ticks[offset :: round((nl - offset) / 2)]


def plot_eeg_emg_epochs(
    ax: matplotlib.axes, samples_per_epoch: int, samples_shown: int, epochs_to_show: int
):
    """Plot epoch boundary markers on EEG/EMG data

    :param ax: axes on which to plot
    :param samples_per_epoch: number of samples in each epoch
    :param samples_shown: number of samples of data displayed
    :param epochs_to_show: number of epochs displayed
    :return: handle to plot object
    """
    return ax.plot(
        np.repeat(np.arange(samples_per_epoch, samples_shown, samples_per_epoch), 2),
        np.tile(np.array([-1.5, 1.5, 1.5, -1.5]), epochs_to_show)[
            : epochs_to_show * 2 - 2
        ],
        "gray",
        linewidth=1,
    )[0]
