from PySide6.QtWidgets import QWidget, QVBoxLayout

import matplotlib

matplotlib.use("QtAgg")

from matplotlib.backends.backend_qtagg import (
    FigureCanvasQTAgg,
    NavigationToolbar2QT as NavigationToolbar,
)
from matplotlib.figure import Figure

from skellyforge.freemocap_utils.postprocessing_widgets.visualization_widgets.skeleton_builder import (
    mediapipe_model_info, get_index_by_name
)

import numpy as np


class TimeSeriesPlotCanvas(FigureCanvasQTAgg):

    def __init__(self, parent=None, width=15, height=4, dpi=100):
        fig = Figure(figsize=(width, height * 0.7), dpi=dpi)
        self.x_ax = fig.add_subplot(311)
        self.y_ax = fig.add_subplot(312)
        self.z_ax = fig.add_subplot(313)

        fig.subplots_adjust(hspace=0.4)
        fig.tight_layout(pad=2.5)

        super(TimeSeriesPlotCanvas, self).__init__(fig)


class TimeSeriesPlotterWidget(QWidget):

    def __init__(self):
        super().__init__()

        self._layout = QVBoxLayout()
        self.setLayout(self._layout)

        self.lines = []

        self.fig, self.axes_list = self.initialize_skeleton_plot()

        toolbar = NavigationToolbar(self.fig, self)
        self._layout.addWidget(toolbar)
        self._layout.addWidget(self.fig)

    def initialize_skeleton_plot(self):
        fig = TimeSeriesPlotCanvas(self, width=12, height=9, dpi=100)
        self.x_ax = fig.figure.axes[0]
        self.y_ax = fig.figure.axes[1]
        self.z_ax = fig.figure.axes[2]

        self.axes_list = [self.x_ax, self.y_ax, self.z_ax]
        return fig, self.axes_list

    def update_plot(
        self,
        marker_to_plot: str,
        original_freemocap_data: np.ndarray,
        processed_freemocap_data: np.ndarray,
        reset_axes=True,
        landmark_names: list = mediapipe_model_info["body_landmark_names"],
    ):
        axes_names = ["X Axis (mm)", "Y Axis (mm)", "Z Axis (mm)"]

        joint_index = get_index_by_name(name=marker_to_plot, landmark_names=landmark_names)

        for line in self.lines:
            line.remove()
        self.lines = []

        for dimension, (ax, ax_name) in enumerate(zip(self.axes_list, axes_names)):
            if reset_axes:
                ax.cla()

            (line1,) = ax.plot(
                original_freemocap_data[:, joint_index, dimension],
                label="Original Data",
                alpha=0.8,
                color="red",
            )
            self.lines.append(line1)
            if processed_freemocap_data is not None:
                (line2,) = ax.plot(
                    processed_freemocap_data[:, joint_index, dimension],
                    label="Processed Data",
                    alpha=0.6,
                    color="blue",
                )
                self.lines.append(line2)

            ax.set_ylabel(ax_name)

            if dimension == 2:  # put the xlabel only on the last plot
                ax.set_xlabel("Frame #")
            ax.legend()

        self.fig.figure.suptitle(f"{marker_to_plot} trajectory")

        self.fig.figure.canvas.draw_idle()
