time_algorithm.py

from io import StringIO
import itertools
import sys
import time
import warnings
from typing import Any, Dict, List, Optional, Tuple

import numpy as np
import pandas as pd
from scipy.interpolate import LinearNDInterpolator, NearestNDInterpolator
import tqdm

import race_ordering
import race_selection

# =====================================================================
# 1. GRID CONFIGURATION & BENCHMARKING ENGINE
# =====================================================================
GRID_RANGES: Dict[str, List[Any]] = {
    "p_count": [5, 8, 12, 15, 20, 25],
    "n_races": [3, 5, 7],
    "strict_R": [False, True],
    "strictness_tolerance": [1, 3, 6],
    "cuts_per_rotation": [4, 0, 1, 2, 3],
    "num_stencils": [20, 0, 10, 30, 50],
    "pool_size": [3, 0, 1, 2, 5, 10],
    "restarts": [10, 0, 5, 20],
    "gap_variance_weight": [0.1],
    "steps": [10000, 1000, 2000, 3000, 5000],
    "num_climbers": [3000, 10000, 20000, 50000],
}

SEEDS: List[int] = [42, 123, 999]

NUMERIC_FEATURES: List[str] = [
    "p_count",
    "n_races",
    "strict_R",
    "strictness_tolerance",
    "cuts_per_rotation",
    "num_stencils",
    "pool_size",
    "restarts",
    "gap_variance_weight",
    "steps",
    "num_climbers",
]
NUMERIC_FEATURES.reverse()


# Non-loggable defaults:
race_selection.ILP_MAX_TIME_LIMIT = 5.0
race_ordering.MILP_TIMEOUT = 60.0


class NullIO(StringIO):
    def write(self, txt):
        pass


def execute_single_run(config: Dict[str, Any], seed: int) -> float:
    """Executes a single algorithm run and returns execution time in seconds."""
    p_count = config["p_count"]
    n_races = config["n_races"]

    kwargs_selection = {
        "strict_R": config["strict_R"],
        "strictness_tolerance": config["strictness_tolerance"],
        "cuts_per_rotation": config["cuts_per_rotation"],
        "num_stencils": config["num_stencils"],
        "pool_size": config["pool_size"],
        "restarts": config["restarts"],
        "seed": seed,
        "gap_variance_weight": config["gap_variance_weight"],
    }

    kwargs_ordering = {
        "steps": config["steps"],
        "num_climbers": config["num_climbers"],
    }

    start_time = time.perf_counter()

    # Step 1: Race selection
    A, variance, C, super_pools = race_selection.get_best_unordered_races(
        p_count, n_races, **kwargs_selection
    )
    # Step 2: Race ordering
    permutation_order = race_ordering.find_best_race_order(
        seed, A, n_races, **kwargs_ordering
    )
    _ = race_ordering.permute_incidence_matrix(A, permutation_order)

    return time.perf_counter() - start_time


def generate_grid_configurations(grid: Dict[str, List[Any]]) -> List[Dict[str, Any]]:
    """Generates valid hyperparameter combinations, bypassing redundant variations."""
    keys = list(grid.keys())
    values = list(grid.values())
    configs = []

    for combination in itertools.product(*values):
        config = dict(zip(keys, combination))
        # Skip redundant tolerance variations when strict_R is True
        if config["strict_R"] and config["strictness_tolerance"] != grid["strictness_tolerance"][0]:
            continue
        configs.append(config)

    return configs


def _compile_dataframe_from_map(data_map: Dict[Tuple, Dict[str, Any]]) -> pd.DataFrame:
    """Converts the internal tracking dictionary into a formatted DataFrame."""
    records = []
    for entry in data_map.values():
        cfg = entry["config"]
        runtimes = entry["runtimes"]
        valid_times = [t for t in runtimes if pd.notna(t)]
        errors = sum(1 for t in runtimes if pd.isna(t))
        mean_time = float(np.mean(valid_times)) if valid_times else np.nan
        success_rate = (len(valid_times) / len(runtimes)) if runtimes else np.nan

        records.append({
            **cfg,
            "runtimes": runtimes,
            "mean_runtime": mean_time,
            "error_count": errors,
            "success_rate": success_rate,
        })
    return pd.DataFrame(records)


