import pandas as pd
import os

# Get the directory where the code file is located
code_dir = os.path.dirname(os.path.abspath(__file__))

# Get the parent directory of the code directory
parent_dir = os.path.dirname(code_dir)
qc_csv = os.path.join(parent_dir, "input", "AIDAqc_ouptut_for_data", "Votings.csv")
df_qc = pd.read_csv(qc_csv)

# Define a function to extract the timepoint information
def extract_timepoint(path):
    parts = path.split(os.sep)
    for part in parts:
        if 'ses' in part:
            return part

# Define a function to extract the timepoint number
def extract_timepoint_number(tp):
    if 'Baseline' in tp:
        return 0
    else:
        return int(tp.split('ses-')[-1].replace("P",""))
    
# Define the session mapping dictionary
session_mapping = {
    0: 0,
    1: 3, 2: 3, 3: 3,
    4: 3, 5: 3, 6: 7, 7: 7,
    8: 7, 9: 7, 10: 7, 11: 14, 12: 14,
    13: 14, 14: 14, 15: 14, 16: 14, 17: 14, 18: 14, 19: 21,
    20: 21, 21: 21, 22: 21, 23: 21, 24: 21, 25: 21, 26: 28,
    27: 28, 28: 28, 29: 28, 30: 28 , 42:42, 43:42, 56:56, 57:56
}

# Define a function to extract the subject ID
def extract_subject_id(path):
    parts = path.split(os.sep)
    for part in parts:
        if 'sub-' in part and '.nii.' not in part:
            return part

# Remove rows containing "brkraw" or "DN" in the "Pathes" column
df_qc = df_qc[~df_qc['Pathes'].str.contains('brkraw|DN')]

# Create the "tp" column
df_qc['tp'] = df_qc['Pathes'].apply(lambda x: extract_timepoint(x))

# Define a function to map tp_int using session_mapping
def map_merged_timepoint(tp_int):
    return session_mapping.get(tp_int, tp_int)

# Create the "tp_int" column
df_qc['tp_int'] = df_qc['tp'].apply(lambda x: extract_timepoint_number(x))

# Create the "merged_timepoint" column
df_qc['merged_timepoint'] = df_qc['tp_int'].apply(map_merged_timepoint)

# Add the "subjectID" column
df_qc['subjectID'] = df_qc['Pathes'].apply(lambda x: extract_subject_id(x))

# Save the DataFrame as a CSV file
output_csv = os.path.join(parent_dir, "input", "AIDAqc_ouptut_for_data", "voting_remapped.csv")
df_qc.to_csv(output_csv, index=False)

print("DataFrame saved as 'voting_remapped.csv'")