# Options: 
# 1. what types of features to use as input (GT/learned, audio/video). Which img view is most occluded?
# 2. what types of models to use for learning (CNN/ SVM/ RF)
# 3. hyperparams (no of layers, kernel type, no. of trees, how long to train)
# 4. train-test split
# 5. train/test mode
# 6. granularity of output labels (each mistake individually, 0-1 (yes/no))


## load data
import pandas as pd
import numpy as np
from sklearn.model_selection import train_test_split
from sklearn.model_selection import KFold
from sklearn.ensemble import RandomForestClassifier
from sklearn.model_selection import cross_val_score, cross_val_predict
from sklearn import metrics
import matplotlib.pyplot as plt
import argparse
import random

class Seg():
    def __init__(self, args):
        self.task = args.task
        self.demo_type = args.demo_type
        self.split_data = args.split
        self.label_type = args.label_type
        self.viz = args.viz
        self.feature_type = args.feature_type

        self.csv_file = '../data/rf_data_'+self.task+'_'+self.demo_type+'.csv'

        # Instantiate model with 1000 decision trees
        self.rf = RandomForestClassifier(n_estimators = 100) #1000

        self.gt_feats = ['utt_dur','utt_presence_No','utt_presence_Yes',  \
            'silence_presence_label_No', 'silence_presence_label_Yes', 'normal_speech_presence_No', 'normal_speech_presence_Yes',\
            'frustration_presence_label_No', 'frustration_presence_label_Yes','encourage_presence_label_No','encourage_presence_label_Yes',\
            'laughter_presence_No','laughter_presence_Yes', 'surprise_presence_No','surprise_presence_Yes',\
            'obj_sound_label_No', 'obj_sound_label_Yes']

        self.prosody_feats = ['mean_loudness', 'max_loudness', 'mean_pitch', 'max_pitch', 'mean_energy', 'max_energy', \
        'total_energy', 'word_count', 'word_rate', 'pause']

        # self.gt_prosody_feats = ['utt_dur'] + self.prosody_feats
        self.gt_prosody_feats = self.gt_feats + self.prosody_feats

        self.train_user_count, self.test_user_count = 14, 4
        self.users = ['user2', 'user3', 'user4', 'user5', 'user6', 'user7', 'user8', 'user9', 'user10',\
            'user11', 'user12','user14', 'user15', 'user16', 'user17', 'user18', 'user19', 'user20']

    def diff_list(self, li1, li2):
        li_dif = [i for i in li1 + li2 if i not in li1 or i not in li2]
        return li_dif

    def load_pase_feats(self):
        pase_feats = np.load('../data/pase_feats_'+self.task+'_'+self.demo_type+'.npy')
        pase_feats =  pase_feats.item()
        # print(len(list(pase_feats.keys())))
        pase_rows = []
        for img_id in pase_feats:
            # print(pase_feats[img_id].shape)
            row = [img_id] + (pase_feats[img_id]).tolist()
            pase_rows.append(row)

        pase_column_names = ['pase_'+str(i) for i in range(len(pase_rows[0]))]
        self.pase_feats = pase_column_names[1:]
        df = pd.DataFrame(pase_rows, columns = ["img_id"]+pase_column_names[1:])
        return df

    def load_resnet_feats(self):
        resnet_feats = np.load('../data/resnet_feats_'+self.task+'_'+self.demo_type+'.npy')
        resnet_feats =  resnet_feats.item()
        # print(len(list(resnet_feats.keys())))
        resnet_rows = []
        for img_id in resnet_feats:
            # print(np.squeeze(resnet_feats[img_id]).shape)
            row = [img_id] + (np.squeeze(resnet_feats[img_id])).tolist()
            resnet_rows.append(row)

        resnet_column_names = ['resnet_'+str(i) for i in range(len(resnet_rows[0]))]
        self.resnet_feats = resnet_column_names[1:]
        df = pd.DataFrame(resnet_rows, columns = ["img_id"]+resnet_column_names[1:])
        return df

    def load_video_feats(self):
        video_feats = np.load('../data/video_feats/vid_feats_'+self.task+'_'+self.demo_type+'_0.5.npy')
        video_feats =  video_feats.item()
        # print(len(list(resnet_feats.keys())))
        video_rows = []
        for img_id in video_feats:
            # print(np.squeeze(resnet_feats[img_id]).shape)
            row = [img_id] + (np.squeeze(video_feats[img_id])).tolist()
            video_rows.append(row)

        video_column_names = ['resnet_'+str(i) for i in range(len(video_rows[0]))]
        self.video_feats = video_column_names[1:]
        df = pd.DataFrame(video_rows, columns = ["img_id"]+video_column_names[1:])
        return df


    def prepare_data(self):
        # read data 
        feats = pd.read_csv(self.csv_file)
        
        if self.feature_type=='prosody':
            self.feat_list = self.gt_prosody_feats
            self.features = feats
            
        elif self.feature_type=='vision':
            resnet_df = self.load_resnet_feats()
            self.feat_list = self.resnet_feats
            resnet_prosody = pd.merge(feats, resnet_df, on='img_id', how='outer')
            self.features = resnet_prosody

        elif self.feature_type=='video':
            video_df = self.load_video_feats()
            self.feat_list = self.video_feats
            video_prosody = pd.merge(feats, video_df, on='img_id', how='outer')
            self.features = video_prosody

        elif self.feature_type=='prosody_vision':
            resnet_df = self.load_resnet_feats()
            self.feat_list = self.gt_prosody_feats+self.resnet_feats
            resnet_prosody = pd.merge(feats, resnet_df, on='img_id', how='outer')
            self.features = resnet_prosody

        elif self.feature_type=='pase':
            pase_df = self.load_pase_feats()
            self.feat_list = self.gt_prosody_feats+self.pase_feats
            pase_prosody = pd.merge(feats, pase_df, on='img_id', how='outer')
            self.features = pase_prosody

        elif self.feature_type=='pase_vision':
            resnet_df = self.load_resnet_feats()
            pase_df = self.load_pase_feats()
            self.feat_list = self.pase_feats+self.resnet_feats
            pase_vision = pd.merge(pase_df, resnet_df, on='img_id', how='outer')
            self.features = pase_vision

        elif self.feature_type=='prosody_pase_vision':
            resnet_df = self.load_resnet_feats()
            pase_df = self.load_pase_feats()
            self.feat_list = self.gt_prosody_feats+self.pase_feats+self.resnet_feats
            prosody_pase = pd.merge(feats, pase_df, on='img_id', how='outer')
            prosody_pase_vision = pd.merge(prosody_pase, resnet_df, on='img_id', how='outer')
            self.features = prosody_pase_vision

        # print(pase_prosody.shape[0],pase_df.shape[0],features.shape[0])
        # print(resnet_prosody.shape[0],resnet_df.shape[0],features.shape[0])
        # print(pase_prosody.describe())
        print(self.features.head(5))

        # count rows in dataframe
        # prosody_img_ids = set(features['img_id'].tolist())
        # pase_img_ids = set(pase_df['img_id'].tolist())
        # print(prosody_img_ids-pase_img_ids)

        # seg_labels = features.iloc[:,16] # presence
        if self.label_type=='seg_presence':
            self.label_id = "seg_presence_label"
        if self.label_type=='err_presence':
            self.label_id = "err_presence_label"
        elif self.label_type=='quality':
            self.label_id = "seg_quality_label"
            self.features = self.features[self.features[self.label_id].isin(['great','poor'])]
        elif self.label_type=='precision':
            self.label_id = "seg_precision_type"
            self.features = self.features[self.features[self.label_id].isin(['non-precision','precision'])]
        self.labels = self.features[[self.label_id]]
        # label = self.label_id
        # self.labels = self.features.label

        self.features= self.features.drop('img_id', axis = 1)
        # self.features= self.features.drop(self.label_id, axis = 1)
        float_feats = [self.label_id,'user_id']+self.prosody_feats
        if 'vision' in self.feature_type:
            float_feats+=self.resnet_feats
        if 'pase' in self.feature_type:
            float_feats+=self.pase_feats
        self.features = pd.get_dummies(self.features, columns=[col for col in self.features.columns if col not in float_feats])

        # self.features = pd.get_dummies(self.features)

        # select features for seg classification
        self.sel_features = self.features.filter(self.feat_list, axis=1)
        self.feature_list = list(self.sel_features.columns)

        
    def data_stats(self):
        print('Training Features Shape:', self.train_features.shape)
        print('Training Labels Shape:', self.train_labels.shape)
        print('Testing Features Shape:', self.test_features.shape)
        print('Testing Labels Shape:', self.test_labels.shape)
        

    def train_test(self):
        # Train the model on training data
        # self.rf.fit(self.train_features, self.train_labels.values.ravel())
        perf = []

        

        if self.split_data=='random':
            # feature_list = list(features.columns)

            # # convert features to numpy array
            # self.sel_features = np.array(self.sel_features)
            
            # self.train_features, self.test_features, self.train_labels, self.test_labels = train_test_split(self.sel_features, self.labels, test_size = 0.25)#, random_state = 42)
            kf = KFold(n_splits=5) # Define the split - into 5 folds 
            kf.get_n_splits(self.sel_features) # returns the number of splitting iterations in the cross-validator

            # trained_models = []
            for train_index, test_index in kf.split(self.sel_features):
                # print('TRAIN:', train_index, 'TEST:', test_index)
                X_train, X_test = self.sel_features[train_index], self.sel_features[test_index]
                y_train, y_test = self.labels[train_index], self.labels[test_index]
                
                X_train, X_test = np.array(X_train), np.array(X_test)

                model = self.rf.fit(X_train, y_train.values.ravel())
                # trained_models.append(model)

                # Use the forest's predict method on the test data
                # predictions = self.rf.predict(X_test)
                score = model.score(X_test, y_test.values.ravel())
                # accuracy = 100*score
                perf.append(score)

        elif self.split_data=='user':
            splits = [4,4,4,3,3]
            splits = random.shuffle(splits)
            train_users, test_users = [], []
            left = self.users
            for i in range(5):
                test = random.sample(left,splits[i])
                train = self.diff_list(self.users,test)

                train_users.append(train)
                test_users.append(test)

                left = self.diff_list(left,test)

            for tr_user_list,te_user_list in zip(train_users,test_users):
            
                X_train = self.sel_features[self.sel_features["user_id"].isin(tr_user_list)]
                X_test = self.sel_features[self.sel_features["user_id"].isin(te_user_list)]

                y_train = X_train[[self.label_id]]
                y_test = X_test[[self.label_id]]

                X_train= X_train.drop('img_id', axis = 1)
                X_train= X_train.drop(self.label_id, axis = 1)
                X_train = pd.get_dummies(X_train)

                X_test= X_test.drop('img_id', axis = 1)
                X_test= X_test.drop(self.label_id, axis = 1)
                X_test = pd.get_dummies(X_test)

                self.train_features = X_train.filter(self.feat_list, axis=1)
                self.feature_list = list(self.train_features.columns)
                self.train_features = np.array(X_train)

                X_test = X_test.filter(self.feat_list, axis=1)
                X_test = np.array(X_test)

                model = self.rf.fit(X_train, y_train.values.ravel())
                score = model.score(X_test, y_test.values.ravel())
                perf.append(score)


        # scores = cross_val_score(self.rf, self.sel_features, self.labels.values.ravel(), cv=5)
        # accuracy = 100*np.mean(scores)
        # predictions = cross_val_predict(self.rf, self.sel_features, self.labels, cv=6)
        # accuracy = metrics.r2_score(self.labels, predictions)
        accuracy = np.mean(perf) * 100
        return accuracy

    def viz_feats(self):
        ## Plot feature importance for GT
        # Get numerical feature importances
        importances = list(self.rf.feature_importances_)
        print(len(importances),len(self.feature_list))
        # List of tuples with variable and importance
        feature_importances = [(feature, round(importance, 2)) for feature, importance in zip(self.feature_list, importances)]
        # Sort the feature importances by most important first
        feature_importances = sorted(feature_importances, key = lambda x: x[1], reverse = True)

        # Set the style
        plt.style.use('fivethirtyeight')
        # list of x locations for plotting
        x_values = list(range(len(importances)))
        # Make a bar chart
        plt.bar(x_values, importances, orientation = 'vertical')
        # Tick labels for x axis
        # TODO: update name of features
        plt.xticks(x_values, self.feature_list, rotation='vertical')
        # Axis labels and title
        plt.ylabel('Importance'); plt.xlabel('Variable'); plt.title('Variable Importances')
        plt.show()


    def learn(self):
        self.prepare_data()
        accuracy = self.train_test()
        
        if self.viz:
            self.viz_feats()
        return accuracy

def main():
    parser = argparse.ArgumentParser()
    parser.add_argument('-t','--task',type=str,default='box',help='[box,cutting]')
    parser.add_argument('-d','--demo_type',type=str,default='video',help='[video,kt]')
    parser.add_argument('-s','--split',type=str,default='random',help='train-test split criteria. Options: [user,random]')
    parser.add_argument('-l','--label_type',type=str,default='seg_presence',help='[seg_presence,precision,quality,err_presence]')
    parser.add_argument('-v','--viz',action='store_true',default=False)
    parser.add_argument('-f','--feature_type',type=str,default='prosody',help='[prosody,prosody_vision,vision,pase_vision,pase,prosody_pase_vision]')

    args = parser.parse_args()

    seg = Seg(args)
    accuracy = seg.learn()
    print('Accuracy: {}%'.format(round(accuracy, 2)))

if __name__ == '__main__':
    main()