"""

print(f"Loading dataset from: {data_path}")

# Load the parquet file

df = pd.read_parquet(data_path)

n_samples = len(df)

print(f"Dataset contains {n_samples:,} samples")

print(f"Dataset columns: {list(df.columns)}")

# Try to find the MoA labels column

moa_column = None

possible_moa_columns = ['moa-fine', 'moa-fine-encoded', 'moa', 'moa_fine', 'labels']

for col in possible_moa_columns:

if col in df.columns:

moa_column = col

break

if moa_column is None:

print("Warning: No MoA column found. Using random stratification.")

# Create dummy labels for stratification

labels = np.random.randint(0, 10, n_samples)

else:

print(f"Using column '{moa_column}' for stratified splitting")

labels = df[moa_column].values

# Handle string labels

if isinstance(labels[0], str):

unique_labels = np.unique(labels)

print(f"Found {len(unique_labels)} unique string labels, converting to integers")

label_to_idx = {label: idx for idx, label in enumerate(unique_labels)}

labels = np.array([label_to_idx[label] for label in labels])

print(f"Label distribution: {np.bincount(labels)}")

train_ratio: float = 0.8, val_ratio: float = 0.1, test_ratio: float = 0.1,

random_state: int = 42) -> Dict[str, List[int]]:

"""

# Validate ratios

total_ratio = train_ratio + val_ratio + test_ratio

if abs(total_ratio - 1.0) > 1e-6:

raise ValueError(f"Ratios must sum to 1.0, got {total_ratio}")

print(f"Creating splits with ratios - Train: {train_ratio}, Val: {val_ratio}, Test: {test_ratio}")

# Create index array

indices = np.arange(n_samples)

# First split: separate test set

if test_ratio > 0:

train_val_indices, test_indices = train_test_split(

indices,

test_size=test_ratio,

random_state=random_state,

stratify=labels

)

else:

train_val_indices = indices

test_indices = np.array([])

# Second split: separate train and validation from remaining data

if val_ratio > 0 and len(train_val_indices) > 0:

# Calculate validation ratio relative to remaining data

remaining_ratio = train_ratio + val_ratio

val_ratio_adjusted = val_ratio / remaining_ratio

train_indices, val_indices = train_test_split(

train_val_indices,

test_size=val_ratio_adjusted,

random_state=random_state,

stratify=labels[train_val_indices]

)

else:

train_indices = train_val_indices

val_indices = np.array([])

# Create the splits dictionary

splits = {

'train': train_indices.tolist(),

'val': val_indices.tolist(),

'test': test_indices.tolist()

}

# Print split statistics

print(f"Split sizes:")

print(f" Train: {len(splits['train'])} samples ({len(splits['train'])/n_samples:.1%})")

print(f" Val: {len(splits['val'])} samples ({len(splits['val'])/n_samples:.1%})")

print(f" Test: {len(splits['test'])} samples ({len(splits['test'])/n_samples:.1%})")

# Verify no overlap and all samples covered

all_indices = set(splits['train'] + splits['val'] + splits['test'])

assert len(all_indices) == n_samples, "Not all samples are covered by splits"

assert len(all_indices) == len(splits['train']) + len(splits['val']) + len(splits['test']), "Overlapping indices detected"

print("✓ Split validation passed")

metadata: Dict = None) -> None:

