from PySide6.QtCore import Signal
from PySide6.QtWidgets import QWidget, QVBoxLayout

import matplotlib

matplotlib.use("QtAgg")
from matplotlib.backends.backend_qtagg import FigureCanvasQTAgg
from matplotlib.figure import Figure

import numpy as np

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


class SkeletonViewWidget(QWidget):

    session_folder_loaded_signal = Signal()

    def __init__(self, plot_title: str):
        super().__init__()

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

        self.plot_title = plot_title
        self.fig, self.ax = self.initialize_skeleton_plot()
        self._layout.addWidget(self.fig)

        self.skeleton_loaded = False

        self.current_xlim = None
        self.current_ylim = None
        self.current_zlim = None

    def load_skeleton(self, skeleton_3d_data: np.ndarray, connections: list = mediapipe_model_info["body_connections"]):

        self.skeleton_3d_data = skeleton_3d_data
        self.skeleton = build_skeleton(
            skeleton_3d_frame_marker_xyz=self.skeleton_3d_data,
            pose_estimation_connections=connections,
        )
        self.reset_skeleton_3d_plot()

        self.skeleton_loaded = True

    def initialize_skeleton_plot(self):
        fig = Mpl3DPlotCanvas(self, width=5, height=4, dpi=100)
        ax = fig.figure.axes[0]

        ax.set_title(self.plot_title)
        return fig, ax

    def reset_skeleton_3d_plot(self):
        self.ax.cla()
        self.calculate_axes_means(self.skeleton_3d_data)
        self.skel_x, self.skel_y, self.skel_z = self.get_x_y_z_data(0)
        self.plot_skel(0, self.skel_x, self.skel_y, self.skel_z)

    def calculate_axes_means(self, skeleton_3d_data):
        self.mx_skel = np.nanmean(skeleton_3d_data[:, 0:33, 0])
        self.my_skel = np.nanmean(skeleton_3d_data[:, 0:33, 1])
        self.mz_skel = np.nanmean(skeleton_3d_data[:, 0:33, 2])
        self.skel_3d_range = 900

    def plot_skel(self, frame_number, skel_x, skel_y, skel_z):
        self.ax.scatter(skel_x, skel_y, skel_z)
        self.plot_skeleton_bones(frame_number)
        if self.current_xlim:
            self.ax.set_xlim([self.current_xlim[0], self.current_xlim[1]])
            self.ax.set_ylim([self.current_ylim[0], self.current_ylim[1]])
            self.ax.set_zlim([self.current_zlim[0], self.current_zlim[1]])
        else:
            self.ax.set_xlim(
                [self.mx_skel - self.skel_3d_range, self.mx_skel + self.skel_3d_range]
            )
            self.ax.set_ylim(
                [self.my_skel - self.skel_3d_range, self.my_skel + self.skel_3d_range]
            )
            self.ax.set_zlim(
                [self.mz_skel - self.skel_3d_range, self.mz_skel + self.skel_3d_range]
            )

        self.ax.set_title(self.plot_title)
        self.fig.figure.canvas.draw()

    def plot_skeleton_bones(self, frame_number):
        this_frame_skeleton_data = self.skeleton[frame_number]
        for connection in this_frame_skeleton_data:
            line_start_point = connection[0]
            line_end_point = connection[1]

            bone_x, bone_y, bone_z = (
                [line_start_point[0], line_end_point[0]],
                [line_start_point[1], line_end_point[1]],
                [line_start_point[2], line_end_point[2]],
            )

            self.ax.plot(bone_x, bone_y, bone_z)

    def get_x_y_z_data(self, frame_number: int):
        skel_x = self.skeleton_3d_data[frame_number, :, 0]
        skel_y = self.skeleton_3d_data[frame_number, :, 1]
        skel_z = self.skeleton_3d_data[frame_number, :, 2]

        return skel_x, skel_y, skel_z

    def replot(self, frame_number: int):
        skel_x, skel_y, skel_z = self.get_x_y_z_data(frame_number)
        self.current_xlim = self.ax.get_xlim()
        self.current_ylim = self.ax.get_ylim()
        self.current_zlim = self.ax.get_zlim()
        self.ax.cla()
        self.plot_skel(frame_number, skel_x, skel_y, skel_z)
        # self.label.setText(str(frame_number))


class Mpl3DPlotCanvas(FigureCanvasQTAgg):

    def __init__(self, parent=None, width=4, height=4, dpi=100):
        fig = Figure(figsize=(width, height), dpi=dpi)
        self.axes = fig.add_subplot(111, projection="3d")
        super(Mpl3DPlotCanvas, self).__init__(fig)