def run_benchmark_grid(
    grid_ranges: Dict[str, List[Any]] = GRID_RANGES,
    seeds: List[int] = SEEDS,
    checkpoint_filepath: str = "benchmark_checkpoint.parquet",
    checkpoint_interval: int = 100,
    existing_df: Optional[pd.DataFrame] = None,
) -> pd.DataFrame:
    """Runs grid search benchmark with outer seed iteration and periodic checkpointing.

    Accepts an existing DataFrame to append/resume runs.
    """
    configs = generate_grid_configurations(grid_ranges)
    config_keys = [tuple(cfg[k] for k in NUMERIC_FEATURES) for cfg in configs]

    # Map: tuple_key -> {"config": cfg_dict, "runtimes": []}
    data_map: Dict[Tuple, Dict[str, Any]] = {}

    # Load state from existing DataFrame if provided
    if existing_df is not None and not existing_df.empty:
        for _, row in existing_df.iterrows():
            key = tuple(row[k] for k in NUMERIC_FEATURES)
            r_list = list(row["runtimes"]) if isinstance(row["runtimes"], (list, np.ndarray)) else []
            data_map[key] = {
                "config": {k: row[k] for k in NUMERIC_FEATURES},
                "runtimes": r_list,
            }

    # Initialize unvisited configs
    for key, cfg in zip(config_keys, configs):
        if key not in data_map:
            data_map[key] = {"config": cfg, "runtimes": []}

    total_evaluations = len(seeds) * len(configs)
    completed_evaluations = 0

    print(f"Starting grid search: {len(seeds)} seeds x {len(configs)} configs = {total_evaluations} total runs.")
    pbar = tqdm.tqdm(total=total_evaluations, desc="Benchmarking")

    for cfg, key in zip(configs, config_keys):
        for seed in seeds:
            # Suppress stdout during run
            sys.stdout = NullIO()
            try:
                elapsed = execute_single_run(cfg, seed)
                data_map[key]["runtimes"].append(elapsed)
            except Exception:
                data_map[key]["runtimes"].append(np.nan)
            finally:
                sys.stdout = sys.__stdout__

            completed_evaluations += 1
            pbar.update(1)

            # Save Checkpoint every checkpoint_interval evaluations
            if completed_evaluations % checkpoint_interval == 0:
                current_df = _compile_dataframe_from_map(data_map)
                save_benchmark_results(current_df, checkpoint_filepath)

    pbar.close()
    final_df = _compile_dataframe_from_map(data_map)
    save_benchmark_results(final_df, checkpoint_filepath)
    return final_df


# =====================================================================
# 2. DATAFRAME MERGING & PERSISTENCE
# =====================================================================
def merge_benchmark_dfs(df1: pd.DataFrame, df2: pd.DataFrame) -> pd.DataFrame:
    """Combines two benchmark DataFrames by merging runtime lists across identical hyperparameter configurations."""
    if df1 is None or df1.empty:
        return df2.copy() if df2 is not None else pd.DataFrame()
    if df2 is None or df2.empty:
        return df1.copy()

    combined = pd.concat([df1, df2], ignore_index=True)
    grouped_records = []

    for _, group in combined.groupby(NUMERIC_FEATURES, as_index=False):
        all_runtimes = []
        for r_item in group["runtimes"]:
            if isinstance(r_item, (list, np.ndarray)):
                all_runtimes.extend(r_item)
            elif pd.notna(r_item):
                all_runtimes.append(r_item)

        valid_times = [t for t in all_runtimes if pd.notna(t)]
        errors = sum(1 for t in all_runtimes if pd.isna(t))
        mean_time = float(np.mean(valid_times)) if valid_times else np.nan
        success_rate = (len(valid_times) / len(all_runtimes)) if all_runtimes else np.nan

        cfg_dict = group.iloc[0][NUMERIC_FEATURES].to_dict()
        cfg_dict["runtimes"] = all_runtimes
        cfg_dict["mean_runtime"] = mean_time
        cfg_dict["error_count"] = errors
        cfg_dict["success_rate"] = success_rate

        grouped_records.append(cfg_dict)

    return pd.DataFrame(grouped_records)


