import matplotlib
import matplotlib.pyplot as plt
from mpl_toolkits.mplot3d import Axes3D
import os
import sys
import time
from shutil import copyfile

import tensorflow as tf
import numpy as np
import pandas as pd
from tqdm import tqdm
import time
import datetime
module_path = os.path.abspath(os.path.join('../'))
if module_path not in sys.path:
    sys.path.append(module_path)
import pymo

from pymo import features
from groovenet_helpers import *
from audio_extractor import *
import nbstripout


def generate(sess, model, init_frames, audio, p):
    num_gen = len(audio)
    gen_sample = []
    
    gen_sample.extend(init_frames)
    
    print('Generating %d frames: '%(num_gen))

    for f in range(num_gen-1):
        frame_audio = [audio[f]]
                
        init_x = np.expand_dims(np.asarray(gen_sample)[-p.n_steps:,:], axis=0)

        y_pred = sess.run(model.outputs, feed_dict={model.M: np.asarray(init_x), model.A: frame_audio})

        gen_sample.append(y_pred[0])
    
    return np.asarray(gen_sample)



def save_trial(_gen_mocap, _in_audio, _gen_params, _params, model):
    name = _gen_params['name']
    timestamp = time.strftime('%Y%m%d-%H%M')
    
    parent_dir = os.path.join(model.trials_dir,'%s-ep%d'%(_params.name, _gen_params['gen_epoch']))
    
    trial_root_dir = os.path.join(parent_dir,'%s-%s'%(name,timestamp))
    out_dir = os.path.join(trial_root_dir, 'outputs')
    code_dir = os.path.join(trial_root_dir, 'code')
    model_dir = os.path.join(trial_root_dir, 'model')
    
    if not os.path.exists(parent_dir):
        os.mkdir(parent_dir)
        
    if not os.path.exists(trial_root_dir):
        os.mkdir(trial_root_dir)
    if not os.path.exists(out_dir):
        os.mkdir(out_dir)
    if not os.path.exists(code_dir):
        os.mkdir(code_dir)

    
  
    with open(os.path.join(trial_root_dir, 'params.json'), 'w') as outfile:
        json.dump(_params, outfile)
    
    with open(os.path.join(trial_root_dir, 'gen-info.json'), 'w') as outfile:
        json.dump(_gen_params, outfile)
        
    with open(os.path.join(parent_dir, 'params.json'), 'w') as outfile:
        json.dump(_params, outfile)
    
    with open(os.path.join(parent_dir, 'gen-info.json'), 'w') as outfile:
        json.dump(_gen_params, outfile)
        
    save_output(name, _gen_mocap[0], _in_audio['audio'], _in_audio['audio_samplerate'], 1/60, _gen_params, out_dir)
    
    ## code
#     copyfile(gen_params['checkpoint'], os.path.join(trials_dir, os.path.basename(gen_params['checkpoint'])))
    
    with open(os.path.join(trial_root_dir, 'pipes.pkl'), 'wb') as outfile:
        pickle.dump({'mocap_pipe': model.mocap_pipe,'model.audio_pipe': model.audio_pipe}, outfile);
    
    
    fig=plt.figure(figsize=(18,10))
    _=plt.subplot(211)
    plt.title('Input Audio')
    _=plt.plot(_in_audio['features_windowed'][:,-1,:])

    _=plt.subplot(212)
    plt.title('Movement (Generated)')
    _=plt.plot(_gen_mocap[0].values)
    
    plt.savefig((os.path.join(out_dir, 'plot.png')))
    
    
    
    inM = tf.identity(model.M, 'inM')
    inA = tf.identity(model.A, 'inA')
    outM = tf.identity(model.outputs, 'outputs')

    with tf.Session() as sess:       

        print('Loading the model checkpoint')
        model.saver.restore(sess, _gen_params['checkpoint'])

        tf.saved_model.simple_save(sess, model_dir, inputs={'inM': inM, 'inA': inA}, outputs={'outputs': outM})
        

        
def gen_some(chk_path, inputs, gen_epoch, p, model):
    
    chk_path = '%s-%d'%(chk_path, gen_epoch)
    outs = []
    for (a_file, start, length) in inputs:
        print('Input file: %s'%a_file)
        gen_params = {}
        gen_params['gen_epoch'] = gen_epoch
        gen_params['name'] = os.path.basename(a_file)
        gen_params['input_filename'] = a_file
        gen_params['start_second'] = start # in seconds
        gen_params['gen_length'] = length # in seconds
        gen_params['checkpoint'] = chk_path
        #     saver.restore(sess, saver.last_checkpoints[-1])

        print('Extracting audio features for the input audio file')
        in_audio = load_and_proc_audio_file(gen_params['input_filename'], gen_params['start_second'], gen_params['gen_length'], model.audio_pipe, p)
