dataset_processor.py

import os
import shutil
import signal
import traceback
from data_utils import find_fold_files, parse_keel_dat, prepare_datasets, normalize_name
from feature_selection import compute_po_statistic, genetic_algorithm
from evaluation import evaluate_test_performance
from config import BASE_DIR, DATASET_TIMEOUT


def timeout_handler(signum, frame):
    raise TimeoutError("Dataset processing timed out")


def process_dataset(folder_name, idx, total_datasets):
    signal.signal(signal.SIGALRM, timeout_handler)
    signal.alarm(DATASET_TIMEOUT)

    try:
        folder_path = os.path.join(BASE_DIR, folder_name)
        if not os.path.isdir(folder_path):
            return folder_name, "N/A"

        print(f"\n[{idx}/{total_datasets}] Processing dataset: {folder_name}")

        fold_zip, train_file, test_file = find_fold_files(folder_path)

        if not fold_zip or not train_file or not test_file:
            reason = "No -fold.zip found" if not fold_zip else "Missing tra or tst file"
            print(f"  Skipping: {reason}")
            if fold_zip:
                extract_dir = os.path.join(folder_path, "extracted_folds")
                shutil.rmtree(extract_dir, ignore_errors=True)
            return folder_name, "N/A"

        dataset_name = normalize_name(folder_name)
        df_train, class_col = parse_keel_dat(train_file)
        df_test, _ = parse_keel_dat(test_file)

        X_train, y_train, X_test, y_test, numeric_cols, target_class = prepare_datasets(
            df_train, df_test, class_col
        )

        print(f"  Dataset shape: {X_train.shape}, Target class: {target_class}")
        print(f"  Class distribution - Train: {sum(y_train)}/{len(y_train)}, "
              f"Test: {sum(y_test)}/{len(y_test)}")

        print("  Computing PO statistic...")
        feature_probs = compute_po_statistic(X_train, y_train)

        print("  Running Genetic Algorithm...")
        selected_features = genetic_algorithm(X_train, y_train, feature_probs, numeric_cols)

        auc = evaluate_test_performance(
            X_train, y_train, X_test, y_test, selected_features, numeric_cols
        )

        if isinstance(auc, float):
            n_selected = sum(selected_features)
            n_total = len(selected_features)
            print(f"  Test AUC: {auc:.4f} with {n_selected}/{n_total} features")

        print(f"  Completed {dataset_name} with AUC: {auc}")

        extract_dir = os.path.join(folder_path, "extracted_folds")
        shutil.rmtree(extract_dir, ignore_errors=True)

        return dataset_name, auc

    except TimeoutError:
        print(f"  Dataset {folder_name} timed out after {DATASET_TIMEOUT}s")
        extract_dir = os.path.join(folder_path, "extracted_folds")
        shutil.rmtree(extract_dir, ignore_errors=True)
        return folder_name, "Timeout"

    except Exception as e:
        print(f"  Error in {folder_name}: {str(e)}")
        traceback.print_exc()
        return folder_name, f"Error: {type(e).__name__}"

    finally:
        signal.alarm(0)


def process_dataset_wrapper(args):
    folder_name, idx, total_datasets = args
    return process_dataset(folder_name, idx, total_datasets)