# Copyright (c) 2016 Electric Power Research Institute, Inc.
# All rights reserved.
#
# Redistribution and use in source and binary forms, with or without
# modification, are permitted provided that the following conditions are met:
#
# 1. Redistributions of source code must retain the above copyright notice,
#    this list of conditions and the following disclaimer.
#
# 2. Redistributions in binary form must reproduce the above copyright notice,
#    this list of conditions and the following disclaimer in the documentation
#    and/or other materials provided with the distribution.
#
# 3. Neither the name of the EPRI nor the names of its contributors may be used
#    to endorse or promote products derived from this software without specific
#    prior written permission.
#
# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
# AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
# IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE
# ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE
# LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR
# CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF
# SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS
# INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN
# CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE)
# ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE
# POSSIBILITY OF SUCH DAMAGE.

import logging

from matplotlib.backend_bases import MouseButton
from matplotlib.figure import Figure
import numpy
from numpy.typing import NDArray
from PySide6 import QtGui, QtWidgets

from . import util
from .. import export_script

logger = logging.getLogger(__name__)


def get_mean_std_dev_in_range(
        start_time: float, end_time: float, x: NDArray[numpy.float64],
        y: NDArray[numpy.float64]) -> tuple[float, float]:
    start_index: int = max(0, numpy.searchsorted(
        x, start_time, side="right").item() - 1)
    end_index: int = min(len(x), numpy.searchsorted(
        x, end_time, side="left").item() + 1)

    data = y[start_index:end_index]

    mean: float = numpy.mean(data, axis=0).item()
    std_dev: float = numpy.std(data, axis=0).item()

    return mean, std_dev


def display_mean_and_std_dev(
        start_time: float, end_time: float,
        data: list[tuple[NDArray[numpy.float64], NDArray[numpy.float64],
                         str, util.UnitInfo]],
        title: str, parent: QtWidgets.QWidget) -> None:

    value_strings: list[str] = []

    for x, y, name, unit in data:
        mean, std_dev = get_mean_std_dev_in_range(start_time, end_time, x, y)

        base_str = f"mean={mean} {unit.utf8}, std dev={std_dev} {unit.utf8}"

        if name:
            value_strings.append(f"{name}: {base_str}")
        else:
            value_strings.append(base_str)

    text = "\n".join(value_strings)

    QtWidgets.QMessageBox.information(parent, title, text)


def add_mean_std_dev_summary_buttons(fig: Figure, axes: list, data: list[
        tuple[NDArray[numpy.float64], NDArray[numpy.float64], str,
              util.UnitInfo]]) -> None:
    rects = {ax: None for ax in axes}  # used in the callback

    toolbar: QtWidgets.QToolBar = fig.canvas.toolbar  # type: ignore
    total_action_text = QtWidgets.QApplication.translate(
        "RMSTotalAction", "Mean and Std Dev (All)")
    total_action = QtGui.QAction(total_action_text, toolbar)
    total_action.setIcon(QtGui.QIcon.fromTheme("multimeter_analog"))

    range_action_text = QtWidgets.QApplication.translate(
        "RMSRangeAction", "Mean and Std Dev (Range)")
    range_action = QtGui.QAction(range_action_text, toolbar)
    range_action.setIcon(QtGui.QIcon.fromTheme("measuring_cup"))

    def measure_total() -> None:
        if not data:
            return  # no data

        start_time = min(numpy.min(d[0]).item() for d in data)
        end_time = max(numpy.max(d[0]).item() for d in data)

        display_mean_and_std_dev(start_time, end_time, data, total_action_text,
                                 toolbar)

    def measure_range() -> None:
        QtWidgets.QMessageBox.information(
            toolbar, "Select Measurement Range",
            "Select measurement range.\nUndo = middle click / backspace")

        # grab the points from the user
        points = fig.ginput(n=2, timeout=0, mouse_add=MouseButton.LEFT,
                            mouse_stop=None,  # type: ignore
                            mouse_pop=MouseButton.MIDDLE)
        if len(points) != 2:
            return  # user cancelled
        start_time = min(p[0] for p in points)
        end_time = max(p[0] for p in points)

        # show a box with the freq range
        for ax in axes:
            rect = rects[ax]
            if rect is not None:
                rect.remove()
                rects[ax] = None
                rect = None
            rect = ax.axvspan(start_time, end_time, facecolor='g', alpha=0.5)
            rects[ax] = rect
        fig.canvas.draw()

        display_mean_and_std_dev(start_time, end_time, data, range_action_text,
                                 toolbar)

    total_action.triggered.connect(measure_total)
    range_action.triggered.connect(measure_range)

    toolbar.addAction(total_action)
    toolbar.addAction(range_action)


