create_batch_datasets.py

import json
import csv
import sys
import random
from pathlib import Path
from collections import defaultdict

SOURCE_JSON_PATH = Path('datasets/Syn-STARTS_datasets/all_scenarios.json')
OUTPUT_DIR_PATH = SOURCE_JSON_PATH.parent / 'batch'

RANDOM_SEED = 42

def save_to_csv(data_to_write, output_file_path):
    csv_headers = ['question', 'triage_zone']
    processed_data_for_csv = []
    for item in data_to_write:
        if 'patient_description' in item and 'triage_tag' in item:
            processed_data_for_csv.append({
                'question': item['patient_description'],
                'triage_zone': item['triage_tag']
            })
    
    if not processed_data_for_csv:
        print(f"⚠️  Warning: No valid data to write to '{output_file_path}'")
        return

    try:
        output_file_path.parent.mkdir(parents=True, exist_ok=True)
        with open(output_file_path, 'w', newline='', encoding='utf-8') as csv_file:
            writer = csv.DictWriter(csv_file, fieldnames=csv_headers)
            writer.writeheader()
            writer.writerows(processed_data_for_csv)
        print(f"✅ Successfully created CSV file: '{output_file_path}'")
    except IOError as e:
        print(f"❌ Error: Could not write to CSV file at '{output_file_path}': {e}")


def main():
    print("🚀 Starting dataset generation...")
    random.seed(RANDOM_SEED)

    try:
        with open(SOURCE_JSON_PATH, 'r', encoding='utf-8') as json_file:
            all_data = json.load(json_file)
    except FileNotFoundError:
        print(f"❌ Error: Input file not found at '{SOURCE_JSON_PATH}'")
        sys.exit(1)
    except json.JSONDecodeError:
        print(f"❌ Error: Could not decode JSON from '{SOURCE_JSON_PATH}'. Check format.")
        sys.exit(1)

    data_by_tag = defaultdict(list)
    for record in all_data:
        raw_tag = record.get('triage_tag')
        if raw_tag and isinstance(raw_tag, str):
            normalized_tag = raw_tag.strip().title()
            data_by_tag[normalized_tag].append(record)
    
    for tag in data_by_tag:
        random.shuffle(data_by_tag[tag])
        print(f"Found {len(data_by_tag[tag])} records for tag: '{tag}' (after normalization)")

    print("\n--- Generating Uniform Distribution Datasets ---")
    uniform_sizes = [12, 20, 56, 100, 200]
    num_sets = 10
    tags_in_order = ['Green', 'Yellow', 'Red', 'Black']

    for size in uniform_sizes:
        print(f"\nProcessing size: {size}")
        if size % len(tags_in_order) != 0:
            print(f"⚠️  Warning: Size {size} is not divisible by the number of tags ({len(tags_in_order)}). Skipping.")
            continue
        
        samples_per_tag = size // len(tags_in_order)
        
        total_required_per_tag = samples_per_tag * num_sets
        can_proceed = True
        for tag in tags_in_order:
            if len(data_by_tag[tag]) < total_required_per_tag:
                print(f"❌ Error: Not enough data for tag '{tag}' to create {num_sets} non-overlapping sets of size {size}.")
                print(f"   (Required: {total_required_per_tag}, Available: {len(data_by_tag[tag])})")
                can_proceed = False
        if not can_proceed:
            continue

        for i in range(1, num_sets + 1):
            current_set_data = []
            for tag in tags_in_order:
                start_index = (i - 1) * samples_per_tag
                end_index = i * samples_per_tag
                current_set_data.extend(data_by_tag[tag][start_index:end_index])

            filename = f"syn_{size}_{i-1}.csv"
            output_path = OUTPUT_DIR_PATH / filename
            save_to_csv(current_set_data, output_path)

    print("\n--- Generating Specific Distribution (TRIAGEa) Datasets ---")
    specific_distribution = {'Green': 18, 'Yellow': 11, 'Red': 22, 'Black': 3}
    
    specific_task_pointers = {tag: 0 for tag in data_by_tag}
    
    for i in range(1, num_sets + 1):
        current_set_data = []
        try:
            for tag, count in specific_distribution.items():
                start_index = specific_task_pointers[tag]
                end_index = start_index + count

                if end_index > len(data_by_tag[tag]):
                    raise ValueError(f"Not enough unique data for tag '{tag}' to create all TRIAGEa sets.")
                
                current_set_data.extend(data_by_tag[tag][start_index:end_index])
                specific_task_pointers[tag] = end_index

            filename = f"syn_TRIAGEa_{i-1}.csv"
            output_path = OUTPUT_DIR_PATH / filename
            save_to_csv(current_set_data, output_path)

        except ValueError as e:
            print(f"❌ Error generating TRIAGEa set {i}: {e}")
            break
    
    print(f"\n✨ All dataset creation complete in directory '{OUTPUT_DIR_PATH}'.")


if __name__ == '__main__':
    main()