import pickle
import scipy.io.wavfile
from sklearn.pipeline import Pipeline
from sklearn.preprocessing import StandardScaler
import math

import os
import sys
module_path = os.path.abspath(os.path.join('../../PyMO/'))
module_path = os.path.abspath(os.path.join('../'))
if module_path not in sys.path:
    sys.path.append(module_path)

from essentia.standard import MonoWriter

from pymo.parsers import BVHParser
from pymo.preprocessing import *
from pymo.viz_tools import *
from pymo.writers import *
# from pymo.features import *

from audio_extractor import *

GN_DATA_ROOT = '../groovenet-data/'

def idx_to_data(X, idx):    
    return np.asarray([X[fi][s:e] for fi, s, e in idx])

# from https://goodcode.io/articles/python-dict-object/
class objdict(dict):
    def __getattr__(self, name):
        if name in self:
            return self[name]
        else:
            raise AttributeError("No such attribute: " + name)

    def __setattr__(self, name, value):
        self[name] = value

    def __delattr__(self, name):
        if name in self:
            del self[name]
        else:
            raise AttributeError("No such attribute: " + name)
            

def get_wave_segment(audio_track, start_i, length, audio_fr, target_fr):
    coef = audio_fr/target_fr
    segment_s = int(math.floor(start_i - 1) * coef)
    segment_e = int(segment_s + math.floor(length * coef))
#     print(segment_s)
#     print(segment_e)
    
    return audio_track[segment_s:segment_e]

def write_mov_aud_segments(mocap, audio_track, start_i, audio_fr, target_fr, name, out_dir):
    length = mocap.values.shape[0]
    
    a = get_wave_segment(audio_track, start_i, length, sr, audio_fr)
    scipy.io.wavfile.write(os.path.join(out_dir,'%s.wav'%name), sr, a)
    
    cols = list(mocap.values.columns)
    data_csv = mocap.values.to_csv(index=False, columns=cols)
    
    with open(os.path.join(out_dir,'%s.csv'%name), 'w') as outf:
        outf.write(data_csv)
        
    return (os.path.join(out_dir,'%s.wav'%name), os.path.join(out_dir,'%s.csv'%name))

def nb_play_groove(mocap, audio, audio_rate, frame_time=1/30, scale=1, camera_z=500, out_dir='./'):
    
    name = 'rnn_out'
    mp, ap = os.path.join(out_dir,'%s.csv'%name), os.path.join(out_dir,'%s.wav'%name)
    # scipy.io.wavfile.write(ap, audio_rate, audio)
    MonoWriter(filename=ap, sampleRate=audio_rate, format='wav')(audio)
    
    cols = list(mocap.values.columns)
    data_csv = mocap.values.to_csv(index=False, columns=cols)
    
    with open(mp, 'w') as outf:
        outf.write(data_csv)
    
    

    url = '/files/players/groovenet/groovenet_player.html?mov=outs/%s&aud=outs/%s'%('%s.csv'%name, '%s.wav'%name)
    url += '&movfmt=csv&frame_time=%f&scale=%f&cz=%d'%(frame_time, scale, camera_z)
    
    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)
    
def nb_play_mocap_(mocap, frame_time=1/30, scale=1, camera_z=500, out_dir='./'):
    
    name = 'rnn_out'
    mp, ap = os.path.join(out_dir,'%s.csv'%name), os.path.join(out_dir,'%s.wav'%name)
    
    cols = list(mocap.values.columns)
    data_csv = mocap.values.to_csv(index=False, columns=cols)
    
    with open(mp, 'w') as outf:
        outf.write(data_csv)

    url = '/files/players/groovenet/groovenet_player.html?mov=outs/%s&aud=outs/%s'%('%s.csv'%name, '%s.wav'%name)
    url += '&movfmt=csv&frame_time=%f&scale=%f&cz=%d'%(frame_time, scale, camera_z)
    
    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)

def save_output(name, mocap, audio, audio_rate, frame_time, gen_params, out_dir):    
    mp, ap = os.path.join(out_dir,'mov.csv'), os.path.join(out_dir,'music.wav')
    # scipy.io.wavfile.write(ap, audio_rate, audio)
    MonoWriter(filename=ap, sampleRate=audio_rate, format='wav')(audio)
    
    cols = list(mocap.values.columns)
    data_csv = mocap.values.to_csv(index=False, columns=cols)
    
    with open(mp, 'w') as outf:
        outf.write(data_csv)

def save_output_bvh(name, mocap, audio, audio_rate, frame_time, gen_params, out_dir):    
    mp, ap = os.path.join(out_dir,'mov.bvh'), os.path.join(out_dir,'music.wav')
    # scipy.io.wavfile.write(ap, audio_rate, audio)
    MonoWriter(filename=ap, sampleRate=audio_rate, format='wav')(audio)
    
    from pymo.writers import BVHWriter

    bw = BVHWriter()

    with open(mp, 'w') as ofile:
        bw.write(mocap, ofile)
        

def save_outputs_(mocap, audio, exp_name, params, arch_code, gen_code, audio_rate, frame_time=1/30, out_dir='./'):
    
    name = exp_name
    
    os.mkdir(os.path.join(out_dir, name))
    
    out_dir = os.path.join(out_dir, name)
    
    mp = os.path.join(out_dir,'%s.csv'%name)
    ap = os.path.join(out_dir,'%s.wav'%name)
    pp = os.path.join(out_dir,'%s.json'%name)
            
    scipy.io.wavfile.write(ap, audio_rate, audio)
    
    cols = list(mocap.values.columns)
    data_csv = mocap.values.to_csv(index=False, columns=cols)
    
    with open(mp, 'w') as outf:
        outf.write(data_csv)
        
    metadata = {'params': params,
               'arch_code': arch_code,
               'gen_code': gen_code}
    
    with open(pp, 'w') as outf:
        json.dump(metadata, outf)
        

