import os
import glob
import pandas as pd
import matplotlib.pyplot as plt
import math
import seaborn as sns
import warnings
from collections import defaultdict


class NostressProcessing():
    def __init__(self, participant_path, weight_df, conditions, results_path, time_difference):
        """_summary_

        Args:
            participant_path (str): _description_
            weight_df (Dataframe): dataframe containing the weights of all the participants
            conditions (list): list containing the name of the columns with the experiments or the conditions of the experiments 
            results_path (str): the path to save the results
            time_difference (int): the time difference between the excel file containg the timing of the experiment and the actigraph files (UTC+0).
                                      the given value will be added to the time in the actigraph data
        """
        self.participant_path = participant_path
        self.operators_weight = weight_df
        self.files_li = []
        # includes the different experiments done. (can be different conditions (Temperature, Time, noise), or different experiments compeletely)
        self.conditions = conditions
        self.results_path = results_path
        self.time_difference= int(time_difference)
        return
    
    
    def synchronize_data(self):
        """
        function to sychronize the data with the experiments.
        for each operator in each of the test cases, read the waist, wrist and ankle actigraph data.
        Using the excel file that contains the inforamtion about the timing of the experiments, the 3 data files are sychronized and the needed data columns from all the files are
        combined into one file. this file is then used to get the energy expenditure.
        """  
        excel_file = self.participant_path+'performance.xlsx'
        waist_file = self.participant_path +'waist\\ACTiLife_DataTable.csv'
        wrist_file = self.participant_path +'wrist\\ACTiLife_DataTable.csv'
        ankle_file = self.participant_path +'ankle\\ACTiLife_DataTable.csv'
        
        waist = pd.read_csv(waist_file, skiprows=10)
        wrist = pd.read_csv(wrist_file, skiprows=10)
        ankle = pd.read_csv(ankle_file, skiprows=10)
        #waist
        waist_col = waist.columns
        waist_col = [w[1:] if w[0]==' ' else w for w in waist_col]
        waist_col = waist_col[:2] + [w+'_waist' for w in waist_col[2:]]
        waist.columns = waist_col
        # wrist
        wrist_col = wrist.columns
        wrist_col = [w[1:] if w[0]==' ' else w for w in wrist_col]
        wrist_col = wrist_col[:2] + [w+'_wrist' for w in wrist_col[2:]]
        wrist.columns = wrist_col
        # ankle
        ankle_col = ankle.columns
        ankle_col = [w[1:] if w[0]==' ' else w for w in ankle_col]
        ankle_col = ankle_col[:2] + [w+'_ankle' for w in ankle_col[2:]]
        ankle.columns = ankle_col 
        
        # Convert the "Time" column to datetime objects (if it's not already) 
        # this should be changed to include the date. A date column should also be added to the excel sheet.
        waist['Time'] = pd.to_datetime(waist['Time'], format="%H:%M:%S")
        wrist['Time'] = pd.to_datetime(wrist['Time'], format="%H:%M:%S")
        ankle['Time'] = pd.to_datetime(wrist['Time'], format="%H:%M:%S")
        ### todo: add time zone
        # if self.time_difference:
            # Add 1 hour to the time values
            # print('time delta  = +1h')
        waist['Time'] = waist['Time'] + pd.Timedelta(hours=self.time_difference)
        wrist['Time'] = wrist['Time'] + pd.Timedelta(hours=self.time_difference)
        ankle['Time'] = ankle['Time'] + pd.Timedelta(hours=self.time_difference)
        waist['Time'] = waist['Time'].dt.time
        wrist['Time'] = wrist['Time'].dt.time
        ankle['Time'] = ankle['Time'].dt.time
        # v[-1] contains the path for the excel file for sychronization of the data
        subject_df = self.sychronize_operator_data(excel_file, waist, wrist, ankle)
        # the path to save the resulting data frame containing all the sychronized raw data.
        subject_csv_path = self.results_path + '\\' + excel_file.split('\\')[-2]
        if not os.path.exists(subject_csv_path):
            # Create the directory if it doesn't exist
            os.makedirs(subject_csv_path)
        # save the data frame in a csv file
        actigraph_data = subject_csv_path+'\\'+excel_file.split('\\')[-2]+'_actigraph_data.csv'
        self.actigraph_file = actigraph_data
        subject_df.to_csv(actigraph_data, index=False)
        return
           
                    
    def sychronize_operator_data(self, excel_path, waist, wrist, ankle):
        """reads the actigraph data files (waist, wrist and ankle). Using the excel file with the start and end time of each of the experiments,
        extracts the corresponding sections from the files, concatenates them in a dataframe.

        Args:
            excel_path (str): excel file with a sheets named "Performance"
            waist (Dataframe): waist 1sec DataTable
            wrist (Dataframe): wrist 1sec DataTable
            ankle (Dataframe): ankle 1sec DataTable

        Returns:
            _type_: _description_
        """
        dataframes_list = []
        # read the experiment data shet containing the timing and conditions of the experiments
        operator_df = pd.read_excel(excel_path, sheet_name='Performance', header=[0, 1])
        # Convert the contents of the 'Operator' column to lowercase
        operator_df['Operator'] = operator_df['Operator'].apply(lambda x: x.str.lower() if x.dtype == 'object' else x)
        for index, row in operator_df.iterrows():
            start = row[('Summary', 'Starting time')]
            finish = row[('Summary', 'Finishing time')]
            error = row[('Summary', 'Total errors')]
            if error in ['-','_','']:
                error = None
            # skip the iteration if the order doesn't have start or end time.
            if pd.isna(start) or pd.isna(finish):
                continue
            try:
                start = pd.to_datetime(start, format="%H:%M:%S").time()
                finish = pd.to_datetime(finish, format="%H:%M:%S").time()        
            except:
                start = pd.to_datetime(start).time()
                finish = pd.to_datetime(finish).time()
            # get the number of seconds to compare with the lengh of the dataframe. 
            # Convert time values to timedelta objects
            start_time = pd.to_timedelta(str(start))
            finish_time = pd.to_timedelta(str(finish))
            # Calculate the time difference
            time_difference = finish_time - start_time
            # Convert the time difference to seconds
            seconds_difference = int(time_difference.total_seconds() + 1)
            # get the data sections for waist and wrist
            waist_section = waist[(waist['Time'] >= start) & (waist['Time'] <= finish)]
            wrist_section = wrist[(wrist['Time'] >= start) & (wrist['Time'] <= finish)]
            ankle_section = ankle[(ankle['Time'] >= start) & (ankle['Time'] <= finish)]
            if len(waist_section)==0 or len(wrist_section)==0 or len(ankle_section)==0:
                print(f'Warning: data missing. The time interval of the experiment wasn\'t found in the wrist, waist or ankle file')
            if len(waist_section)!=seconds_difference:
                print(f'Warning: missmatch between data rows and time internval\n-number of data rows = {len(waist_section)}\n-number of seconds = {seconds_difference}')
                
            # concatinate the data for both waist and wrist. (remove the first two columns from waist, data and time)
            full_df = pd.concat([wrist_section, waist_section.iloc[:, 2:], ankle_section.iloc[:, 2:]], axis=1)
            full_df.rename(columns={'Time':'Time(hh:mm:ss)'}, inplace=True)

            # add the date and time columns
            df_add = full_df.iloc[:, :2]
            # add the columns from the row (from the excel file)
            # experiments
            if self.conditions[0]=='Experiment':
                df_add['Experiment'] = row[('Experiment', slice(None))].values[0]
            else:
                for cond in self.conditions:
                    df_add[cond] = row[('Conditions', cond)]
            df_add['Test_number'] = row[('Test number', slice(None))].values[0]
            # df_add['Order_number'] = order[-1]
            df_add['Error'] = error
            # add the rest of the data columns
            df_add = pd.concat([df_add, full_df.iloc[:, 2:]], axis=1)
            # add the Questionnaire
            for i in range(1,8,1):
                df_add['Q'+str(i)] = row[('Questionnaire items', i)]
            dataframes_list.append(df_add)
            # update result_df by concatination with df_add
        result_df = pd.concat(dataframes_list, ignore_index=True)
        return result_df
    

    def operator_energy_expenditure(self):
        """computes the energy expediture for an operator. the resutling data frame is saved into a csv file in the corresponding operator directory

        Args:
            operator_path (str): parth to the data file of an operator containing the actigraph data

        Returns:
            dataframe: the dataframe containing the energy expenditure data, as well as the stress score and the errors
        """
        df = pd.read_csv(self.actigraph_file)
        
        df.dropna(subset=['Test_number'], inplace=True)
        df['Time(hh:mm:ss)'] = pd.to_datetime(df['Time(hh:mm:ss)'], format='%H:%M:%S')
        df.set_index(pd.DatetimeIndex(df['Time(hh:mm:ss)']))
        
        df_min = pd.DataFrame()
        cols = ['Test_number', 'Error'] + self.conditions + ['Q1', 'Q2', 'Q3', 'Q4', 'Q5', 'Q6', 'Q7']
        for col in cols:
            df_min[col] = df.resample('T', on='Time(hh:mm:ss)')[col].first()
        # actigraph data resampleing
        cols = ['Axis1_wrist', 'Axis2_wrist', 'Axis3_wrist', 
                'Axis1_waist', 'Axis2_waist', 'Axis3_waist', 
                'Axis1_ankle', 'Axis2_ankle', 'Axis3_ankle']
        for col in cols:
            df_min[col] = df.resample('T', on='Time(hh:mm:ss)')[col].sum()
            
        #####to calculate VM from counts/min
        for index_min, row in df_min.iterrows():
            # wrist
            axis1 = df_min.loc[index_min,'Axis1_wrist']
            axis2 = df_min.loc[index_min,'Axis2_wrist']
            axis3 = df_min.loc[index_min,'Axis3_wrist']
            df_min.loc[index_min,'VM_wrist'] = math.sqrt(axis1**2 + axis2**2 + axis3**2)
            # waist
            axis1 = df_min.loc[index_min,'Axis1_waist']
            axis2 = df_min.loc[index_min,'Axis2_waist']
            axis3 = df_min.loc[index_min,'Axis3_waist']
            df_min.loc[index_min,'VM_waist'] = math.sqrt(axis1**2 + axis2**2 + axis3**2)
            # ankle
            axis1 = df_min.loc[index_min,'Axis1_ankle']
            axis2 = df_min.loc[index_min,'Axis2_ankle']
            axis3 = df_min.loc[index_min,'Axis3_ankle']
            df_min.loc[index_min,'VM_ankle'] = math.sqrt(axis1**2 + axis2**2 + axis3**2)

        df_test = pd.DataFrame()
        df_test= df_min.groupby(['Test_number']+self.conditions).mean() #df_min.groupby(df_min['Test_number']).mean()

        ##### A) Swartz(2000) equations for multitask activities:METs #kcal/min considering wrist and waist counts
        for index_test, row in df_test.iterrows():
            wrist_counts = df_test.loc[index_test,'VM_wrist']
            waist_counts = df_test.loc[index_test,'VM_waist']
            df_test.loc[index_test,'Swartz_wrist'] = ((3.195 + (0.0001314*wrist_counts))*3.5*self.operators_weight/200)#kcal/min
            df_test.loc[index_test,'Swartz_hip'] = ((2.606 + (0.0006863*waist_counts))*3.5*self.operators_weight/200) #kcal/min
            df_test.loc[index_test,'Swartz_hipwrist'] = ((2.245 + (0.000679*waist_counts) + (0.0001165*wrist_counts))*3.5*self.operators_weight/200) #kcal/min

        #Stress score
        df_test['Q2*'] = 7 - df_test['Q2']
        df_test['Q4*'] = 7 - df_test['Q4']
        df_test['Q5*'] = 7 - df_test['Q5']

        df_test['Score'] = df_test['Q1'] + df_test['Q2*'] + df_test['Q3'] + df_test['Q4*'] + df_test['Q5*'] + df_test['Q6'] + df_test['Q7']
        df_test['Stress_Score'] = df_test['Score'] / 49
        df_test['Stress_Score(%)'] = df_test['Stress_Score'] * 100
        df_test = df_test.round({'Stress_Score(%)':0})
        df_test.reset_index(inplace=True)
        # df_test['Error'] = df_error['Error'].to_numpy()
        
        save_path = '\\'.join(self.actigraph_file.split('\\')[:-1]) + '\\' + self.actigraph_file.split('\\')[-2] + '_energy_expenditure.csv'
        
        # reorder the columns
        df_test = df_test[['Test_number'] + self.conditions + ['Axis1_wrist', 'Axis2_wrist', 'Axis3_wrist', 
                                                              'Axis1_waist', 'Axis2_waist', 'Axis3_waist', 
                                                              'Axis1_ankle', 'Axis2_ankle', 'Axis3_ankle', 
                                                              'VM_wrist', 'VM_waist', 'Swartz_wrist', 'Swartz_hip', 'Swartz_hipwrist', 
                                                              'Q1', 'Q2', 'Q3', 'Q4', 'Q5', 'Q6', 'Q7', 'Q2*', 'Q4*', 'Q5*',
                                                              'Score', 'Stress_Score', 'Stress_Score(%)', 'Error']]
        df_test.to_csv(save_path, index=False)
        self.energy_expenditure_file = save_path
        return df_test
        

    def add_incrementing_index(self, lst):
        """If the same experiment is done multiple times witht the same operator, add an increasing index to each one

        Args:
            lst (list): list containing the experiments done by the operator

        Returns:
            list: list with an index added to the duplicate experiments
        """
        count_dict = defaultdict(int)
        result = []
        for item in lst:
            if lst.count(item) > 1:
                count_dict[item] += 1
                if item[-1]=='\n':
                    result.append(f'{item[:-1]} {count_dict[item]}')
                else:
                    result.append(f'{item} {count_dict[item]}')
            else:
                result.append(item)
        return result
    
    def plot_energy(self, y='Stress_Score(%)', y_scale=(0,100)):
        """function to plot the energy expenditure with the stress score or the number of errors.

        Args:
            y (str, optional): specifies the column to plot from the datafrme (Stress_Score or Error) in addition to the energy expenditure (Swartz_hipwrist). Defaults to 'Stress_Score(%)'.
            y_scale (tuple, optional): The scale to use for the second y-axis. Defaults to (0,100).
        """
        # Load CSV data into a DataFrame
        operator_data = pd.read_csv(self.energy_expenditure_file)
        # Set the style for Seaborn
        sns.set(style="whitegrid")
        # Define figure size and font size
        figure_size = (12, 6) 
        font_size = 12
        font_size_title = 16

        # Get the corresponding x-axis values for the operator (the experiments or conditions of the experiments)
        x = []
        for conditions in operator_data[self.conditions].values:
            condition = ''
            for c in conditions:
                condition +=  c + '\n'
            x.append(condition)
        x = self.add_incrementing_index(x)
        
        # Create a new figure and axis for each operator
        fig, ax1 = plt.subplots(figsize=figure_size)

        # Set the style and color palette for Seaborn
        sns.set(style="whitegrid")
        sns.set_palette("Set1")

        # Plot the first y-axis data for the current operator
        sns.lineplot(x=range(len(x)), y=operator_data['Swartz_hipwrist'], label=f'Energy Expenditure (kcal/min)', ax=ax1, color='b',
                        linestyle=':', marker='D', markersize=5, markeredgecolor='black')

        # Create a second y-axis
        ax2 = ax1.twinx()

        # Plot the second y-axis data for the current operator
        sns.lineplot(x=range(len(x)), y=operator_data[y], label=y.replace('_',' '), markers=True, ax=ax2, color='r')

        # Set the x-axis labels based on the combinations of Temperature, Noise, and Time
        if len(self.conditions)>1:
            ax1.set_xlabel('Experiment '+str(self.conditions).replace("'",''), fontsize=font_size_title)
        else:
            ax1.set_xlabel('Experiment', fontsize=font_size_title)
            
        ax1.set_xticks(range(len(x)))
        if len(x)>4:
            ax1.set_xticklabels(x, rotation=45, fontsize=font_size, rotation_mode='anchor', ha='right')
        else:
            ax1.set_xticklabels(x, rotation=0, fontsize=font_size, rotation_mode='anchor', ha='center')

        ax1.set_ylabel('Energy Expenditure (kcla/min)', fontsize=font_size_title)
        ax2.set_ylabel(y.replace('_',' '), fontsize=font_size_title)

        # Add a legend
        ax1.legend(loc='upper left')
        ax2.legend(loc='lower left')

        # Set y-axis limits
        ax1.set_ylim(0, 8)
        ax2.set_ylim(y_scale[0], y_scale[1])
        
        file_li = self.energy_expenditure_file.split('\\')
        participant_name = file_li[-2]
        
        plt.tight_layout()
        title = 'Energy expenditure vs '+ y.replace('_',' ') +', ' + participant_name
        plt.title(title, fontsize=font_size_title+2)
        plt.tight_layout()
        
        save_dir = self.results_path + '\\' + participant_name + '\\plots'
        # Check if the directory exists
        if not os.path.exists(save_dir):
            # Create the directory if it doesn't exist
            os.makedirs(save_dir)
        plt.savefig(f'{save_dir}\\{participant_name}_{title}.png')
        return


