quick_benchmark.py

#!/usr/bin/env python3
"""
Quick benchmark script - test detectors with comprehensive metrics.

Usage:
    python quick_benchmark.py --data-dir ./data
    python quick_benchmark.py --data-dir ./data --samples 1000
    python quick_benchmark.py --data-dir ./data --top 5
"""

import os
import sys
import subprocess
import argparse
from pathlib import Path
import time
import json
import warnings
warnings.filterwarnings('ignore')

import pandas as pd
import numpy as np
from sklearn.metrics import (
    roc_auc_score, accuracy_score, recall_score, f1_score, confusion_matrix
)


# Recommended detectors by expected performance
RECOMMENDED_ORDER = [
    'meta_stacking_7models',
    'mlp_ensemble_deep_features',
    'xgb_70_statistical',
    'stacking_200_features',
    'quad_model_ensemble',
    'xgb_tuned_regularization',
    'kolmogorov_smirnov_xgb',
    'wavelet_lstm',
    'knn_wavelet',
    'hypothesis_testing_pure',
]


def calculate_all_metrics(y_true, y_prob, threshold=0.5):
    """Calculate metrics: ROC AUC, F1, Accuracy, Recall, TP, FP, TN, FN."""
    y_true = np.array(y_true).astype(int)
    y_prob = np.array(y_prob)
    y_pred = (y_prob >= threshold).astype(int)

    metrics = {}

    # Confusion Matrix
    tn, fp, fn, tp = confusion_matrix(y_true, y_pred, labels=[0, 1]).ravel()
    metrics['TP'] = int(tp)
    metrics['FP'] = int(fp)
    metrics['TN'] = int(tn)
    metrics['FN'] = int(fn)

    # Core metrics
    metrics['accuracy'] = round(accuracy_score(y_true, y_pred), 4)
    metrics['recall'] = round(recall_score(y_true, y_pred, zero_division=0), 4)
    metrics['f1_score'] = round(f1_score(y_true, y_pred, zero_division=0), 4)

    # ROC AUC
    try:
        metrics['roc_auc'] = round(roc_auc_score(y_true, y_prob), 4)
    except:
        metrics['roc_auc'] = None

    return metrics


def run_single_detector(detector_name: str, base_dir: str, data_dir: str,
                        max_samples: int = 2000) -> dict:
    """Run a single detector and return comprehensive results."""

    detector_dir = os.path.join(base_dir, detector_name)
    model_path = os.path.join(detector_dir, 'model.joblib')

    result = {
        'detector': detector_name,
        'status': 'pending'
    }

    try:
        start = time.time()

        # Train
        train_cmd = [
            sys.executable,
            os.path.join(detector_dir, 'main.py'),
            '--mode', 'train',
            '--data-dir', data_dir,
            '--model-path', model_path,
            '--max-samples', str(max_samples)
        ]

        print(f"  Training...", end=' ', flush=True)
        train_result = subprocess.run(
            train_cmd,
            capture_output=True,
            text=True,
            timeout=600,
            cwd=detector_dir
        )

        if train_result.returncode != 0:
            print(f"FAILED")
            result['status'] = 'train_error'
            result['error'] = train_result.stderr[:200]
            return result

        train_time = time.time() - start
        result['train_time'] = round(train_time, 1)

        # Get predictions via infer
        print(f"Evaluating...", end=' ', flush=True)
        eval_start = time.time()

        # Load test data
        y_test = pd.read_parquet(f"{data_dir}/y_test.reduced.parquet").squeeze()

        # Run inference and capture predictions
        infer_script = f'''
import sys
sys.path.insert(0, "{detector_dir}")
import pandas as pd
import json

from main import infer

predictions = []
series_ids = []

for series_id, prob in infer("{data_dir}", "{model_path}"):
    predictions.append(float(prob))
    series_ids.append(series_id)

print(json.dumps({{"predictions": predictions, "series_ids": [str(s) for s in series_ids]}}))
'''

        infer_result = subprocess.run(
            [sys.executable, '-c', infer_script],
            capture_output=True,
            text=True,
            timeout=300,
            cwd=detector_dir
        )

        if infer_result.returncode != 0:
            print(f"EVAL FAILED")
            result['status'] = 'eval_error'
            result['error'] = infer_result.stderr[:200]
            return result

        # Parse predictions
        try:
            output_lines = infer_result.stdout.strip().split('\n')
            json_line = [l for l in output_lines if l.startswith('{')][-1]
            pred_data = json.loads(json_line)
            predictions = pred_data['predictions']
            series_ids = pred_data['series_ids']
        except Exception as e:
            print(f"PARSE FAILED")
            result['status'] = 'parse_error'
            result['error'] = str(e)
            return result

        # Get true labels
        true_labels = []
        for sid in series_ids:
            try:
                true_labels.append(y_test.loc[int(sid)] if sid.isdigit() else y_test.loc[sid])
            except:
                true_labels.append(y_test.iloc[0])  # Fallback

        # Calculate metrics
        metrics = calculate_all_metrics(true_labels, predictions)
        result.update(metrics)

        eval_time = time.time() - eval_start
        result['eval_time'] = round(eval_time, 1)
        result['total_time'] = round(time.time() - start, 1)
        result['status'] = 'success'

        print(f"ROC AUC: {result['roc_auc']:.4f}, F1: {result['f1_score']:.4f}, "
              f"TP:{result['TP']} FP:{result['FP']} TN:{result['TN']} FN:{result['FN']} "
              f"({result['total_time']:.1f}s)")

    except subprocess.TimeoutExpired:
        print("TIMEOUT")
        result['status'] = 'timeout'
    except Exception as e:
        print(f"ERROR: {e}")
        result['status'] = 'error'
        result['error'] = str(e)

    return result


