# import IPython
import os

import matplotlib.animation as animation
import matplotlib.colors as colors
import matplotlib.patheffects as pe
import matplotlib.pyplot as plt
import numpy as np
import pandas as pd


def save_fig(fig_id, tight_layout=True):
    if tight_layout:
        plt.tight_layout()
    plt.savefig(fig_id + ".png", format="png", dpi=300)


def draw_stickfigure(mocap_track, frame, data=None, joints=None, draw_names=False, ax=None, figsize=(8, 8)):
    if ax is None:
        fig = plt.figure(figsize=figsize)
        ax = fig.add_subplot(111)

    if joints is None:
        joints_to_draw = mocap_track.skeleton.keys()
    else:
        joints_to_draw = joints

    if data is None:
        df = mocap_track.values
    else:
        df = data

    for joint in joints_to_draw:
        ax.scatter(
            x=df["%s_Xposition" % joint][frame], y=df["%s_Yposition" % joint][frame], alpha=0.6, c="b", marker="o"
        )

        parent_x = df["%s_Xposition" % joint][frame]
        parent_y = df["%s_Yposition" % joint][frame]

        children_to_draw = [c for c in mocap_track.skeleton[joint]["children"] if c in joints_to_draw]

        for c in children_to_draw:
            child_x = df["%s_Xposition" % c][frame]
            child_y = df["%s_Yposition" % c][frame]
            ax.plot([parent_x, child_x], [parent_y, child_y], "k-", lw=2)

        if draw_names:
            ax.annotate(joint, (df["%s_Xposition" % joint][frame] + 0.1, df["%s_Yposition" % joint][frame] + 0.1))

    return ax


def draw_stickfigure3d(mocap_track, frame, data=None, joints=None, draw_names=False, ax=None, figsize=(8, 8)):
    from mpl_toolkits.mplot3d import Axes3D

    if ax is None:
        fig = plt.figure(figsize=figsize)
        ax = fig.add_subplot(111, projection="3d")

    if joints is None:
        joints_to_draw = mocap_track.skeleton.keys()
    else:
        joints_to_draw = joints

    if data is None:
        df = mocap_track.values
    else:
        df = data

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

        ax.scatter(xs=parent_x, ys=parent_y, zs=parent_z, alpha=0.6, c="b", marker="o")

        children_to_draw = [c for c in mocap_track.skeleton[joint]["children"] if c in joints_to_draw]

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

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

        if draw_names:
            ax.text(x=parent_x + 0.1, y=parent_y + 0.1, z=parent_z + 0.1, s=joint, color="rgba(0,0,0,0.9")

    return ax


def sketch_move(mocap_track, data=None, ax=None, figsize=(16, 8)):
    if ax is None:
        fig = plt.figure(figsize=figsize)
        ax = fig.add_subplot(111)

    if data is None:
        data = mocap_track.values

    for frame in range(0, data.shape[0], 4):
        #         draw_stickfigure(mocap_track, f, data=data, ax=ax)

        for joint in mocap_track.skeleton.keys():
            children_to_draw = [c for c in mocap_track.skeleton[joint]["children"]]

            parent_x = data["%s_Xposition" % joint][frame]
            parent_y = data["%s_Yposition" % joint][frame]

            frame_alpha = frame / data.shape[0]

            for c in children_to_draw:
                child_x = data["%s_Xposition" % c][frame]
                child_y = data["%s_Yposition" % c][frame]

                ax.plot([parent_x, child_x], [parent_y, child_y], "-", lw=1, color="gray", alpha=frame_alpha)


def render_mp4(mocap_track, filename, data=None, ax=None, axis_scale=50, elev=45, azim=45):
    if ax is None:
        fig = plt.figure(figsize=(10, 10))
        ax = fig.add_subplot(111, projection="3d")
        ax.set_xlim3d(-axis_scale, axis_scale)
        ax.set_zlim3d(0, axis_scale)
        ax.set_ylim3d(-axis_scale, axis_scale)
        ax.grid(True)
        ax.set_axis_off()

        ax.view_init(elev=elev, azim=azim)

        xs = np.linspace(-200, 200, 50)
        ys = np.linspace(-200, 200, 50)
        X, Y = np.meshgrid(xs, ys)
        Z = np.zeros(X.shape)

        wframe = ax.plot_wireframe(X, Y, Z, rstride=2, cstride=2, color="grey", lw=0.2)

        # fig = plt.figure(figsize=figsize)
        # ax = fig.add_subplot(111)

    if data is None:
        data = mocap_track.values

    fps = int(np.round(1 / mocap_track.framerate))
    lines = []
    lines.append(
        [
            plt.plot(
                [0, 0],
                [0, 0],
                [0, 0],
                color="red",
                lw=2,
                path_effects=[pe.Stroke(linewidth=3, foreground="black"), pe.Normal()],
            )[0]
            for _ in range(len(mocap_track.skeleton.keys()))
        ]
    )

    def animate(frame):
        changed = []
        j = 0
        for joint in mocap_track.skeleton.keys():
            children_to_draw = [c for c in mocap_track.skeleton[joint]["children"]]

            parent_x = data["%s_Xposition" % joint][frame]
            parent_y = data["%s_Yposition" % joint][frame]
            parent_z = data["%s_Zposition" % joint][frame]

            # frame_alpha = frame/data.shape[0]

            for c in children_to_draw:
                child_x = data["%s_Xposition" % c][frame]
                child_y = data["%s_Yposition" % c][frame]
                child_z = data["%s_Zposition" % c][frame]

                lines[0][j].set_data(np.array([[child_x, parent_x], [-child_z, -parent_z]]))
                lines[0][j].set_3d_properties(np.array([child_y, parent_y]))

            changed += lines
            j += 1

        return changed

    plt.tight_layout()

    ani = animation.FuncAnimation(fig, animate, np.arange(data.shape[0]), interval=1000 / fps)

    if filename != None:
        ani.save(filename, fps=fps, bitrate=13934)
        ani.event_source.stop()
        del ani
        plt.close()
    try:
        plt.show()
        plt.save()
    except AttributeError as e:
        pass