def save_benchmark_results(df: pd.DataFrame, filepath: str = "benchmark_results.parquet") -> None:
    """Saves the benchmark DataFrame to disk using Parquet compression."""
    df_to_save = df.copy()
    df_to_save["strict_R"] = df_to_save["strict_R"].astype(int)
    df_to_save.to_parquet(filepath, compression="snappy", engine="pyarrow")


def load_benchmark_results(filepath: str = "benchmark_results.parquet") -> pd.DataFrame:
    """Loads benchmark DataFrame from disk."""
    df = pd.read_parquet(filepath, engine="pyarrow")
    df["strict_R"] = df["strict_R"].astype(bool)
    return df


# =====================================================================
# 3. INTERPOLATION & EXTRAPOLATION ENGINE
# =====================================================================
class RaceBenchmarkPredictor:
    """Estimates execution time for unseen points using linear interpolation with nearest-neighbor fallback for extrapolation."""

    def __init__(self, benchmark_df: pd.DataFrame, feature_cols: List[str] = NUMERIC_FEATURES):
        self.feature_cols = feature_cols
        self.df = benchmark_df.dropna(subset=["mean_runtime"]).copy()
        self.df["strict_R"] = self.df["strict_R"].astype(float)
        self.X = self.df[self.feature_cols].values.astype(float)
        self.y = self.df["mean_runtime"].values.astype(float)

        # Primary ND Linear Interpolator
        self.linear_interp = LinearNDInterpolator(self.X, self.y)
        # Fallback ND Nearest Interpolator (for extrapolation outside convex hull)
        self.nearest_interp = NearestNDInterpolator(self.X, self.y)

    def predict(self, point: Dict[str, Any]) -> float:
        """Predicts estimated runtime (seconds) for a given hyperparameter dictionary."""
        vector = []
        for col in self.feature_cols:
            val = point.get(col, 0)
            if isinstance(val, bool):
                val = float(val)
            vector.append(float(val))

        target_point = np.array([vector])
        pred = self.linear_interp(target_point)[0]

        if np.isnan(pred):
            warnings.warn("Point is outside known grid scope. Using extrapolation (nearest neighbor).")
            pred = self.nearest_interp(target_point)[0]

        return float(pred)


# =====================================================================
# 4. EXECUTION EXAMPLE
# =====================================================================
if __name__ == "__main__":
    # 1. Run grid search (will save checkpoints every 2500 runs to 'benchmark_checkpoint.parquet')
    results_df = run_benchmark_grid()
    print("Finished benchmarking.")

    # 2. Example: Merge newly collected data with an existing dataset
    # loaded_df = load_benchmark_results("previous_results.parquet")
    # merged_df = merge_benchmark_dfs(loaded_df, results_df)

    # 3. Predict runtime for a query point
    predictor = RaceBenchmarkPredictor(results_df)
    sample_query = {
        "p_count": 10,
        "n_races": 4,
        "strict_R": False,
        "strictness_tolerance": 2,
        "cuts_per_rotation": 2,
        "num_stencils": 15,
        "pool_size": 4,
        "restarts": 8,
        "gap_variance_weight": 0.1,
        "steps": 2500,
        "num_climbers": 7500,
    }

    predicted_time = predictor.predict(sample_query)
    print(f"\nEstimated Runtime for query point: {predicted_time:.4f} seconds")