def gn_load_data(pickle_file):
    with open(pickle_file, 'rb') as inf:
        saved_data = pickle.load(inf)

    file_list = saved_data['file_list']
    # euler_data = saved_data['euler_data']
    pos_data = saved_data['pos_data']
    audio_ess_train = saved_data['audio_ess_train']
    audio_ess_test = saved_data['audio_ess_test']
    
    return file_list, pos_data, audio_ess_train, audio_ess_test

def gn_load_waves():
    wav_data_train = []
    wav_data_train.append(scipy.io.wavfile.read('../groovenet-data/audiofiles/sm-house-01_16.wav'))
    wav_data_train.append(scipy.io.wavfile.read('../groovenet-data/audiofiles/sm-house-01_16.wav'))
    wav_data_train.append(scipy.io.wavfile.read('../groovenet-data/audiofiles/sm-house-02_16.wav'))
    wav_data_train.append(scipy.io.wavfile.read('../groovenet-data/audiofiles/sm-house-03_16.wav'))

    wav_data_test = {}
    wav_data_test['Christian'] = scipy.io.wavfile.read('../groovenet-data/audiofiles/ChristianLoeffler-Athlete.wav')
    wav_data_test['PanSonic'] = scipy.io.wavfile.read('../groovenet-data/audiofiles/PanSonic-Wanyugo.wav')
    wav_data_test['Pender'] = scipy.io.wavfile.read('../groovenet-data/audiofiles/PenderStreetSteppers-BubbleWorld.wav')
    wav_data_test['SM4'] = scipy.io.wavfile.read('../groovenet-data/audiofiles/sm-house-04_16.wav')
    wav_data_test['Grant'] = scipy.io.wavfile.read('../groovenet-data/audiofiles/Grant-Bend.wav')
    
    return wav_data_train, wav_data_test


def gn_preprocess(pos_data, audio_ess_train):
    train_X = pos_data
    train_aud = audio_ess_train

    #######################
    mocap_pipe = Pipeline([
        ('rcpn', RootCentricPositionNormalizer()),
        ('delta', RootTransformer('abdolute_translation_deltas')),
        ('const', ConstantsRemover()),
        ('np', Numpyfier()),
    #     ('down', DownSampler(2)),
        ('stdscale', ListStandardScaler())
    ])

    audio_pipe = Pipeline([
        ('stdscale', ListStandardScaler())
    ])

    piped_train_X = mocap_pipe.fit_transform(train_X)
    piped_train_A = audio_pipe.fit_transform(train_aud)

    
    return piped_train_X, piped_train_A, mocap_pipe, audio_pipe



def gn_make_seqs(piped_train_X, piped_train_A, p):
    train_data_idx = []
    train_predictions_idx = []
    train_labels_steps = []

    if not hasattr(p, 'n_steps'):
        p.n_steps = p.input_len
        p.window_gap = p.shift_len



    for fileindex, m in enumerate(piped_train_X):
        mm = m
        for i in range(0, len(mm)-p.n_steps-1, p.window_gap):
            train_data_idx.append([fileindex, i,i+p.n_steps])
            train_predictions_idx.append([fileindex, i+1,i+p.n_steps+1])


   
    train_data_idx = np.asarray(train_data_idx)

    train_predictions_idx = np.asarray(train_predictions_idx)
    
    return train_data_idx, train_predictions_idx




def gn_extract_subseq(seq_index, st_in, length, data_indexes, piped_X, piped_A, wav_data):
    train_seq = seq_index
    train_ids = data_indexes[data_indexes[:,0]==seq_index]
    
    X_sub = idx_to_data(piped_X, train_ids[st_in:st_in+length])
    A_sub = idx_to_data(piped_A, train_ids[st_in:st_in+length])
    aud_rate, gen_wav = wav_data[seq_index]
    W_sub = get_wave_segment(gen_wav, st_in, length, aud_rate, target_fr=60)
    
    return A_sub, W_sub, X_sub, aud_rate




def load_and_proc_audio_file(input_filename, start_second, gen_length, audio_pipe, p):
    in_audio = extract_audio_segment_jules(input_filename, start_second, gen_length)
    
    in_audio_feat = audio_pipe.transform([in_audio['features']])[0]
    in_aud_feat_w = []
    for i in range(0, in_audio_feat.shape[0]-p.n_steps):
        in_aud_feat_w.append(in_audio_feat[i:i+p.n_steps])

    in_audio['features_windowed'] = np.array(in_aud_feat_w)

    return in_audio


def mirror_mocap_side(track):
    new_track = track.clone()
    new_df = pd.DataFrame(index=track.values.index)

    for c in track.values.columns:
        if "Xposition" in c:
            new_df[c] = pd.Series(data=-track.values[c], index=new_df.index)
        else:
            new_df[c] = pd.Series(data=track.values[c], index=new_df.index)


    new_track.values = new_df
    
    return new_track



def gn_train_test_split(ratio, input_tracks, min_len):
    train_tracks = []
    test_tracks = []
    
    for track in input_tracks:
        size = track.shape[0]
        test_len = int(size * ratio)
        
        assert(size - test_len - min_len > min_len)
        
        s = np.random.randint(min_len, size - test_len - min_len)
        
        train_tracks.append(track[0:s,:])
        test_tracks.append(track[s:s+test_len,:])
        train_tracks.append(track[s+test_len:,:])
    
    return np.asarray(train_tracks), np.asarray(test_tracks)


def gn_make_seqs2(input_tracks, target_tracks, p):
    input_tensors = []
    target_tensors = []


    for t in range(input_tracks.shape[0]):
        input_track = input_tracks[t]
        target_track = target_tracks[t]
        
        last_index = input_track.shape[0]-p.input_len-p.target_len+p.target_overlap
        
        for i in range(0, last_index, p.shift_len):
            tg_start = i+p.input_len-p.target_overlap
            
            input_tensors.append(input_track[i:i+p.input_len,:])
            target_tensors.append(target_track[tg_start:tg_start+p.target_len ,:])
           


    
    return np.asarray(input_tensors), np.asarray(target_tensors)