#         model.n_mocap_dim = 84

        ## Init type
        zero_init = np.zeros([p.n_steps, model.n_mocap_dim])
        rnd_init = np.random.rand(p.n_steps, model.n_mocap_dim) * 5 - 2.5
        X_init = rnd_init
        
#         in_audio['features_windowed'] = np.clip(in_audio['features_windowed'], -10, 10)
        
        ## Gen
        with tf.Session() as sess:
            print('Loading the model checkpoint')
            model.saver.restore(sess, gen_params['checkpoint'])

            start_time = time.time()
            gen_X = generate(sess, model, X_init, in_audio['features_windowed'], p)
            end_time = time.time()

            print('Generation took: %f seconds, %f frames/seconds'%(end_time-start_time, (gen_params['gen_length']*60)/(end_time-start_time)))


        gen_X_noinit = gen_X[p.n_steps-1:,:]
        gen_mocap = model.mocap_pipe.inverse_transform([gen_X_noinit])
        
    return gen_mocap[0]


####################

def axvlines(xs, ax=None, line_style='-', **plot_kwargs):
    """
    Draw vertical lines on plot
    :param xs: A scalar, list, or 1D array of horizontal offsets
    :param ax: The axis (or none to use gca)
    :param plot_kwargs: Keyword arguments to be passed to plot
    :return: The plot object corresponding to the lines.
    """
    if ax is None:
        ax = plt.gca()
    xs = np.array((xs, ) if np.isscalar(xs) else xs, copy=False)
    lims = ax.get_ylim()
    x_points = np.repeat(xs[:, None], repeats=3, axis=1).flatten()
    y_points = np.repeat(np.array(lims + (np.nan, ))[None, :], repeats=len(xs), axis=0).flatten()
    plot = ax.plot(x_points, y_points, line_style, scaley = False, **plot_kwargs)
    return plot


def output_gen(chk_path, gen_epoch, p, model, a_file, start, length, random_init=True, X_init=None):
    gen_params = {}
    gen_params['gen_epoch'] = gen_epoch
    gen_params['name'] = os.path.basename(a_file)
    gen_params['input_filename'] = a_file
    gen_params['start_second'] = start # in seconds
    gen_params['gen_length'] = length # in seconds
    gen_params['checkpoint'] = '%s-%d'%(chk_path, gen_epoch)
    #     saver.restore(sess, saver.last_checkpoints[-1])

    print('Extracting audio features for the input audio file')
    in_audio = load_and_proc_audio_file(gen_params['input_filename'], gen_params['start_second'], gen_params['gen_length'], model.audio_pipe, p)
#         model.n_mocap_dim = 84

    ## Init type
    zero_init = np.zeros([p.n_steps, model.n_mocap_dim])
    rnd_init = np.random.rand(p.n_steps, model.n_mocap_dim) * 5 - 2.5
    
    if X_init is None:
        if random_init:
            X_init = rnd_init
        else:
            X_init = zero_init

#         in_audio['features_windowed'] = np.clip(in_audio['features_windowed'], -10, 10)

    ## Gen
    with tf.Session() as sess:
        print('Loading the model checkpoint')
        model.saver.restore(sess, gen_params['checkpoint'])

        start_time = time.time()
        gen_X = generate(sess, model, X_init, in_audio['features_windowed'], p)
        end_time = time.time()

        print('Generation took: %f seconds, %f frames/seconds'%(end_time-start_time, (gen_params['gen_length']*60)/(end_time-start_time)))

    
    
    gen_X_noinit = gen_X[p.n_steps-1:,:]
        
    gen_mocap = model.mocap_pipe.inverse_transform([gen_X_noinit])
    save_trial(gen_mocap, in_audio, gen_params, p, model)
    return in_audio, gen_mocap[0], gen_X