def time_analysis(parent: QtWidgets.QWidget) -> None:
    ret = util.load_batch(parent)
    if ret is None:
        return  # error message displayed, or user cancelled
    file_infos, header = ret

    channel_index = util.choose_channel(header, parent)
    if channel_index is None:
        return  # error message displayed, or user cancelled

    batch_info = util.decode_batch(file_infos, header, [channel_index], parent)
    if batch_info is None:
        return  # error message displayed, or user cancelled

    batch_info = util.get_datetime_subset(batch_info, parent)
    if batch_info is None:
        return  # error message displayed, or user cancelled

    ret = util.warn_about_lost_packets(batch_info, parent)
    if not ret:
        return  # user cancelled

    sequence_info, unit_info = util.sequence_data(
        batch_info, parent, chunk=False)

    time, data, _start, _end = sequence_info[channel_index][0]

    show_legend = False

    if data.shape[1] != 1:
        subchannel_index = util.choose_subchannel(
            batch_info[channel_index], allow_all=True, parent=parent)
        if subchannel_index is None:
            return  # user cancelled
        elif subchannel_index != -1:
            data = data[:, subchannel_index:subchannel_index+1]
        else:
            show_legend = True

    # wrapper around matplotlib.pyplot.subplots()
    fig, ax = export_script.subplots_wrapper()

    ax.plot(time, data)

    channel_data = batch_info[channel_index]
    subchannel_names = channel_data.channel_decoder.subchannel_names

    mean_std_data: list[tuple[NDArray[numpy.float64], NDArray[numpy.float64],
                              str, util.UnitInfo]] = []
    for i, name in enumerate(subchannel_names):
        mean_std_data.append((time, data[:, i], name,
                              unit_info[channel_index]))

    ax.set_xlabel("Time (s)")
    unit_str = unit_info[channel_index].utf8
    if unit_str:
        ax.set_ylabel("Magnitude ({})".format(unit_str))
    else:
        ax.set_ylabel("Magnitude")

    ax.get_xaxis().get_major_formatter().set_useOffset(False)  # type: ignore
    ax.get_yaxis().get_major_formatter().set_useOffset(False)  # type: ignore

    if show_legend:
        ax.legend(subchannel_names)

    add_mean_std_dev_summary_buttons(fig, [ax], mean_std_data)

    fig.show()


def synchronous_time_analysis(parent: QtWidgets.QWidget) -> None:
    ret = util.load_batch(parent)
    if ret is None:
        return  # error message displayed, or user cancelled
    file_infos, header = ret

    channel_indexes = util.choose_synchronous_channels(header, parent)
    if channel_indexes is None:
        return  # error message displayed, or user cancelled

    batch_info = util.decode_batch(file_infos, header, channel_indexes, parent)
    if batch_info is None:
        return  # error message displayed, or user cancelled

    batch_info = util.get_datetime_subset(batch_info, parent)
    if batch_info is None:
        return  # error message displayed, or user cancelled

    ret = util.warn_about_lost_packets(batch_info, parent, once=True)
    if not ret:
        return  # user cancelled

    sequence_info, unit_info = util.sequence_data(
        batch_info, parent, chunk=False, unscaled_units=True)

    # group channels by unit type
    plot_groups: list[list[int]] = []
    groups_by_unit_type: dict[int, list[int]] = {}
    for channel_index in sorted(channel_indexes):
        channel = header['channels'][channel_index]
        if channel.unit_type == 0:
            # don't share plots between channels with unknown units
            plot_groups.append([channel_index])
        elif channel.unit_type in groups_by_unit_type:
            # already have a group made
            plot_group = groups_by_unit_type[channel.unit_type]
            plot_group.append(channel_index)
        else:
            # need a new group
            plot_group = [channel_index]
            plot_groups.append(plot_group)
            groups_by_unit_type[channel.unit_type] = plot_group

    # wrapper around matplotlib.pyplot.subplots()
    fig, ax_list = export_script.subplots_wrapper(len(plot_groups),
                                                  squeeze=False, sharex=True)

    mean_std_data: list[tuple[NDArray[numpy.float64], NDArray[numpy.float64],
                              str, util.UnitInfo]] = []

    for i, ax in enumerate(ax_list[:, 0]):
        plot_group = plot_groups[i]

        legend_strs = []

        for channel_index in plot_group:
            time, data, _start, _end = sequence_info[channel_index][0]
            ax.plot(time, data)

            channel_data = batch_info[channel_index]
            subchannel_names = channel_data.channel_decoder.subchannel_names
            legend_strs.extend(subchannel_names)

            for i, name in enumerate(subchannel_names):
                mean_std_data.append((time, data[:, i], name,
                                      unit_info[channel_index]))

        ax.set_xlabel("Time (s)")
        unit_str = unit_info[plot_group[0]].utf8
        if unit_str:
            ax.set_ylabel("Magnitude ({})".format(unit_str))
        else:
            ax.set_ylabel("Magnitude")

        ax.get_xaxis().get_major_formatter().set_useOffset(False)
        ax.get_yaxis().get_major_formatter().set_useOffset(False)

        ax.legend(legend_strs)

    add_mean_std_dev_summary_buttons(fig, list(ax_list[:, 0]), mean_std_data)

    fig.show()