class EssDescriptorSel(BaseEstimator, TransformerMixin):
    def __init__(self, keys):
        self.keys = keys
        
    
    def fit(self, X, y=None):    

        return self
    
    def transform(self, X, y=None):
        Q = []
        for t in X:
            feats = np.concatenate([t[k] for k in self.keys], axis=1)
            Q.append(feats)
    
        return Q
        
        
    def inverse_transform(self, X, copy=None):
        return X




def gn_load_data2(mirror_mocap=False, data_version='v1_3'):
    pickle_path = os.path.join(GN_DATA_ROOT, 'groovenet_essall+mocap%s_%s.pkl'%('+mirror' if mirror_mocap else '', data_version))

    if os.path.exists(pickle_path):
         with open(pickle_path, 'rb') as _pp:
            saved_data = pickle.load(_pp)

            return saved_data['pos_data'], saved_data['audio_ess'], saved_data['file_list'], saved_data['aud_feat_keys']

    else:
        sm1_all = extract_wrapper_o0('../groovenet-data/sm-house-01.wav', DEF_PARAMS)['all']
        sm2_all = extract_wrapper_o0('../groovenet-data/sm-house-02.wav', DEF_PARAMS)['all']
        sm3_all = extract_wrapper_o0('../groovenet-data/sm-house-03.wav', DEF_PARAMS)['all']


        file_list, pos_data, _, _ = gn_load_data('../groovenet-data/groovenet_mocap+ess_v1_2.pkl')

        audio_ess = [sm1_all, sm1_all, sm2_all, sm3_all]

        names = ['sm-house1a', 'sm-house1b', 'sm-house2', 'sm-house3']
        if mirror_mocap:
            
            aug_pos_data = []
            aug_ess_data = []
            aug_names = []
            aug_pos_data.extend(pos_data)
            aug_ess_data.extend(audio_ess)
            aug_names.extend(names)

            for i in range(4):
                aug_pos_data.append(mirror_mocap_side(pos_data[i]))
                aug_ess_data.append(audio_ess[i])
                aug_names.append(names[i])

            pos_data = aug_pos_data
            audio_ess = aug_ess_data
            names = aug_names

        save_data = {
            'pos_data': pos_data,
            'audio_ess': audio_ess,
            'file_list': names,
            'aud_feat_keys': list(sm1_all.keys())
        }

        with open(pickle_path, 'wb') as _pp:
            pickle.dump(save_data, _pp)


        return pos_data, audio_ess, names, sm1_all.keys()



"""
Loads the groovenet data
valid params are: 'pos', 'exp', and 'eul'
mirror_mocap is only available for 'pos'
"""
def gn_load_data3(param, mirror_mocap=False, data_version='v1_4'):
    
    pickle_path = os.path.join(GN_DATA_ROOT, 'groovenet_essall+%s_mocap%s_%s.pkl'%(param, '+mirror' if mirror_mocap else '', data_version))
    
    if os.path.exists(pickle_path):
        print('Loading the data from %s'%pickle_path)
        with open(pickle_path, 'rb') as _pp:
            saved_data = pickle.load(_pp)
    else:
        assert("what?")
        
    
    if param == 'exp':
        return saved_data['exp_data'], saved_data['audio_ess'], saved_data['file_list'], saved_data['aud_feat_keys']
    elif param == 'eul':        
        return saved_data['euler_data'], saved_data['audio_ess'], saved_data['file_list'], saved_data['aud_feat_keys']
    elif param == 'pos':        
        return saved_data['pos_data'], saved_data['audio_ess'], saved_data['file_list'], saved_data['aud_feat_keys']
    elif param == 'exp_euroot':
        return saved_data['exp_euroot_data'], saved_data['audio_ess'], saved_data['file_list'], saved_data['aud_feat_keys']
    else:
        raise('Valid param types are "exp", "eul", and "pos"')
        

class TaylorRootRep(BaseEstimator, TransformerMixin):  
    def __init__(self):       
        
        self.start_face = 0
        self.start_x = 0
        self.start_z = 0


    def fit(self, X, y=None):
        return self

    def transform(self, X, y=None):
        import matlab.engine
        eng = matlab.engine.start_matlab()
        Q = []

        for track in X:
            new_track = track.clone()
            new_df = track.values.copy()
            
            org_root = track.values[['Hips_Xposition', 'Hips_Yposition', 'Hips_Zposition', 'Hips_alpha', 'Hips_beta', 'Hips_gamma']].values
                    
            new_df.drop(['Hips_Xposition', 'Hips_Zposition', 'Hips_alpha', 'Hips_beta', 'Hips_gamma'], axis=1, inplace=True)


            eng = matlab.engine.start_matlab()
            ty_rep, start_face = eng.ty_pre(matlab.double(org_root.tolist()), nargout=2)
            self.start_face = start_face

            ty_rep = np.asarray(ty_rep)

            new_df['root_rot_a'] = pd.Series(data=ty_rep[:,0], index=new_df.index)
            new_df['root_rot_b'] = pd.Series(data=ty_rep[:,1], index=new_df.index)

            new_df['root_rot_vert_delta'] = pd.Series(data=ty_rep[:,2], index=new_df.index)

            new_df['root_delta_x'] = pd.Series(data=ty_rep[:,3], index=new_df.index)

            new_df['root_delta_z'] = pd.Series(data=ty_rep[:,4], index=new_df.index)
            

            new_track.values = new_df

            Q.append(new_track)
        
        return Q

    def inverse_transform(self, X, copy=None):
        import matlab.engine
        eng = matlab.engine.start_matlab()
        Q = []

        for track in X:
            new_track = track.clone()
            new_df = track.values.copy()

            ty_rep = track.values[['root_rot_a', 'root_rot_b', 'root_rot_vert_delta', 'root_delta_x', 'root_delta_z', 'Hips_Yposition']].values
            
            new_df.drop(['root_rot_a', 'root_rot_b', 'root_rot_vert_delta', 'root_delta_x', 'root_delta_z'], axis=1, inplace=True)
           

            backrep = eng.ty_post(matlab.double(ty_rep.tolist()), self.start_face, self.start_x, self.start_z, nargout=1)

            backrep = np.asarray(backrep)
        
            new_df['Hips_Xposition'] = pd.Series(data=backrep[:,0], index=new_df.index)
            new_df['Hips_Zposition'] = pd.Series(data=backrep[:,2], index=new_df.index)

            new_df['Hips_alpha'] = pd.Series(data=backrep[:,3], index=new_df.index)
            new_df['Hips_beta'] = pd.Series(data=backrep[:,4], index=new_df.index)
                       
            new_df['Hips_gamma'] = pd.Series(data=backrep[:,5], index=new_df.index)
           

            new_track.values = new_df

            Q.append(new_track)
        
        return Q

    