"""

output_path = Path(output_path)

output_path.parent.mkdir(parents=True, exist_ok=True)

# Prepare the data to save

data_to_save = {

'splits': splits,

'metadata': metadata or {}

}

# Add generation timestamp

from datetime import datetime

data_to_save['metadata']['generated_at'] = datetime.now().isoformat()

data_to_save['metadata']['total_samples'] = len(splits['train']) + len(splits['val']) + len(splits['test'])

# Save to JSON

with open(output_path, 'w') as f:

json.dump(data_to_save, f, indent=2)

parser = argparse.ArgumentParser(description="Generate train/val/test splits for PRISM datasets")

# Input options

group = parser.add_mutually_exclusive_group(required=True)

group.add_argument('--config', type=str, help='Path to PRISM config file')

group.add_argument('--data-path', type=str, help='Direct path to dataset parquet file')

# Output options

parser.add_argument('--output', type=str, help='Output directory or file path')

parser.add_argument('--filename', type=str, default='train_val_test_splits.json',

help='Output filename (default: train_val_test_splits.json)')

# Split ratio options

parser.add_argument('--train-ratio', type=float, default=0.8, help='Training set ratio (default: 0.8)')

parser.add_argument('--val-ratio', type=float, default=0.1, help='Validation set ratio (default: 0.1)')

parser.add_argument('--test-ratio', type=float, default=0.1, help='Test set ratio (default: 0.1)')

# Other options

parser.add_argument('--random-seed', type=int, default=42, help='Random seed for reproducibility (default: 42)')

parser.add_argument('--force', action='store_true', help='Overwrite existing split files')

args = parser.parse_args()

# Determine data path

if args.config:

if Config is None:

print("Error: Cannot load config file without prism.configs module")

sys.exit(1)

config = Config(config_files=[args.config])

data_path = config.get('data.processed_parquet_path') or config.get('data.raw_parquet_path')

if not data_path:

print("Error: No data path found in config file")

sys.exit(1)

# Determine output path from config if not specified

if not args.output:

args.output = config.get('data.splits_path', 'data/splits/train_val_test_splits.json')

else:

data_path = args.data_path

# Determine output path if not specified

if not args.output:

data_dir = Path(data_path).parent

args.output = data_dir / 'splits' / args.filename

# Ensure data file exists

if not Path(data_path).exists():

print(f"Error: Data file not found: {data_path}")

sys.exit(1)

# Determine final output path

output_path = Path(args.output)

if output_path.is_dir():

output_path = output_path / args.filename

# Check if output file already exists

if output_path.exists() and not args.force:

print(f"Error: Output file already exists: {output_path}")

print("Use --force to overwrite")

sys.exit(1)

print(f"=== PRISM Dataset Split Generator ===")

print(f"Data source: {data_path}")

print(f"Output: {output_path}")

print(f"Split ratios: {args.train_ratio:.1%} train, {args.val_ratio:.1%} val, {args.test_ratio:.1%} test")

print(f"Random seed: {args.random_seed}")

print()

try:

# Load dataset information

n_samples, labels = load_dataset_info(data_path)

# Create splits

splits = create_stratified_splits(

n_samples=n_samples,

labels=labels,

train_ratio=args.train_ratio,

val_ratio=args.val_ratio,

test_ratio=args.test_ratio,

random_state=args.random_seed

)

# Prepare metadata

metadata = {

'source_file': str(data_path),

'random_seed': args.random_seed,

'train_ratio': args.train_ratio,

'val_ratio': args.val_ratio,

'test_ratio': args.test_ratio,

'n_samples': n_samples,

'n_unique_labels': len(np.unique(labels))

}

# Save splits (but extract just the splits for DataPreparer compatibility)

splits_only = {

'train': splits['train'],

'val': splits['val'],

'test': splits['test']

}

# Save the splits in DataPreparer-compatible format

with open(output_path, 'w') as f:

json.dump(splits_only, f, indent=2)

print(f"\n✓ Splits successfully generated and saved to: {output_path}")

print(f"\nTo use these splits with DataPreparer:")

print(f"1. Ensure your config file has:")

print(f" data.splits_path: \"{output_path}\"")

print(f"2. The DataPreparer will automatically load these splits when prepare_train_val_test() is called")

# Also save a detailed version with metadata

metadata_output = output_path.with_suffix('.metadata.json')

save_splits_json(splits, metadata_output, metadata)

print(f"3. Detailed metadata saved to: {metadata_output}")

except Exception as e:

print(f"Error: {e}")