ahmet / multilayer-perceptron

import argparse
import sys

import pandas as pd

from train import load_data


def parse_args():
    parser = argparse.ArgumentParser(
        description="Split dataset into training and validation sets"
    )
    parser.add_argument("--dataset", required=True, help="Path to the raw csv file")
    parser.add_argument(
        "--ratio", type=float, default=0.8, help="Train set ratio (default: 0.8)"
    )
    parser.add_argument(
        "--seed", type=int, default=42, help="Random seed for reproducibility"
    )
    parser.add_argument(
        "--train-output",
        default="data_training.csv",
        help="Output path for the training set",
    )
    parser.add_argument(
        "--valid-output",
        default="data_validation.csv",
        help="Output path for the validation set",
    )
    args = parser.parse_args()

    if not 0 < args.ratio < 1:
        parser.error(f"Ratio must be between 0 and 1, got {args.ratio}")
    if args.train_output == args.valid_output:
        parser.error("Train and validation output paths must be different")

    return args


def clean_data(data: pd.DataFrame) -> pd.DataFrame:
    data = data.drop(columns=data.columns[0])
    if data.isna().any().any():
        raise ValueError("Dataset contains missing values")
    labels = data[data.columns[0]]
    if not labels.isin(["M", "B"]).all():
        raise ValueError("Diagnosis column must contain only 'M' or 'B'")
    return data


def stratified_split(data: pd.DataFrame, ratio: float, seed: int):
    label_col = data.columns[0]
    train_parts, valid_parts = [], []
    for label in data[label_col].unique():
        group = data[data[label_col] == label].sample(frac=1, random_state=seed)
        split_idx = int(len(group) * ratio)
        train_parts.append(group.iloc[:split_idx])
        valid_parts.append(group.iloc[split_idx:])
    train = pd.concat(train_parts).sample(frac=1, random_state=seed)
    valid = pd.concat(valid_parts).sample(frac=1, random_state=seed)
    return train.reset_index(drop=True), valid.reset_index(drop=True)


def print_summary(name: str, data: pd.DataFrame):
    counts = data[data.columns[0]].value_counts()
    print(
        f"{name}: {len(data)} samples "
        f"(M: {counts.get('M', 0)}, B: {counts.get('B', 0)})"
    )


def main():
    args = parse_args()
    try:
        data = clean_data(load_data(args.dataset))
    except ValueError as e:
        print(f"Error: {e}", file=sys.stderr)
        sys.exit(1)
    print(f"Loaded {len(data)} samples from {args.dataset}")

    train, valid = stratified_split(data, args.ratio, args.seed)
    train.to_csv(args.train_output, index=False, header=False)
    valid.to_csv(args.valid_output, index=False, header=False)

    print_summary("train", train)
    print_summary("valid", valid)
    print(f"> saved '{args.train_output}' and '{args.valid_output}' to disk")


if __name__ == "__main__":
    main()