def print_results_table(results):
    """Print results table with ROC AUC, F1, Accuracy, Recall, TP, FP, TN, FN."""

    print("\n" + "="*110)
    print("RESULTS")
    print("="*110)

    # Sort by ROC AUC
    results_sorted = sorted(results, key=lambda x: x.get('roc_auc') or 0, reverse=True)

    # Results table
    print(f"\n{'Rank':<4} {'Detector':<32} {'ROC AUC':<9} {'F1':<9} {'Accuracy':<10} {'Recall':<9} {'TP':<6} {'FP':<6} {'TN':<6} {'FN':<6}")
    print("-"*110)

    for i, r in enumerate(results_sorted, 1):
        if r['status'] != 'success':
            print(f"{i:<4} {r['detector']:<32} {'N/A':<9} {'N/A':<9} {'N/A':<10} {'N/A':<9} {'N/A':<6} {'N/A':<6} {'N/A':<6} {r['status']}")
            continue

        roc = f"{r['roc_auc']:.4f}" if r.get('roc_auc') else "N/A"
        f1 = f"{r['f1_score']:.4f}" if r.get('f1_score') else "N/A"
        acc = f"{r['accuracy']:.4f}" if r.get('accuracy') else "N/A"
        rec = f"{r['recall']:.4f}" if r.get('recall') else "N/A"
        tp = str(r.get('TP', 'N/A'))
        fp = str(r.get('FP', 'N/A'))
        tn = str(r.get('TN', 'N/A'))
        fn = str(r.get('FN', 'N/A'))
        print(f"{i:<4} {r['detector']:<32} {roc:<9} {f1:<9} {acc:<10} {rec:<9} {tp:<6} {fp:<6} {tn:<6} {fn:<6}")

    # Summary
    successful = [r for r in results if r['status'] == 'success']
    if successful:
        print("\n--- Best Performers ---")
        best_roc = max(successful, key=lambda x: x.get('roc_auc') or 0)
        best_f1 = max(successful, key=lambda x: x.get('f1_score') or 0)

        print(f"Best ROC AUC:  {best_roc['detector']} ({best_roc['roc_auc']:.4f})")
        print(f"Best F1 Score: {best_f1['detector']} ({best_f1['f1_score']:.4f})")

        total_time = sum(r.get('total_time', 0) for r in successful)
        print(f"Total time:    {total_time:.1f}s")


def save_results(results, output_path):
    """Save results to CSV."""
    df = pd.DataFrame(results)

    # Reorder columns
    priority_cols = [
        'detector', 'status', 'roc_auc', 'f1_score', 'accuracy', 'recall',
        'TP', 'FP', 'TN', 'FN',
        'train_time', 'eval_time', 'total_time', 'error'
    ]

    cols = [c for c in priority_cols if c in df.columns]
    cols += [c for c in df.columns if c not in cols]
    df = df[cols]

    df.to_csv(output_path, index=False)
    print(f"\nResults saved to: {output_path}")


def main():
    parser = argparse.ArgumentParser(description="Quick benchmark for detectors")
    parser.add_argument("--data-dir", required=True, help="Data directory")
    parser.add_argument("--samples", type=int, default=2000, help="Training samples")
    parser.add_argument("--top", type=int, default=None, help="Only run top N recommended")
    parser.add_argument("--detectors", nargs="+", help="Specific detectors to run")
    parser.add_argument("--all", action="store_true", help="Run all detectors")
    parser.add_argument("--output", default="quick_results.csv", help="Output CSV file")

    args = parser.parse_args()

    base_dir = os.path.dirname(os.path.abspath(__file__))
    data_dir = os.path.abspath(args.data_dir)

    # Determine detectors to run
    if args.detectors:
        detectors = args.detectors
    elif args.all:
        detectors = sorted([d for d in os.listdir(base_dir)
                           if os.path.isdir(os.path.join(base_dir, d))
                           and os.path.exists(os.path.join(base_dir, d, 'main.py'))])
    elif args.top:
        detectors = RECOMMENDED_ORDER[:args.top]
    else:
        detectors = RECOMMENDED_ORDER[:5]

    print("="*80)
    print("QUICK BENCHMARK - Comprehensive Metrics")
    print("="*80)
    print(f"Data dir: {data_dir}")
    print(f"Samples: {args.samples}")
    print(f"Detectors: {len(detectors)}")
    print("="*80)

    results = []
    for i, detector in enumerate(detectors, 1):
        print(f"\n[{i}/{len(detectors)}] {detector}")
        result = run_single_detector(detector, base_dir, data_dir, args.samples)
        results.append(result)

    # Print comprehensive results
    print_results_table(results)

    # Save results
    output_path = os.path.join(base_dir, args.output)
    save_results(results, output_path)


if __name__ == "__main__":
    main()