import pandas as pd
from sklearn.model_selection import StratifiedKFold

csv_path = '/workspace/odelia-data-scripts/csvs/combined_metadata.csv'
df = pd.read_csv(csv_path)

# Ensure Lesion is integer
df['Lesion'] = df['Lesion'].astype(int)

# Step 1: Get max lesion label per patient
patient_labels = df.groupby('PatientID')['Lesion'].max().reset_index()

# Step 2: Prepare data for stratified k-fold
patient_ids = patient_labels['PatientID'].values
labels = patient_labels['Lesion'].values

skf = StratifiedKFold(n_splits=5, shuffle=True, random_state=42)

# Loop through folds
for fold, (train_idx, val_idx) in enumerate(skf.split(patient_ids, labels), 1):
    train_patients = patient_ids[train_idx]
    val_patients = patient_ids[val_idx]

    # Filter rows for the current split
    train_df = df[df['PatientID'].isin(train_patients)][['UID', 'Lesion', 'source']]
    val_df = df[df['PatientID'].isin(val_patients)][['UID', 'Lesion', 'source']]

    # Save to CSV
    train_df.rename(columns={'UID': 'patient_id', 'Lesion': 'label'}).to_csv(
        f'/workspace/odelia-data-scripts/odelia/train_ids_fold{fold}.csv', index=False
    )
    val_df.rename(columns={'UID': 'patient_id', 'Lesion': 'label'}).to_csv(
        f'/workspace/odelia-data-scripts/odelia/val_ids_fold{fold}.csv', index=False
    )

print("5-fold cross-validation splits saved.")