def gn_load_data2020_1(data_version='v2_0', smooth_audio=False, data_root='2020/sync_v1.1', framerate=30, use_saved_data=True):
    _data_root = os.path.join(GN_DATA_ROOT, data_root)

    params = {}
    params['audio_samplerate'] = 44100
    params['framerate'] = framerate
    params['hop_size'] = int(params['audio_samplerate'] / params['framerate'])
    params['window_size'] = params['hop_size'] * 4

    pickle_path = os.path.join(GN_DATA_ROOT, 'groovenet_essall%s+eul+pos_%s.pkl'%('_smooth' if smooth_audio else '', data_version))

    if use_saved_data and os.path.exists(pickle_path):
         with open(pickle_path, 'rb') as _pp:
            saved_data = pickle.load(_pp)

            return saved_data

    else:
        bvh_parser = BVHParser()

        eu2ex = MocapParameterizer('expmap')
        eu2pos = MocapParameterizer('position')

        audio_ess = {}
        euler_data = {}
        exp_data = {}
        pos_data = {}
        file_list = []

        for track in os.listdir(_data_root):
            print('processing track: %s'%track)
            a = extract_wrapper_o0(os.path.join(_data_root, track, '%s.mp3'%track), params, smooth_audio)['all']
            eu = bvh_parser.parse(os.path.join(_data_root, track, '%s.bvh'%track))

            # for j in eu.traverse():
            #     if j == eu.root_name:
            #         continue
            #     if 'Xposition' in eu.skeleton[j]['channels']:
            #         eu.skeleton[j]['channels'].remove('Xposition')
            #         eu.skeleton[j]['channels'].remove('Yposition')
            #         eu.skeleton[j]['channels'].remove('Zposition')                

            print(a['lowLevel.barkbands'].shape)
            print(eu.values.shape)
            # input('....')
            ex = 0#eu2ex.fit_transform([eu])[0]
            pp = eu2pos.fit_transform([eu])[0]
            
            # with open(os.path.join(data_root, track, '%s.pkl'%track), 'wb') as _pp:
            #     pickle.dump({
            #         'track_name': track,
            #         'audio_ess': a,
            #         'euler_data': eu,
            #         'pos_data': pp,
            #         'exp_data': ex
            #     }, _pp)


            file_list.append(track)
            audio_ess[track] = a
            euler_data[track] = eu
            exp_data[track] = ex
            pos_data[track] = pp


        save_data = {
            'euler_data': euler_data,
            'pos_data': pos_data,
            'exp_data': exp_data,
            'audio_ess': audio_ess,
            'file_list': file_list,
            'aud_feat_keys': list(audio_ess[file_list[0]].keys())
        }

        with open(pickle_path, 'wb') as _pp:
            pickle.dump(save_data, _pp)


        return save_data


class TaylorRootRep_(BaseEstimator, TransformerMixin):  
    def __init__(self):       
        
        self.start_face = 0
        self.start_x = 0
        self.start_z = 0


    def fit(self, X, y=None):
        return self

    def transform(self, X, y=None):
        import matlab.engine
        eng = matlab.engine.start_matlab()
        Q = []

        for track in X:
            new_track = track.clone()
            new_df = track.values.copy()
            
            org_root = track.values[['Hips_Xposition', 'Hips_Yposition', 'Hips_Zposition', 'Hips_alpha', 'Hips_beta', 'Hips_gamma']].values
                    
            new_df.drop(['Hips_Xposition', 'Hips_Zposition', 'Hips_alpha', 'Hips_beta', 'Hips_gamma'], axis=1, inplace=True)


            eng = matlab.engine.start_matlab()
            ty_rep, start_face = eng.ty_pre(matlab.double(org_root.tolist()), nargout=2)
            self.start_face = start_face

            ty_rep = np.asarray(ty_rep)

            new_df['root_rot_a'] = pd.Series(data=ty_rep[:,0], index=new_df.index)
            new_df['root_rot_b'] = pd.Series(data=ty_rep[:,1], index=new_df.index)

            new_df['root_rot_vert_delta'] = pd.Series(data=ty_rep[:,2], index=new_df.index)

            new_df['root_delta_x'] = pd.Series(data=ty_rep[:,3], index=new_df.index)

            new_df['root_delta_z'] = pd.Series(data=ty_rep[:,4], index=new_df.index)
            

            new_track.values = new_df

            Q.append(new_track)
        
        return Q

    def inverse_transform(self, X, copy=None):
        import matlab.engine
        eng = matlab.engine.start_matlab()
        Q = []

        for track in X:
            new_track = track.clone()
            new_df = track.values.copy()

            ty_rep = track.values[['root_rot_a', 'root_rot_b', 'root_rot_vert_delta', 'root_delta_x', 'root_delta_z', 'Hips_Yposition']].values
            
            new_df.drop(['root_rot_a', 'root_rot_b', 'root_rot_vert_delta', 'root_delta_x', 'root_delta_z'], axis=1, inplace=True)
           

            backrep = eng.ty_post(matlab.double(ty_rep.tolist()), self.start_face, self.start_x, self.start_z, nargout=1)

            backrep = np.asarray(backrep)
        
            new_df['Hips_Xposition'] = pd.Series(data=backrep[:,0], index=new_df.index)
            new_df['Hips_Zposition'] = pd.Series(data=backrep[:,2], index=new_df.index)

            new_df['Hips_alpha'] = pd.Series(data=backrep[:,3], index=new_df.index)
            new_df['Hips_beta'] = pd.Series(data=backrep[:,4], index=new_df.index)
                       
            new_df['Hips_gamma'] = pd.Series(data=backrep[:,5], index=new_df.index)
           

            new_track.values = new_df

            Q.append(new_track)
        
        return Q


