from sklearn.base import BaseEstimator, TransformerMixin
import numpy as np
import pandas as pd

# ## The Transformer
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

   