def overlaid_time_analysis(parent: QtWidgets.QWidget) -> None:
    sequences = []
    subchannel_indexes = []
    unit_types = []
    unit_names = []
    unit_infos: list[util.UnitInfo] = []
    names = []

    while True:
        # load a batch of files
        ret = util.load_batch(parent)
        if ret is None:
            return  # error message displayed, or user cancelled
        file_infos, header = ret

        # choose a channel
        channel_index = util.choose_channel(header, parent)
        if channel_index is None:
            return  # error message displayed, or user cancelled

        batch_info = util.decode_batch(
            file_infos, header, [channel_index], parent)
        if batch_info is None:
            return  # error message displayed, or user cancelled

        subchannel_index = util.choose_subchannel(
            batch_info[channel_index], allow_all=False, parent=parent)
        if subchannel_index is None:
            return  # user cancelled
        subchannel_indexes.append(subchannel_index)

        batch_info = util.get_datetime_subset(batch_info, parent)
        if batch_info is None:
            return  # error message displayed, or user cancelled

        ret = util.warn_about_lost_packets(batch_info, parent)
        if not ret:
            return  # user cancelled

        sequence_info, unit_info = util.sequence_data(
            batch_info, parent, chunk=False, unscaled_units=True)

        sequence = sequence_info[channel_index][0]

        sequences.append(sequence)
        unit_types.append(header['channels'][channel_index].unit_type)
        unit_names.append(unit_info[channel_index].utf8)
        unit_infos.append(unit_info[channel_index])

        # ask the user for a name for this set
        msg = "Display name for this batch (#{}):".format(len(sequences))
        name, ok = QtWidgets.QInputDialog.getText(parent, "Batch Name", msg)
        if not ok:
            return  # user cancelled

        names.append(name.strip())

        # ask the user if they want to load more
        msg = "Load more batches? Loaded {} so far".format(len(sequences))
        ret = QtWidgets.QMessageBox.question(
            parent, "More Batches?", msg,
            buttons=(QtWidgets.QMessageBox.StandardButton.Yes |
                     QtWidgets.QMessageBox.StandardButton.No),
            defaultButton=QtWidgets.QMessageBox.StandardButton.No)
        if ret != QtWidgets.QMessageBox.StandardButton.Yes:
            break

    # group channels by unit type
    plot_groups: list[list[int]] = []
    groups_by_unit_type: dict[int, list[int]] = {}
    for batch_index, unit_type in enumerate(unit_types):
        if unit_type == 0:
            # don't share plots between channels with unknown units
            plot_groups.append([batch_index])
        elif unit_type in groups_by_unit_type:
            # already have a group made
            plot_group = groups_by_unit_type[unit_type]
            plot_group.append(batch_index)
        else:
            # need a new group
            plot_group = [batch_index]
            plot_groups.append(plot_group)
            groups_by_unit_type[unit_type] = plot_group

    # wrapper around matplotlib.pyplot.subplots()
    fig, ax_list = export_script.subplots_wrapper(len(plot_groups),
                                                  squeeze=False, sharex=True)

    mean_std_data: list[tuple[NDArray[numpy.float64], NDArray[numpy.float64],
                              str, util.UnitInfo]] = []

    for i, ax in enumerate(ax_list[:, 0]):
        plot_group = plot_groups[i]

        for batch_index in plot_group:
            time, data, _start, _end = sequences[batch_index]
            subchannel_index = subchannel_indexes[batch_index]
            ax.plot(time, data[:, subchannel_index], label=names[batch_index])

            mean_std_data.append((time, data[:, subchannel_index],
                                  names[batch_index], unit_infos[batch_index]))

        ax.legend()

        ax.set_xlabel("Time (s)")
        unit_str = unit_names[plot_group[0]]
        if unit_str:
            ax.set_ylabel("Magnitude ({})".format(unit_str))
        else:
            ax.set_ylabel("Magnitude")

        ax.get_xaxis().get_major_formatter().set_useOffset(False)
        ax.get_yaxis().get_major_formatter().set_useOffset(False)

    add_mean_std_dev_summary_buttons(fig, list(ax_list[:, 0]), mean_std_data)

    fig.show()