# cache_path = '%/cached_param/'%GN_DATA_ROOT

# class CachedMocapParam(BaseEstimator, TransformerMixin):
#     def __init__(self, param_type):
#         self.param_type = param_type

#     def fit(self, X, y=None):
#         return self

#     def transform(self, X, y=None):
#         for track
#         os.path.exists()
        
        
#         if self.param_type == 'euler':
#             return X
#         elif self.param_type == 'expmap':
#             return self._to_expmap(X)
#         elif self.param_type == 'exp_euroot':
#             return self._to_exp_euroot(X)
#         elif self.param_type == 'quat':
#             return X
#         elif self.param_type == 'position':
#             return self._to_pos(X)
#         else:
#             raise UnsupportedParamError('Unsupported param: %s. Valid param types are: euler, quat, expmap, position' % self.param_type)
# #        return X
    
#     def inverse_transform(self, X, copy=None): 
#         pass

class ParMocapParameterizer(BaseEstimator, TransformerMixin):
    def __init__(self, param_type = 'euler', n_jobs=4):
        '''
        
        param_type = {'euler', 'quat', 'expmap', 'position', 'exp_euroot'}
        '''
        
        self.param_type = param_type
        self.n_jobs = n_jobs
        

    def fit(self, X, y=None):
        return self

    def transform(self, X, y=None):
        if self.param_type == 'euler':
            return X
        elif self.param_type == 'expmap':
            pfn = self._to_expmap
        elif self.param_type == 'exp_euroot':
            pfn = self._to_exp_euroot
        elif self.param_type == 'quat':
            return X
        elif self.param_type == 'position':
            pfn = self._to_pos
        else:
            raise UnsupportedParamError('Unsupported param: %s. Valid param types are: euler, quat, expmap, position' % self.param_type)
