from matplotlib import pyplot as plt
from matplotlib.widgets import Slider
import numpy as np

from resources.pymo.pymo.preprocessing import MocapParameterizer


class BVHAnimation:
    def __init__(self):
        self.bvh_data = None

    def animate(self):
        BVHAnimation.init_plot(BVHAnimation.positions(self.bvh_data)[0])

    @staticmethod
    def positions(bvh_data):
        # convert rotations to positions for each joint
        mp = MocapParameterizer('position')
        return mp.fit_transform([bvh_data])

    @staticmethod
    def init_plot(positions):
        fig = plt.figure(figsize=(8, 8))
        ax = fig.add_subplot(111, projection='3d')

        axcolor = 'lightgoldenrodyellow'
        axframe = plt.axes([0.25, 0.1, 0.65, 0.03], facecolor=axcolor)

        bvh_number_frames = positions.values.shape[0]
        sframe = Slider(axframe, 'Frame', 1, bvh_number_frames, valinit=1, valstep=1)

        joints_to_draw = positions.skeleton.keys()
        lines = []
        points = []

        df = positions.values

        for joint in joints_to_draw:
            parent_x = df['%s_Xposition' % joint][0]
            parent_y = df['%s_Yposition' % joint][0]
            parent_z = df['%s_Zposition' % joint][0]
            # ^ In mocaps, Y is the up-right axis

            points.append(ax.plot([parent_x],
                                  [parent_y],
                                  [parent_z],
                                  linestyle="", c='b', marker='o'))

            children_to_draw = [c for c in positions.skeleton[joint]['children'] if c in joints_to_draw]

            for c in children_to_draw:
                child_x = df['%s_Xposition' % c][0]
                child_y = df['%s_Yposition' % c][0]
                child_z = df['%s_Zposition' % c][0]
                # ^ In mocaps, Y is the up-right axis

                lines.append(
                    ax.plot([parent_x, child_x], [parent_y, child_y], [parent_z, child_z], 'k-', lw=2, c='black'))

        def update(_):
            """Update function for the slider"""
            frame = int(sframe.val) - 1
            index1 = 0
            index2 = 0

            for up_joint in joints_to_draw:
                up_parent_x = df['%s_Xposition' % up_joint][frame]
                up_parent_y = df['%s_Yposition' % up_joint][frame]
                up_parent_z = df['%s_Zposition' % up_joint][frame]
                # ^ In mocaps, Y is the up-right axis

                points[index1][0].set_data(np.array([up_parent_x]), np.array([up_parent_y]))
                points[index1][0].set_3d_properties([up_parent_z])
                index1 += 1

                up_children_to_draw = [c for c in positions.skeleton[up_joint]['children'] if c in joints_to_draw]

                for c in up_children_to_draw:
                    up_child_x = df['%s_Xposition' % c][frame]
                    up_child_y = df['%s_Yposition' % c][frame]
                    up_child_z = df['%s_Zposition' % c][frame]
                    # ^ In mocaps, Y is the up-right axis

                    lines[index2][0].set_data(np.array([[up_parent_x, up_child_x], [up_parent_y, up_child_y]]))
                    lines[index2][0].set_3d_properties([up_parent_z, up_child_z])
                    index2 += 1

            fig.canvas.draw_idle()

        # update the plot when changing the slider value
        sframe.on_changed(update)

        # ax.set_xlim3d([-200.0, 200.0])
        ax.set_xlabel('X')

        # ax.set_ylim3d([-200.0, 200.0])
        ax.set_ylabel('Y')

        # ax.set_zlim3d([-200.0, 200.0])
        ax.set_zlabel('Z')

        ax.view_init(elev=135, azim=-90)

        plt.show()


bvh_animation = BVHAnimation()