def output_rhythm_sync(chk_path, gen_epoch, p, model):
    tests = [
        #(file, audio_start, audio_len,  beat_start, beat_len, beat_count)
       ('Style Machine - House 2', 'rhythm_sync_sm2.png', '../groovenet-data/audiofiles/sm-house-02.wav', 
        25, 20, 
        19, 56.3, 21),
        
        ('Style Machine - House 1', 'rhythm_sync_sm1.png', '../groovenet-data/audiofiles/sm-house-01.wav', 
        28, 20, 
        8, 56.3, 21),
        
        ('PanSonic', 'rhythm_sync_pansonic.png', '../groovenet-data/audiofiles/PanSonic-Wanyugo.wav',
         60, 20, 
         3, 72, 16),
        
        ('LM100', 'rhythm_sync_lm100.png', '../groovenet-data/audiofiles/LM100.mp3',
         64, 20, 
         47, 68, 17),
    ]
    
    figs = []
    
    
    pad = 10
    for (title, outfile, a_file, audio_start, audio_len,  bstart, blen, bcount) in tests:

        si = (pad) * 60
        ei = si + (audio_len) * 60

        in_audio, mov, _ = output_gen(chk_path, gen_epoch, p, model, a_file, audio_start - pad, audio_len + 10, random_init=True)
                
        ref_move_chl = mov.values['Hips_Yposition'].values[si:ei]

        beats = [bstart+blen*i for i in range(bcount)]
        
        plt.rc('axes', labelsize=18)
        plt.rc('figure', titlesize=22)
        plt.rc('xtick', labelsize=15) 
        plt.rc('ytick', labelsize=14)
        plt.rcParams['axes.edgecolor'] = '#202020'
        
        fig, axs = plt.subplots(2, sharex=False, sharey=False, gridspec_kw={'hspace': 0}, figsize=(15,7), dpi=180)
        plt.suptitle(title, y = 0.97)
        plt.xlabel('Frames - at 60fps')
        
        axs[0].plot(in_audio['audio'][44100*(pad+1):44100*(pad+1) + 44100*audio_len], linewidth=0.75, color='#1987d1')
        axs[0].set_ylabel('audio')
        axs[0].set_yticks(())
        axs[0].yaxis.set_label_coords(-0.075, 0.5)

        axs[1].plot(ref_move_chl, color='#cc3014')
        axs[1].set_ylabel('Hip height (cm)')
        axs[1].yaxis.set_label_coords(-0.075, 0.5)
        

        axvlines([b * ((44100/60)) for b in beats], ax=axs[0], linewidth=1, linestyle='dashed', alpha=0.95, color='#202020')
        axvlines(beats, ax=axs[1], linewidth=1, linestyle='dashed', alpha=0.95, color='#202020')
        
        fig.savefig('siggraph_outputs/%s'%outfile, bbox_inches='tight', pad_inches=0.5)
        
        figs.append(fig)
    
    return figs


def sketch_move(mocap_track, data=None, step = 4, 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], step):
#         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=0.7, color='gray', alpha=0.9)

    
    return fig



def sketch_move_sample(mocap_track, data=None, step = 4, 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], step):
#         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=0.7, color='gray', alpha=0.9)

    
    return fig




def output_init(chk_path, gen_epoch, p, model, inits):
    
    outs = []
    num_plots = 2+len(inits)
    
    audio_file = '../groovenet-data/audiofiles/sm-house-01.wav'
    audio_start = 40
    audio_len = 30
    
    for i in range(num_plots):
       
       
        if i == 0: # zero
            _, mov, g = output_gen(chk_path, gen_epoch, p, model, audio_file, audio_start, audio_len, random_init=False)
        if i == 1: # random
            _, mov, g = output_gen(chk_path, gen_epoch, p, model, audio_file, audio_start, audio_len, random_init=True)
        if i > 1: # training data
            X_init = inits[i-2]

            _, mov, g = output_gen(chk_path, gen_epoch, p, model, audio_file, audio_start, audio_len, X_init=X_init)
    
        outs.append(g)
        
    return outs
    
def output_init_fig(outs):
    num_plots = len(outs)
    
    plt.rc('axes', labelsize=18)
    plt.rc('figure', titlesize=22)
    plt.rc('xtick', labelsize=15) 
    plt.rc('ytick', labelsize=14) 
        
    plt.rcParams['axes.edgecolor'] = '#202020'
    
    fig, axs = plt.subplots(num_plots, sharex=False, sharey=False, gridspec_kw={'hspace': 0}, figsize=(15,10), dpi=180)