#        return X

        import multiprocessing
        pool = multiprocessing.Pool(self.n_jobs)
        result = pool.map(pfn, X)
        print(len(result))
        print(result[0])
        return [x[0] for x in result]
    
    def inverse_transform(self, X, copy=None): 
        if self.param_type == 'euler':
            return X
        elif self.param_type == 'expmap':
            return self._expmap_to_euler(X)
        elif self.param_type == 'exp_euroot':
            return self._exp_euroot_to_euler(X)
        elif self.param_type == 'quat':
            raise UnsupportedParamError('quat2euler is not supported')
        elif self.param_type == 'position':
            print('positions 2 eulers is not supported')
            return X
        else:
            raise UnsupportedParamError('Unsupported param: %s. Valid param types are: euler, quat, expmap, position' % self.param_type)

    def _to_pos(self, X):
        '''Converts joints rotations in Euler angles to joint positions'''
        
        X = [X]
        Q = []
        for track in X:
            channels = []
            titles = []
            euler_df = track.values

            # Create a new DataFrame to store the exponential map rep
            pos_df = pd.DataFrame(index=euler_df.index)

            # Copy the root rotations into the new DataFrame
            # rxp = '%s_Xrotation'%track.root_name
            # ryp = '%s_Yrotation'%track.root_name
            # rzp = '%s_Zrotation'%track.root_name
            # pos_df[rxp] = pd.Series(data=euler_df[rxp], index=pos_df.index)
            # pos_df[ryp] = pd.Series(data=euler_df[ryp], index=pos_df.index)
            # pos_df[rzp] = pd.Series(data=euler_df[rzp], index=pos_df.index)

            # List the columns that contain rotation channels
            rot_cols = [c for c in euler_df.columns if ('rotation' in c)]

            # List the columns that contain position channels
            pos_cols = [c for c in euler_df.columns if ('position' in c)]

            # List the joints that are not end sites, i.e., have channels
            joints = (joint for joint in track.skeleton)
            
            tree_data = {}

            for joint in track.traverse():
                parent = track.skeleton[joint]['parent']

                # Get the rotation columns that belong to this joint
                rc = euler_df[[c for c in rot_cols if joint in c]]

                # Get the position columns that belong to this joint
                pc = euler_df[[c for c in pos_cols if joint in c]]

                # Make sure the columns are organized in xyz order
                if rc.shape[1] < 3:
                    euler_values = [[0,0,0] for f in rc.iterrows()]
                else:
                    euler_values = [[f[1]['%s_Xrotation'%joint], 
                                     f[1]['%s_Yrotation'%joint], 
                                     f[1]['%s_Zrotation'%joint]] for f in rc.iterrows()]

                    ################# in euler angle, the order of rotation axis is very important #####################
                    rotation_order = rc.columns[0][rc.columns[0].find('rotation') - 1] + rc.columns[1][rc.columns[1].find('rotation') - 1] + rc.columns[2][rc.columns[2].find('rotation') - 1] #rotation_order is string : 'XYZ' or'ZYX' or ...
                    ####################################################################################################

                if pc.shape[1] < 3:
                    pos_values = [[0,0,0] for f in pc.iterrows()]
                else:
                    pos_values =[[f[1]['%s_Xposition'%joint], 
                                  f[1]['%s_Yposition'%joint], 
                                  f[1]['%s_Zposition'%joint]] for f in pc.iterrows()]

                #euler_values = [[0,0,0] for f in rc.iterrows()] #for deugging
                #pos_values = [[0,0,0] for f in pc.iterrows()] #for deugging
                
                # Convert the eulers to rotation matrices
                ############################ input rotation order as Rotation class's argument #########################
                rotmats = np.asarray([Rotation([f[0], f[1], f[2]], 'euler', rotation_order=rotation_order, from_deg=True).rotmat for f in euler_values])
                ########################################################################################################
                tree_data[joint]=[
                                    [], # to store the rotation matrix
                                    []  # to store the calculated position
                                 ] 

                if track.root_name == joint:
                    tree_data[joint][0] = rotmats
                    # tree_data[joint][1] = np.add(pos_values, track.skeleton[joint]['offsets'])
                    tree_data[joint][1] = pos_values
                else:
                    # for every frame i, multiply this joint's rotmat to the rotmat of its parent
                    tree_data[joint][0] = np.asarray([np.matmul(rotmats[i], tree_data[parent][0][i]) 
                                                      for i in range(len(tree_data[parent][0]))])

                    # add the position channel to the offset and store it in k, for every frame i
                    k = np.asarray([np.add(pos_values[i],  track.skeleton[joint]['offsets'])
                                    for i in range(len(tree_data[parent][0]))])

                    # multiply k to the rotmat of the parent for every frame i
                    q = np.asarray([np.matmul(k[i], tree_data[parent][0][i]) 
                                    for i in range(len(tree_data[parent][0]))])

                    # add q to the position of the parent, for every frame i
                    tree_data[joint][1] = np.asarray([np.add(q[i], tree_data[parent][1][i])
                                                      for i in range(len(tree_data[parent][1]))])


                # Create the corresponding columns in the new DataFrame
                pos_df['%s_Xposition'%joint] = pd.Series(data=[e[0] for e in tree_data[joint][1]], index=pos_df.index)
                pos_df['%s_Yposition'%joint] = pd.Series(data=[e[1] for e in tree_data[joint][1]], index=pos_df.index)
                pos_df['%s_Zposition'%joint] = pd.Series(data=[e[2] for e in tree_data[joint][1]], index=pos_df.index)


            new_track = track.clone()
            new_track.values = pos_df
            Q.append(new_track)
        return Q

    #TODO: This needs to be checked for other data - Works for ECU 2020 Data
    def _to_expmap(self, X):
        '''Converts Euler angles to Exponential Maps'''
        X = [X]
        def process_track(t):
            track = X[t]
            channels = []
            titles = []
            euler_df = track.values

            # Create a new DataFrame to store the exponential map rep
            exp_df = pd.DataFrame(index=euler_df.index)

            # Copy the root positions into the new DataFrame
            rxp = '%s_Xposition'%track.root_name
            ryp = '%s_Yposition'%track.root_name
            rzp = '%s_Zposition'%track.root_name
            exp_df[rxp] = pd.Series(data=euler_df[rxp], index=exp_df.index)
            exp_df[ryp] = pd.Series(data=euler_df[ryp], index=exp_df.index)
            exp_df[rzp] = pd.Series(data=euler_df[rzp], index=exp_df.index)
            
            # List the columns that contain rotation channels
            rots = [c for c in euler_df.columns if ('rotation' in c and 'Nub' not in c)]

            # List the joints that are not end sites, i.e., have channels
            joints = (joint for joint in track.skeleton if 'Nub' not in joint)

            for joint in joints:
                r = euler_df[[c for c in rots if joint in c]] # Get the columns that belong to this joint
                euler = [[f[1]['%s_Zrotation'%joint], f[1]['%s_Xrotation'%joint], f[1]['%s_Yrotation'%joint]] for f in r.iterrows()] # Make sure the columsn are organized in xyz order
                exps = [Rotation(f, 'euler', from_deg=True, rotation_order='XYZ').to_expmap() for f in euler] # Convert the eulers to exp maps
                
                # Create the corresponding columns in the new DataFrame
    
                exp_df['%s_alpha'%joint] = pd.Series(data=[e[1] for e in exps], index=exp_df.index)
                exp_df['%s_beta'%joint] = pd.Series(data=[e[2] for e in exps], index=exp_df.index)
                exp_df['%s_gamma'%joint] = pd.Series(data=[e[0] for e in exps], index=exp_df.index)

            new_track = track.clone()
            new_track.values = exp_df
            return t, new_track

        


        # with ThreadPoolExecutor(max_workers = 6) as executor:
        #     results = executor.map(process_track, range(len(X)))
        
        Q = []
        for i in range(len(X)):
            _,q = process_track(i)
            print(q.values.shape)
            Q.append(q)
            # Q.append(results[]

        return Q

    def _expmap_to_euler(self, X):
        Q = []
        for track in X:
            channels = []
            titles = []
            exp_df = track.values

            # Create a new DataFrame to store the exponential map rep
            euler_df = pd.DataFrame(index=exp_df.index)

            # Copy the root positions into the new DataFrame
            rxp = '%s_Xposition'%track.root_name
            ryp = '%s_Yposition'%track.root_name
            rzp = '%s_Zposition'%track.root_name
            euler_df[rxp] = pd.Series(data=exp_df[rxp], index=euler_df.index)
            euler_df[ryp] = pd.Series(data=exp_df[ryp], index=euler_df.index)
            euler_df[rzp] = pd.Series(data=exp_df[rzp], index=euler_df.index)
            
            # List the columns that contain rotation channels
            exp_params = [c for c in exp_df.columns if ( any(p in c for p in ['alpha', 'beta','gamma']) and 'Nub' not in c)]

            # List the joints that are not end sites, i.e., have channels
            joints = (joint for joint in track.skeleton if 'Nub' not in joint)

            for joint in joints:
                r = exp_df[[c for c in exp_params if joint in c]] # Get the columns that belong to this joint
                expmap = [[f[1]['%s_gamma'%joint], f[1]['%s_alpha'%joint], f[1]['%s_beta'%joint]] for f in r.iterrows()] # Make sure the columsn are organized in xyz order
                euler_rots = [Rotation(f, 'expmap', rotation_order='XYZ').to_euler(True)[0] for f in expmap] # Convert the eulers to exp maps
                
                # Create the corresponding columns in the new DataFrame
                
                euler_df['%s_Zrotation'%joint] = pd.Series(data=[-e[0] for e in euler_rots], index=euler_df.index)
                euler_df['%s_Xrotation'%joint] = pd.Series(data=[-e[1] for e in euler_rots], index=euler_df.index)
                euler_df['%s_Yrotation'%joint] = pd.Series(data=[-e[2] for e in euler_rots], index=euler_df.index)

            new_track = track.clone()
            new_track.values = euler_df
            Q.append(new_track)

        return Q
    
    def _to_exp_euroot(self, X):
        '''Converts Euler angles to Exponential Maps but keeps the root orientation in Euler'''

        Q = []
        X = [X]
        for track in X:
            channels = []
            titles = []
            euler_df = track.values

            # Create a new DataFrame to store the exponential map rep
            exp_df = pd.DataFrame(index=euler_df.index)

            # Copy the root positions into the new DataFrame
            rxp = '%s_Xposition'%track.root_name
            ryp = '%s_Yposition'%track.root_name
            rzp = '%s_Zposition'%track.root_name
            rxr = '%s_Xrotation'%track.root_name
            ryr = '%s_Yrotation'%track.root_name
            rzr = '%s_Zrotation'%track.root_name
            
            exp_df[rxp] = pd.Series(data=euler_df[rxp], index=exp_df.index)
            exp_df[ryp] = pd.Series(data=euler_df[ryp], index=exp_df.index)
            exp_df[rzp] = pd.Series(data=euler_df[rzp], index=exp_df.index)
            exp_df[rxr] = pd.Series(data=euler_df[rxr], index=exp_df.index)
            exp_df[ryr] = pd.Series(data=euler_df[ryr], index=exp_df.index)
            exp_df[rzr] = pd.Series(data=euler_df[rzr], index=exp_df.index)
            
            # List the columns that contain rotation channels
            rots = [c for c in euler_df.columns if ('rotation' in c and 'Nub' not in c)]

            # List the joints that are not end sites, i.e., have channels
            joints = (joint for joint in track.skeleton if 'Nub' not in joint and track.root_name not in joint)

            for joint in joints:
                r = euler_df[[c for c in rots if joint in c]] # Get the columns that belong to this joint
                euler = [[f[1]['%s_Xrotation'%joint], f[1]['%s_Yrotation'%joint], f[1]['%s_Zrotation'%joint]] for f in r.iterrows()] # Make sure the columsn are organized in xyz order
                exps = [Rotation(f, 'euler', from_deg=True).to_expmap() for f in euler] # Convert the eulers to exp maps
                
                # Create the corresponding columns in the new DataFrame
    
                exp_df['%s_alpha'%joint] = pd.Series(data=[e[0] for e in exps], index=exp_df.index)
                exp_df['%s_beta'%joint] = pd.Series(data=[e[1] for e in exps], index=exp_df.index)
                exp_df['%s_gamma'%joint] = pd.Series(data=[e[2] for e in exps], index=exp_df.index)

            new_track = track.clone()
            new_track.values = exp_df
            Q.append(new_track)

        return Q
    
    def _exp_euroot_to_euler(self, X):
        Q = []
        for track in X:
            channels = []
            titles = []
            exp_df = track.values

            # Create a new DataFrame to store the exponential map rep
            euler_df = pd.DataFrame(index=exp_df.index)

            # Copy the root positions into the new DataFrame
            rxp = '%s_Xposition'%track.root_name
            ryp = '%s_Yposition'%track.root_name
            rzp = '%s_Zposition'%track.root_name
            rxr = '%s_Xrotation'%track.root_name
            ryr = '%s_Yrotation'%track.root_name
            rzr = '%s_Zrotation'%track.root_name
            
            euler_df[rxp] = pd.Series(data=exp_df[rxp], index=exp_df.index)
            euler_df[ryp] = pd.Series(data=exp_df[ryp], index=exp_df.index)
            euler_df[rzp] = pd.Series(data=exp_df[rzp], index=exp_df.index)
            euler_df[rxr] = pd.Series(data=exp_df[rxr], index=exp_df.index)
            euler_df[ryr] = pd.Series(data=exp_df[ryr], index=exp_df.index)
            euler_df[rzr] = pd.Series(data=exp_df[rzr], index=exp_df.index)
            
            # List the columns that contain rotation channels
            exp_params = [c for c in exp_df.columns if ( any(p in c for p in ['alpha', 'beta','gamma']) and 'Nub' not in c)]

            # List the joints that are not end sites, i.e., have channels
            joints = (joint for joint in track.skeleton if 'Nub' not in joint and track.root_name not in joint)

            for joint in joints:
                r = exp_df[[c for c in exp_params if joint in c]] # Get the columns that belong to this joint
                expmap = [[f[1]['%s_alpha'%joint], f[1]['%s_beta'%joint], f[1]['%s_gamma'%joint]] for f in r.iterrows()] # Make sure the columsn are organized in xyz order
                euler_rots = [Rotation(f, 'expmap').to_euler(True)[0] for f in expmap] # Convert the eulers to exp maps
                
                # Create the corresponding columns in the new DataFrame
    
                euler_df['%s_Xrotation'%joint] = pd.Series(data=[-e[0] for e in euler_rots], index=euler_df.index)
                euler_df['%s_Yrotation'%joint] = pd.Series(data=[-e[1] for e in euler_rots], index=euler_df.index)
                euler_df['%s_Zrotation'%joint] = pd.Series(data=[-e[2] for e in euler_rots], index=euler_df.index)

            new_track = track.clone()
            new_track.values = euler_df
            Q.append(new_track)

        return Q

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