def viz_cnn_filter(feature_to_viz, mocap_track, data, gap=25):
    fig = plt.figure(figsize=(16, 4))
    ax = plt.subplot2grid((1, 8), (0, 0))
    ax.imshow(feature_to_viz.T, aspect="auto", interpolation="nearest")

    ax = plt.subplot2grid((1, 8), (0, 1), colspan=7)
    for frame in range(feature_to_viz.shape[0]):
        frame_alpha = 0.2  # frame/data.shape[0] * 2 + 0.2

        for joint_i, joint in enumerate(mocap_track.skeleton.keys()):
            children_to_draw = [c for c in mocap_track.skeleton[joint]["children"]]

            parent_x = data["%s_Xposition" % joint][frame] + frame * gap
            parent_y = data["%s_Yposition" % joint][frame]

            ax.scatter(
                x=parent_x,
                y=parent_y,
                alpha=0.6,
                cmap="RdBu",
                c=feature_to_viz[frame][joint_i] * 10000,
                marker="o",
                s=abs(feature_to_viz[frame][joint_i] * 10000),
            )
            plt.axis("off")
            for c in children_to_draw:
                child_x = data["%s_Xposition" % c][frame] + frame * gap
                child_y = data["%s_Yposition" % c][frame]

                ax.plot([parent_x, child_x], [parent_y, child_y], "-", lw=1, color="gray", alpha=frame_alpha)


def print_skel(X):
    stack = [X.root_name]
    tab = 0
    while stack:
        joint = stack.pop()
        tab = len(stack)
        print("%s- %s (%s)" % ("| " * tab, joint, X.skeleton[joint]["parent"]))
        for c in X.skeleton[joint]["children"]:
            stack.append(c)


# def nb_play_mocap_fromurl(mocap, mf, frame_time=1/30, scale=1, base_url='http://titan:8385'):
# if mf == 'bvh':
# bw = BVHWriter()
# with open('test.bvh', 'w') as ofile:
# bw.write(mocap, ofile)

# filepath = '../notebooks/test.bvh'
# elif mf == 'pos':
# c = list(mocap.values.columns)

# for cc in c:
# if 'rotation' in cc:
# c.remove(cc)
# mocap.values.to_csv('test.csv', index=False, columns=c)

# filepath = '../notebooks/test.csv'
# else:
# return

# url = '%s/mocapplayer/player.html?data_url=%s&scale=%f&cz=200&order=xzyi&frame_time=%f'%(base_url, filepath, scale, frame_time)
# iframe = '<iframe src=' + url + ' width="100%" height=500></iframe>'
# link = '<a href=%s target="_blank">New Window</a>'%url
# return IPython.display.HTML(iframe+link)

# def nb_play_mocap(mocap, mf, meta=None, frame_time=1/30, scale=1, camera_z=500, base_url=None):
# data_template = 'var dataBuffer = `$$DATA$$`;'
# data_template += 'var metadata = $$META$$;'
# data_template += 'start(dataBuffer, metadata, $$CZ$$, $$SCALE$$, $$FRAMETIME$$);'
# dir_path = os.path.dirname(os.path.realpath(__file__))


# if base_url is None:
# base_url = os.path.join(dir_path, 'mocapplayer/playBuffer.html')

# # print(dir_path)

# if mf == 'bvh':
# pass
# elif mf == 'pos':
# cols = list(mocap.values.columns)
# for c in cols:
# if 'rotation' in c:
# cols.remove(c)

# data_csv = mocap.values.to_csv(index=False, columns=cols)

# if meta is not None:
# lines = [','.join(item) for item in meta.astype('str')]
# meta_csv = '[' + ','.join('[%s]'%l for l in lines) +']'
# else:
# meta_csv = '[]'

# data_assigned = data_template.replace('$$DATA$$', data_csv)
# data_assigned = data_assigned.replace('$$META$$', meta_csv)
# data_assigned = data_assigned.replace('$$CZ$$', str(camera_z))
# data_assigned = data_assigned.replace('$$SCALE$$', str(scale))
# data_assigned = data_assigned.replace('$$FRAMETIME$$', str(frame_time))

# else:
# return


# with open(os.path.join(dir_path, 'mocapplayer/data.js'), 'w') as oFile:
# oFile.write(data_assigned)

# url = '%s?&cz=200&order=xzyi&frame_time=%f&scale=%f'%(base_url, frame_time, scale)
# iframe = '<iframe frameborder="0" src=' + url + ' width="100%" height=500></iframe>'
# link = '<a href=%s target="_blank">New Window</a>'%url
# return IPython.display.HTML(iframe+link)