#     plt.suptitle('', y = 0.94)
#     plt.set_cmap('hsv')
    for i in range(num_plots):
        axs[i].set_ylim(-5.,5.)

        color_idx = np.linspace(0.99, 0.2, outs[i].shape[1])
        for ci in range(outs[i].shape[1]):
            axs[i].plot(outs[i][:, ci], linewidth=0.75, alpha=0.8, color=plt.cm.gist_heat(color_idx[ci])) #plt.cm.gist_heat(color_idx[ci])
        
#         axs[i].yaxis.set_label_coords(-0.075, 0.5)
        if i == 0:
            axs[i].set_ylabel('zero')
        elif i == 1:  
            axs[i].set_ylabel('random')
        else:
            axs[i].set_ylabel('training')
            
    plt.xlabel('Frames - at 60fps')
    
    
   
    fig.savefig('siggraph_outputs/output_init_effect.png',bbox_inches='tight', pad_inches=0.25)

    return fig


def output_songs_gen(chk_path, gen_epoch, p, model):
    tests = [
        #(file, audio_start, audio_len,  beat_start, beat_len, beat_count)
       ('Style Machine - House 1', 'output_gen_heard_sm1.png', '../groovenet-data/audiofiles/sm-house-01.wav', 
        25, 20),
        
        ('Style Machine - House 2', 'output_gen_heard_sm2.png', '../groovenet-data/audiofiles/sm-house-02.wav', 
        25, 20),
        
        ('PanSonic-Wanyugo', 'output_gen_unheard_pansonic.png', '../groovenet-data/audiofiles/PanSonic-Wanyugo.wav',
         60, 20),
        
        
        ('LM100', 'output_gen_unheard_lm100.png', '../groovenet-data/audiofiles/LM100.mp3',
         5, 20),
        
        ('Les Jours Tristes', 'output_gen_unheard_lesjours.png', '../groovenet-data/audiofiles/Les Jours Tristes.mp3',
         10, 20),
        
        ('Schubert - Serenade', 'output_gen_unheard_schubert-serenade.png', '../groovenet-data/audiofiles/schubert-serenade.mp3',
         10, 20),
    ]
    
    outs = []
    
    for (title, outfile, a_file, audio_start, audio_len) in tests:

     

        in_audio, mov, _ = output_gen(chk_path, gen_epoch, p, model, a_file, audio_start, audio_len, random_init=True)
        
       
        outs.append((title, outfile, mov))
        
        
    
    return outs


def output_songs_viz(outs, frame_keys):
    assert(len(frame_keys) == len(outs))
    
    figs = []
    
    
    for (i, (title, outfile, mov)) in enumerate(outs):
        num_frames = len(frame_keys[i])
        
        frames = frame_keys[i]
        
        plt.rcParams['axes.edgecolor'] = 'w'
        
        fig = plt.figure(figsize=(num_frames*8,8), dpi=300)
        #plt.suptitle(title, y = 0.97)
        
        
        
        for f in range(num_frames):
            ax = fig.add_subplot(1, num_frames, f+1, projection='3d')
            output_pose3d(mov, frames[f], ax=ax, figsize=(8,8))
            
        fig.savefig('siggraph_outputs/%s'%outfile, bbox_inches='tight', pad_inches=0.25)
        
        figs.append(fig)
    
    return figs
            
            
            
def output_pose3d(mocap_track, frame, data=None, joints=None, draw_names=False, ax=None, figsize=(8,8), eq_aspec=True):
    from mpl_toolkits.mplot3d import Axes3D
    
    if ax is None:
        fig = plt.figure(figsize=figsize)
        ax = fig.gca(projection='3d')
        
#         ax.set_aspect('equal')
        
        plt.tick_params(axis='x',which='both',bottom=False,top=False,labelbottom=False)
        plt.tick_params(axis='y',which='both',bottom=False,top=False,labelbottom=False)
        plt.tick_params(axis='z',which='both',bottom=False,top=False,labelbottom=False)
            
    
    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='#771014', marker='o')

        ax.set_xticks([])
        ax.set_yticks([])
        ax.set_zticks([])
        
        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=3, c='#881014')
            
        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')
    
    if eq_aspec:
        set_axes_equal(ax)
    
    
    ax.xaxis.set_pane_color((1.0, 1.0, 1.0, 1.0))
    ax.yaxis.set_pane_color((1.0, 1.0, 1.0, 1.0))
    
    ax.xaxis.pane.set_edgecolor('w')
    ax.yaxis.pane.set_edgecolor('w')
    
    
#     ax.xaxis.set_edgecolor('w')
#     ax.yaxis.set_edgecolor('w')
    
    
    ax.grid(False)
        
    return ax