class TestMocapParameterizer(BaseEstimator, TransformerMixin):
    def __init__(self, param_type = 'euler', n_jobs=4):
        '''
        
        param_type = {'euler', 'quat', 'expmap', 'position', 'exp_euroot'}
        '''
        
        self.param_type = param_type
        self.n_jobs = n_jobs
        

    def fit(self, X, y=None):
        return self

    def transform(self, X, y=None):
        if self.param_type == 'euler':
            return X
        elif self.param_type == 'expmap':
            pfn = self._to_expmap
        elif self.param_type == 'exp_euroot':
            pfn = self._to_exp_euroot
        elif self.param_type == 'quat':
            return X
        elif self.param_type == 'position':
            pfn = self._to_pos
        else:
            raise UnsupportedParamError('Unsupported param: %s. Valid param types are: euler, quat, expmap, position' % self.param_type)
#        return X

        import multiprocessing
        pool = multiprocessing.Pool(self.n_jobs)
        result = pool.map(pfn, X)
        print(len(result))
        print(result[0])
        return [x[0] for x in result]
  
    def _to_pos(self, X):
        '''Converts joints rotations in Euler angles to joint positions'''
        
        X = [X]
        Q = []
        for track in X:
            channels = []
            titles = []
            euler_df = track.values

            # Create a new DataFrame to store the exponential map rep
            pos_df = pd.DataFrame(index=euler_df.index)

            # Copy the root rotations into the new DataFrame
            # rxp = '%s_Xrotation'%track.root_name
            # ryp = '%s_Yrotation'%track.root_name
            # rzp = '%s_Zrotation'%track.root_name
            # pos_df[rxp] = pd.Series(data=euler_df[rxp], index=pos_df.index)
            # pos_df[ryp] = pd.Series(data=euler_df[ryp], index=pos_df.index)
            # pos_df[rzp] = pd.Series(data=euler_df[rzp], index=pos_df.index)

            # List the columns that contain rotation channels
            rot_cols = [c for c in euler_df.columns if ('rotation' in c)]

            # List the columns that contain position channels
            pos_cols = [c for c in euler_df.columns if ('position' in c)]

            # List the joints that are not end sites, i.e., have channels
            joints = (joint for joint in track.skeleton)
            
            tree_data = {}

            for joint in track.traverse():
                parent = track.skeleton[joint]['parent']

                # Get the rotation columns that belong to this joint
                rc = euler_df[[c for c in rot_cols if joint in c]]

                # Get the position columns that belong to this joint
                pc = euler_df[[c for c in pos_cols if joint in c]]

                # Make sure the columns are organized in xyz order
                if rc.shape[1] < 3:
                    euler_values = [[0,0,0] for f in rc.iterrows()]
                else:
                    euler_values = [[f[1]['%s_Zrotation'%joint], 
                                     f[1]['%s_Xrotation'%joint], 
                                     f[1]['%s_Yrotation'%joint]] for f in rc.iterrows()]

                if pc.shape[1] < 3:
                    pos_values = [[0,0,0] for f in pc.iterrows()]
                else:
                    pos_values =[[f[1]['%s_Zposition'%joint], 
                                  f[1]['%s_Xposition'%joint], 
                                  f[1]['%s_Yposition'%joint]] for f in pc.iterrows()]

                #euler_values = [[0,0,0] for f in rc.iterrows()] #for deugging
                #pos_values = [[0,0,0] for f in pc.iterrows()] #for deugging
                
                # Convert the eulers to rotation matrices
                ############################ input rotation order as Rotation class's argument #########################
                rotmats = np.asarray([Rotation([f[0], f[1], f[2]], 'euler', from_deg=True).rotmat for f in euler_values])
                ########################################################################################################
                tree_data[joint]=[
                                    [], # to store the rotation matrix
                                    []  # to store the calculated position
                                 ] 

                if track.root_name == joint:
                    tree_data[joint][0] = rotmats
                    # tree_data[joint][1] = np.add(pos_values, track.skeleton[joint]['offsets'])
                    tree_data[joint][1] = pos_values
                else:
                    # for every frame i, multiply this joint's rotmat to the rotmat of its parent
                    tree_data[joint][0] = np.asarray([np.matmul(rotmats[i], tree_data[parent][0][i]) 
                                                      for i in range(len(tree_data[parent][0]))])

                    # add the position channel to the offset and store it in k, for every frame i
                    k = np.asarray([np.add(pos_values[i],  track.skeleton[joint]['offsets'])
                                    for i in range(len(tree_data[parent][0]))])

                    # multiply k to the rotmat of the parent for every frame i
                    q = np.asarray([np.matmul(k[i], tree_data[parent][0][i]) 
                                    for i in range(len(tree_data[parent][0]))])

                    # add q to the position of the parent, for every frame i
                    tree_data[joint][1] = np.asarray([np.add(q[i], tree_data[parent][1][i])
                                                      for i in range(len(tree_data[parent][1]))])


                # Create the corresponding columns in the new DataFrame
                pos_df['%s_Xposition'%joint] = pd.Series(data=[e[0] for e in tree_data[joint][1]], index=pos_df.index)
                pos_df['%s_Yposition'%joint] = pd.Series(data=[e[1] for e in tree_data[joint][1]], index=pos_df.index)
                pos_df['%s_Zposition'%joint] = pd.Series(data=[e[2] for e in tree_data[joint][1]], index=pos_df.index)


            new_track = track.clone()
            new_track.values = pos_df
            Q.append(new_track)
        return Q
