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, "output", "quality_control_aidaqc", "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 the session mapping dictionary
session_mapping = {
    "ses-Baseline": 0, "ses-Baseline1": 0, "ses-pre": 0,
    "ses-P1": 3, "ses-P2": 3, "ses-P3": 3, "ses-P4": 3, "ses-P5": 3,
    "ses-P6": 7, "ses-P7": 7, "ses-P8": 7, "ses-P9": 7, "ses-P10": 7,
    "ses-P11": 14, "ses-P12": 14, "ses-P13": 14, "ses-P14": 14, 
    "ses-P15": 14, "ses-P16": 14, "ses-P17": 14, "ses-P18": 14, 
    "ses-P19": 21, "ses-D21": 21, "ses-P20": 21, "ses-P21": 21,
    "ses-P22": 21, "ses-P23": 21, "ses-P24": 21, "ses-P25": 21,
    "ses-P26": 28, "ses-P27": 28, "ses-P28": 28, "ses-P29": 28, 
    "ses-P30": 28, "ses-P42": 42, "ses-P43": 42,
    "ses-P56": 56, "ses-P57": 56, "ses-P58": 56,
    "ses-P151": 28
}

# 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))

# Map the timepoint using the updated session mapping directly
df_qc['merged_timepoint'] = df_qc['tp'].map(session_mapping)

# 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, "output", "quality_control_aidaqc", "voting_remapped.csv")
df_qc.to_csv(output_csv, index=False)

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