if __name__=="__main__":
    # Suppress all warnings
    warnings.filterwarnings("ignore")
    
    # list including the directories of the participants
    participants_dirs = glob.glob(os.path.dirname(__file__) + '\\NO-STRESS ACTiGraph\\*\\')
    # the directory to save the processed data, analysis and plots
    results_dir = os.path.dirname(__file__) + '\\processed_data'
    # get the weights from csv file
    weights_file = os.path.dirname(__file__) + '\\NO-STRESS ACTiGraph\\weights.csv'
    weights = pd.read_csv(weights_file)
    
    # conditions for the participants 01 to 04
    conditions = ['Temperature', 'Noise', 'Time']
    # experiment column for the participants 05 to 07
    experiments = ['Experiment']
    
    for participant in participants_dirs[:4]:
        part_name = participant.split('\\')[-2]
        # there is a 1 time difference between the local time of the experiment (in the excel file) and the actigraph data (UTC) for participants 1 to 4
        processing = NostressProcessing(participant_path=participant, weight_df=weights[part_name].values[0], conditions=conditions, results_path=results_dir, time_difference=1)
        # waist_file, wrist_file, ankle_file, excel_file = processing.get_data_files()
        processing.synchronize_data()
        processing.operator_energy_expenditure()
        # plot energy expenditure vs stress score
        processing.plot_energy(y='Stress_Score(%)', y_scale=(0,100))
        # plot energy expenditure vs number of errors.
        processing.plot_energy(y='Error', y_scale=(0,16))
        
    for participant in participants_dirs[4:]:
        part_name = participant.split('\\')[-2]
        # there is no time difference between the local time of the experiment (in the excel file) and the actigraph data (UTC) for participants 5,6,7
        processing = NostressProcessing(participant_path=participant, weight_df=weights[part_name].values[0], conditions=experiments, results_path=results_dir, time_difference=0)
        # waist_file, wrist_file, ankle_file, excel_file = processing.get_data_files()
        processing.synchronize_data()
        processing.operator_energy_expenditure()
        # plot energy expenditure vs stress score
        processing.plot_energy(y='Stress_Score(%)', y_scale=(0,100))

    input("Press Enter to exit...")