main.py

"""
Evolution Feature Extraction System - Main Runner
"""

import os
import time
import multiprocessing
from tqdm import tqdm
from config import BASE_DIR, MAX_THREADS
from dataset_processor import process_dataset, process_dataset_wrapper


def display_results(results):
    print("\n=== Final Results ===")
    print(f"{'Dataset':<40} {'AUC':>10}")
    print("-" * 50)

    for name, auc in sorted(results.items()):
        auc_str = f"{auc:.4f}" if isinstance(auc, float) else auc
        print(f"{name:<40} {auc_str:>10}")

    valid_aucs = [auc for auc in results.values() if isinstance(auc, float)]
    failed_datasets = [(name, error) for name, error in results.items()
                      if not isinstance(error, float)]

    if valid_aucs:
        print(f"\nStatistics:")
        print(f"Average AUC: {sum(valid_aucs)/len(valid_aucs):.4f}")
        print(f"Min AUC: {min(valid_aucs):.4f}")
        print(f"Max AUC: {max(valid_aucs):.4f}")
        print(f"Successfully processed: {len(valid_aucs)}/{len(results)} datasets")

        if failed_datasets:
            print(f"\nFailed datasets ({len(failed_datasets)}):")
            for name, error in failed_datasets:
                print(f"  - {name}: {error}")
        else:
            print("\nAll datasets processed successfully!")
    else:
        print("\nNo datasets were successfully processed.")


def process_datasets_parallel():
    dataset_folders = [d for d in os.listdir(BASE_DIR)
                      if os.path.isdir(os.path.join(BASE_DIR, d))]
    total_datasets = len(dataset_folders)

    print(f"Processing {total_datasets} datasets with {MAX_THREADS} threads")

    try:
        multiprocessing.set_start_method('spawn', force=True)
    except RuntimeError:
        print("Spawn method already set, continuing")

    results = {}

    try:
        with multiprocessing.Pool(MAX_THREADS) as pool:
            tasks = [(folder, idx + 1, total_datasets)
                    for idx, folder in enumerate(dataset_folders)]

            for dataset_name, auc in tqdm(
                pool.imap(process_dataset_wrapper, tasks),
                total=total_datasets,
                desc="Datasets"
            ):
                if dataset_name in results:
                    print(f"  Warning: Duplicate dataset name {dataset_name}, "
                          f"overwriting previous result")
                results[dataset_name] = auc

    except Exception as e:
        print(f"Error in multiprocessing pool: {str(e)}")
        print("Falling back to sequential processing")
        results = process_datasets_sequential(dataset_folders)

    return results


def process_datasets_sequential(dataset_folders=None):
    if dataset_folders is None:
        dataset_folders = [d for d in os.listdir(BASE_DIR)
                          if os.path.isdir(os.path.join(BASE_DIR, d))]

    total_datasets = len(dataset_folders)
    results = {}

    for idx, folder in enumerate(tqdm(dataset_folders, desc="Datasets"), 1):
        dataset_name, auc = process_dataset(folder, idx, total_datasets)
        results[dataset_name] = auc

    return results


def main():
    print("Evolution Feature Extraction System")
    print("=" * 50)
    print(f"Base directory: {BASE_DIR}")
    print(f"Maximum threads: {MAX_THREADS}")
    print()

    start_time = time.time()

    results = process_datasets_parallel()
    display_results(results)

    elapsed = time.time() - start_time
    print(f"\nTotal execution time: {elapsed:.2f} seconds ({elapsed/60:.2f} minutes)")
    print("Process complete.")


if __name__ == "__main__":
    main()