import os
import shutil
import pandas as pd

melignent = ["MEL", "SCC", "BCC"]
benign = ["NV", "BKL", "DF", "VASC", "AK"]

# Path to the folder containing images after augmentation
output_dir = "dataset/processed_images/output"

# Path to the CSV file with the labels
csv_file = "dataset/diagnoses.csv"  # Update this to the correct path of your CSV file

# Read the CSV file
df = pd.read_csv(csv_file)

# Create folders for each category if they don't exist
malignant_dir = os.path.join(output_dir, "malignant")
benign_dir = os.path.join(output_dir, "benign")
unk_dir = os.path.join(output_dir, "unknown")

# Get a list of all images in the output directory
all_images = os.listdir(output_dir)

# Iterate over each row in the CSV file
for index, row in df.iterrows():
    isic_id = row["image"]

    # Find the matching file in the `output` directory
    matching_files = [f for f in all_images if isic_id in f]  # Looks for any file containing the ISIC ID

    # If matching files are found, move them to the appropriate folder
    for image_name in matching_files:
        image_path = os.path.join(output_dir, image_name)
        if os.path.exists(image_path):
            if row["image"] in melignent:  # Melanoma = malignant
                shutil.move(image_path, os.path.join(malignant_dir, image_name))
            elif row["image"] in benign:
                shutil.move(image_path, os.path.join(benign_dir, image_name))
            else:
                shutil.move(image_path, os.path.join(unk_dir, image_name))