import os
import csv
import shutil
import re
import argparse

def get_label_and_files(positive_dir, negative_dir):
    files_labels = []
    for label, folder in [('yes', positive_dir), ('no', negative_dir)]:
        if not os.path.exists(folder):
            continue
        for fname in os.listdir(folder):
            if fname.endswith('.png'):
                files_labels.append((fname, label, os.path.join(folder, fname)))
    return files_labels

def write_prediction_csv(files_labels, prediction_csv):
    with open(prediction_csv, 'w', newline='') as csvfile:
        writer = csv.writer(csvfile)
        writer.writerow(['filename', 'prediction'])
        for fname, label, _ in sorted(files_labels):
            writer.writerow([fname, label])

def extract_pid(filename):
    match = re.search(r'_T=(\d+)', filename)
    if match:
        return match.group(1)
    return None

def create_all_images_folder(files_labels, all_images_dir):
    if not os.path.exists(all_images_dir):
        os.makedirs(all_images_dir)
    for fname, _, src_path in files_labels:
        pid = extract_pid(fname)
        if pid is None:
            continue
        target_dir = os.path.join(all_images_dir, pid)
        os.makedirs(target_dir, exist_ok=True)
        shutil.copy2(src_path, os.path.join(target_dir, fname))

def main(input_path):
    prediction_csv = os.path.join(os.path.dirname(input_path), 'prediction.csv')
    all_images_dir = os.path.join(os.path.dirname(input_path), 'all-images')
    positive_dir = os.path.join(input_path, 'Positive')
    negative_dir = os.path.join(input_path, 'Negative')

    files_labels = get_label_and_files(positive_dir, negative_dir)
    write_prediction_csv(files_labels, prediction_csv)
    create_all_images_folder(files_labels, all_images_dir)
    print(f'prediction.csv and all-images folder created at: {os.path.dirname(input_path)}')

if __name__ == '__main__':
    parser = argparse.ArgumentParser(description='Process classification images and output predictions.')
    parser.add_argument('input_path', help='Path to the input folder containing Positive and Negative subfolders')
    args = parser.parse_args()
    main(args.input_path)
