

from __future__ import annotations

import argparse
import json
import logging
import math
import time
from collections import defaultdict
from dataclasses import dataclass, field
from pathlib import Path
from typing import Dict, List, Optional, Tuple

import numpy as np

from dpo.core.problem import ContinuousOptimizationProblem
from dpo.core.universal import DPO_Universal, DPO_Presets
from dpo.benchmarks.baselines import PSO, GWO, GA, DE, JADE, WOA, ABC, FA, SA, ACO

# -- simple-hpo-bench ---------------------------------------------------------
try:
    from hpo_benchmarks import NASBench201, HPOBench, HPOLib
    HAS_BENCH = True
except ImportError:
    HAS_BENCH = False
    print("[ERROR] simple-hpo-bench not installed.")
    print("        Run:  pip install simple-hpo-bench")
    raise SystemExit(1)

# -- IOHexperimenter (BBOB benchmarks successor) ------------------------------
try:
    from ioh import get_problem, ProblemClass
    HAS_IOH = True
except ImportError:
    HAS_IOH = False
    print("[WARN] ioh not installed - synthetic benchmarks will be unavailable.")
    print("       Run:  pip install ioh")

# -- matplotlib (optional) ----------------------------------------------------
try:
    import matplotlib
    matplotlib.use("Agg")
    import matplotlib.pyplot as plt
    from matplotlib.ticker import MaxNLocator
    HAS_MPL = True
except ImportError:
    HAS_MPL = False
    print("[WARN] matplotlib not available - all plots will be skipped.")

# -- scipy (optional) ---------------------------------------------------------
try:
    from scipy.stats import wilcoxon, rankdata
    HAS_SCIPY = True
except ImportError:
    HAS_SCIPY = False
    print("[WARN] scipy not available - statistical tests will be skipped.")


# =============================================================================
# Dataset registry
# =============================================================================

@dataclass
class DatasetConfig:
    """Metadata for one benchmark dataset."""
    key:          str          # unique key used throughout
    family:       str          # "nasbench201" | "hpobench" | "hpolib" | "synthetic" | "noisy_synthetic"
    bench_name:   str          # name passed to the benchmark constructor
    display_name: str
    description:  str
    n_dims:       int = 0      # filled in at runtime from search_space
    citation:     str = ""


# --- NASBench201 (3 datasets) ------------------------------------------------
_NAS_CITE = "Dong & Yang, NAS-Bench-201, ICLR 2020"

NASBENCH201_DATASETS: Dict[str, DatasetConfig] = {
    "cifar10": DatasetConfig(
        key="cifar10",  family="nasbench201", bench_name="cifar10",
        display_name="CIFAR-10",
        description="10-class image classification (32×32), 15 625 architectures",
        citation=_NAS_CITE,
    ),
    "cifar100": DatasetConfig(
        key="cifar100", family="nasbench201", bench_name="cifar100",
        display_name="CIFAR-100",
        description="100-class image classification (32×32), 15 625 architectures",
        citation=_NAS_CITE,
    ),
    "imagenet": DatasetConfig(
        key="imagenet", family="nasbench201", bench_name="imagenet",
        display_name="ImageNet-16-120",
        description="120-class downsampled ImageNet (16×16), 15 625 architectures",
        citation=_NAS_CITE,
    ),
}

# --- HPOBench (8 datasets) ---------------------------------------------------
_HPOB_CITE = "Klein & Hutter, Tabular Benchmarks for Joint Architecture and HP Search, 2019"

HPOBENCH_DATASETS: Dict[str, DatasetConfig] = {
    "australian": DatasetConfig(
        key="australian", family="hpobench", bench_name="australian",
        display_name="Australian",
        description="Credit approval (binary classification)",
        citation=_HPOB_CITE,
    ),
    "blood_transfusion": DatasetConfig(
        key="blood_transfusion", family="hpobench", bench_name="blood_transfusion",
        display_name="Blood Transf.",
        description="Blood transfusion service (binary classification)",
        citation=_HPOB_CITE,
    ),
    "car": DatasetConfig(
        key="car", family="hpobench", bench_name="car",
        display_name="Car",
        description="Car evaluation (multi-class classification)",
        citation=_HPOB_CITE,
    ),
    "credit_g": DatasetConfig(
        key="credit_g", family="hpobench", bench_name="credit_g",
        display_name="Credit-G",
        description="German credit (binary classification)",
        citation=_HPOB_CITE,
    ),
    "kc1": DatasetConfig(
        key="kc1", family="hpobench", bench_name="kc1",
        display_name="KC1",
        description="Software defect prediction (binary classification)",
        citation=_HPOB_CITE,
    ),
    "phoneme": DatasetConfig(
        key="phoneme", family="hpobench", bench_name="phoneme",
        display_name="Phoneme",
        description="Spoken vowel classification (binary)",
        citation=_HPOB_CITE,
    ),
    "segment": DatasetConfig(
        key="segment", family="hpobench", bench_name="segment",
        display_name="Segment",
        description="Image segmentation (multi-class classification)",
        citation=_HPOB_CITE,
    ),
    "vehicle": DatasetConfig(
        key="vehicle", family="hpobench", bench_name="vehicle",
        display_name="Vehicle",
        description="Silhouette classification (multi-class)",
        citation=_HPOB_CITE,
    ),
}

# --- HPOLib (4 datasets) -----------------------------------------------------
_HPOL_CITE = "Eggensperger et al., Towards an Empirical Foundation for Assessing HPO, 2013"

HPOLIB_DATASETS: Dict[str, DatasetConfig] = {
    "naval_propulsion": DatasetConfig(
        key="naval_propulsion", family="hpolib", bench_name="naval_propulsion",
        display_name="Naval Propul.",
        description="Naval propulsion regression (minimize MSE)",
        citation=_HPOL_CITE,
    ),
    "parkinsons_telemonitoring": DatasetConfig(
        key="parkinsons_telemonitoring", family="hpolib",
        bench_name="parkinsons_telemonitoring",
        display_name="Parkinsons",
        description="Parkinson's telemonitoring regression (minimize MSE)",
        citation=_HPOL_CITE,
    ),
    "protein_structure": DatasetConfig(
        key="protein_structure", family="hpolib", bench_name="protein_structure",
        display_name="Protein Struct.",
        description="Protein structure regression (minimize MSE)",
        citation=_HPOL_CITE,
    ),
    "slice_localization": DatasetConfig(
        key="slice_localization", family="hpolib", bench_name="slice_localization",
        display_name="Slice Local.",
        description="CT slice localization regression (minimize MSE)",
        citation=_HPOL_CITE,
    ),
}

# --- Synthetic analytic benchmark functions (6 datasets) --------------------
# Using IOHexperimenter BBOB suite (successor to COCO platform)
_SYN_CITE = "Doerr et al., IOHprofiler: A Benchmarking and Profiling Tool for Iterative Optimization Heuristics, arXiv:1810.05281, 2018"

SYNTHETIC_DATASETS: Dict[str, DatasetConfig] = {
    "ackley_20d": DatasetConfig(
        key="ackley_20d",     family="synthetic",  bench_name="bbob_15_1_20",  # Ackley function (f15)
        display_name="Ackley-20D (BBOB)",
        description="BBOB f15 (Ackley, 20D) – highly multimodal, global min at origin",
        citation=_SYN_CITE,
    ),
    "rastrigin_20d": DatasetConfig(
        key="rastrigin_20d",  family="synthetic",  bench_name="bbob_3_1_20",  # Rastrigin function (f3)
        display_name="Rastrigin-20D (BBOB)",
        description="BBOB f3 (Rastrigin, 20D) – ~400 local minima, global min at origin",
        citation=_SYN_CITE,
    ),
    "schwefel_20d": DatasetConfig(
        key="schwefel_20d",   family="synthetic",  bench_name="bbob_20_1_20",  # Schwefel function (f20)
        display_name="Schwefel-20D (BBOB)",
        description="BBOB f20 (Schwefel, 20D) – deceptive landscape, global min near boundary",
        citation=_SYN_CITE,
    ),
    "levy_20d": DatasetConfig(
        key="levy_20d",       family="synthetic",  bench_name="bbob_4_1_20",  # Büche-Rastrigin (f4)
        display_name="Büche-Rastrigin-20D (BBOB)",
        description="BBOB f4 (Büche-Rastrigin, 20D) – several local minima, asymmetric structure",
        citation=_SYN_CITE,
    ),
    "rosenbrock_20d": DatasetConfig(
        key="rosenbrock_20d", family="synthetic",  bench_name="bbob_8_1_20",  # Rosenbrock function (f8)
        display_name="Rosenbrock-20D (BBOB)",
        description="BBOB f8 (Rosenbrock, 20D) – narrow curved valley, global min at origin",
        citation=_SYN_CITE,
    ),
    "griewank_20d": DatasetConfig(
        key="griewank_20d",   family="synthetic",  bench_name="bbob_16_1_20",  # Weierstrass (f16)
        display_name="Weierstrass-20D (BBOB)",
        description="BBOB f16 (Weierstrass, 20D) – dense local-optima grid, highly multimodal",
        citation=_SYN_CITE,
    ),
}

# --- Noisy synthetic benchmarks (3 datasets) -------------------------------
# Using IOHexperimenter BBOB suite with added noise
_NOISY_SYN_CITE = "Doerr et al., IOHprofiler: A Benchmarking and Profiling Tool for Iterative Optimization Heuristics, arXiv:1810.05281, 2018"

NOISY_SYNTHETIC_DATASETS: Dict[str, DatasetConfig] = {
    "noisy_ackley_20d": DatasetConfig(
        key="noisy_ackley_20d", family="noisy_synthetic", bench_name="bbob_15_1_20_noisy",  # Noisy Ackley (f15)
        display_name="Noisy-Ackley-20D (BBOB)",
        description="BBOB f15 (Ackley, 20D) with moderate Gaussian noise",
        citation=_NOISY_SYN_CITE,
    ),
    "noisy_rastrigin_20d": DatasetConfig(
        key="noisy_rastrigin_20d", family="noisy_synthetic", bench_name="bbob_3_1_20_noisy",  # Noisy Rastrigin (f3)
        display_name="Noisy-Rastrigin-20D (BBOB)",
        description="BBOB f3 (Rastrigin, 20D) with moderate Gaussian noise",
        citation=_NOISY_SYN_CITE,
    ),
    "noisy_schwefel_20d": DatasetConfig(
        key="noisy_schwefel_20d", family="noisy_synthetic", bench_name="bbob_20_1_20_noisy",  # Noisy Schwefel (f20)
        display_name="Noisy-Schwefel-20D (BBOB)",
        description="BBOB f20 (Schwefel, 20D) with moderate Gaussian noise",
        citation=_NOISY_SYN_CITE,
    ),
}

# Combined registry
ALL_DATASETS: Dict[str, DatasetConfig] = {
    **NASBENCH201_DATASETS,
    **HPOBENCH_DATASETS,
    **HPOLIB_DATASETS,
    **SYNTHETIC_DATASETS,
    **NOISY_SYNTHETIC_DATASETS,
}

FAMILY_DATASETS = {
    "nasbench201": list(NASBENCH201_DATASETS.keys()),
    "hpobench":   list(HPOBENCH_DATASETS.keys()),
    "hpolib":     list(HPOLIB_DATASETS.keys()),
    "synthetic":  list(SYNTHETIC_DATASETS.keys()),
    "noisy_synthetic": list(NOISY_SYNTHETIC_DATASETS.keys()),
}

# NASBench201 constants
_NAS_OPS   = ["none", "skip_connect", "nor_conv_1x1", "nor_conv_3x3", "avg_pool_3x3"]
_N_OPS     = len(_NAS_OPS)   # 5
_N_EDGES   = 6               # cell edges in NASBench201 DARTS cell


# =============================================================================
# Evaluator base class
# =============================================================================

class BaseEvaluator:
    """
    Common interface for all three benchmark families.

    Subclasses must set:
        param_names  : List[str]
        bounds_low   : np.ndarray  (continuous lower bounds)
        bounds_high  : np.ndarray  (continuous upper bounds)
        metric_name  : str
        direction    : "minimize" | "maximize"

    And implement:
        _raw_query(vec) -> float   (returns the raw benchmark metric)
    """
    param_names:  List[str]
    bounds_low:   np.ndarray
    bounds_high:  np.ndarray
    metric_name:  str
    direction:    str   # "minimize" | "maximize"

    def __init__(self):
        self._tracker: Dict = {"best_score": -1e18, "best_vec": None,
                               "score_history": []}
        self._eval_count: int = 0
        self._eval_budget: Optional[int] = None   # None = unlimited
        self._last_raw: float = 0.0                # cached fallback for over-budget evals

    def reset_tracker(self) -> None:
        self._tracker = {"best_score": -1e18, "best_vec": None,
                         "score_history": []}
        self._eval_count = 0
        self._last_raw = 0.0

    def set_eval_budget(self, budget: int) -> None:
        """Cap total objective evaluations to *budget*.

        Once *budget* real evaluations have been performed, subsequent
        queries return the last observed raw value and do **not** update
        the convergence tracker.  This guarantees that every algorithm
        operates under an identical evaluation budget.
        """
        self._eval_budget = budget

    def raw_to_score(self, raw: float) -> float:
        """Convert raw metric to a direction-corrected score (higher = better)."""
        return raw if self.direction == "maximize" else -raw

    def score_to_fitness(self, score: float) -> float:
        """Score (higher = better)  ->  fitness (lower = better, i.e. -score)."""
        return -score

    def _update_tracker(self, score: float, vec: np.ndarray) -> None:
        if score > self._tracker["best_score"]:
            self._tracker["best_score"] = score
            self._tracker["best_vec"]   = vec.copy()
        # Always record running best so run_dpo can build a score curve
        # without relying on DPO's internal history format.
        self._tracker["score_history"].append(self._tracker["best_score"])

    def objective(self, params: Dict) -> Tuple[float, Dict]:
        """For DPO's ContinuousOptimizationProblem."""
        vec = np.array([params[k] for k in self.param_names])
        raw, real = self._budgeted_query(vec)
        score   = self.raw_to_score(raw)
        fitness = self.score_to_fitness(score)
        if real:
            self._update_tracker(score, vec)
        return fitness, {self.metric_name: raw}

    def make_baseline_objective(self):
        """Return a callable `f(vec) -> fitness` for baselines (np.ndarray input)."""
        def _fn(vec: np.ndarray) -> float:
            raw, real = self._budgeted_query(vec)
            score = self.raw_to_score(raw)
            if real:
                self._update_tracker(score, vec)
            return self.score_to_fitness(score)
        return _fn

    def _budgeted_query(self, vec: np.ndarray) -> Tuple[float, bool]:
        """Query with budget enforcement.

        Returns ``(raw_value, is_real_eval)``.  When the evaluation
        budget has been exhausted, returns the last cached raw value
        and ``False`` so that the tracker is not updated.
        """
        if self._eval_budget is not None and self._eval_count >= self._eval_budget:
            return self._last_raw, False
        raw = self._raw_query(vec)
        self._eval_count += 1
        self._last_raw = raw
        return raw, True

    def _raw_query(self, vec: np.ndarray) -> float:        # pragma: no cover
        raise NotImplementedError


# =============================================================================
# NASBench201 evaluator
# =============================================================================

class NASBench201Evaluator(BaseEvaluator):
    """
    Wraps ``hpo_benchmarks.NASBench201``.

    Search space: 6 edges × 5 ops (none…avg_pool_3x3), continuous [0, 5).
    """

    def __init__(self, bench_name: str, seed: int = 0):
        super().__init__()
        self.bench        = NASBench201(bench_name, seed=seed)
        self.param_names  = list(self.bench.search_space.keys())  # ['Op0'…'Op5']
        self.choices      = list(self.bench.search_space[self.param_names[0]])
        self.bounds_low   = np.zeros(_N_EDGES, dtype=np.float64)
        self.bounds_high  = np.full(_N_EDGES, _N_OPS - 1e-6, dtype=np.float64)
        self.metric_name  = "val_acc"
        self.direction    = "maximize"

    def vec_to_ops(self, vec: np.ndarray) -> Dict[str, str]:
        return {
            k: self.choices[int(np.clip(v, 0, _N_OPS - 1))]
            for k, v in zip(self.param_names, vec)
        }

    def _raw_query(self, vec: np.ndarray) -> float:
        ops = self.vec_to_ops(vec)
        return float(self.bench(ops)["val_acc"]) / 100.0   # fractional

    def best_arch(self) -> Optional[Dict[str, str]]:
        v = self._tracker.get("best_vec")
        return self.vec_to_ops(v) if v is not None else None


# =============================================================================
# HPOBench evaluator
# =============================================================================

class HPOBenchEvaluator(BaseEvaluator):
    """
    Wraps ``hpo_benchmarks.HPOBench``.

    The search space is a dict of ``{param: [choice1, …]}``.
    Continuous encoding: each dim maps to ``[0, n_choices)``, decoded by floor.
    """

    def __init__(self, bench_name: str, seed: int = 0):
        super().__init__()
        self.bench       = HPOBench(bench_name, seed=seed)
        self.param_names = list(self.bench.search_space.keys())
        self._choices    = {k: list(v) for k, v in self.bench.search_space.items()}
        n_choices        = np.array([len(self._choices[k]) for k in self.param_names],
                                    dtype=np.float64)
        self.bounds_low  = np.zeros(len(self.param_names), dtype=np.float64)
        self.bounds_high = n_choices - 1e-9
        # Primary (first) metric
        self._primary_metric = self.bench.metric_names[0]
        self.metric_name     = self._primary_metric
        self.direction       = self.bench.directions[self._primary_metric]  # "maximize"/"minimize"

    def vec_to_params(self, vec: np.ndarray) -> Dict:
        return {
            k: self._choices[k][int(np.clip(v, 0, len(self._choices[k]) - 1))]
            for k, v in zip(self.param_names, vec)
        }

    def _raw_query(self, vec: np.ndarray) -> float:
        params = self.vec_to_params(vec)
        result = self.bench(params)
        return float(result[self._primary_metric])


# =============================================================================
# HPOLib evaluator
# =============================================================================

class HPOLibEvaluator(BaseEvaluator):
    """
    Wraps ``hpo_benchmarks.HPOLib``.

    Same continuous-encoding strategy as ``HPOBenchEvaluator``.
    HPOLib datasets are typically regression benchmarks with a minimize direction.
    """

    def __init__(self, bench_name: str, seed: int = 0):
        super().__init__()
        self.bench       = HPOLib(bench_name, seed=seed)
        self.param_names = list(self.bench.search_space.keys())
        self._choices    = {k: list(v) for k, v in self.bench.search_space.items()}
        n_choices        = np.array([len(self._choices[k]) for k in self.param_names],
                                    dtype=np.float64)
        self.bounds_low  = np.zeros(len(self.param_names), dtype=np.float64)
        self.bounds_high = n_choices - 1e-9
        self._primary_metric = self.bench.metric_names[0]
        self.metric_name     = self._primary_metric
        self.direction       = self.bench.directions[self._primary_metric]

    def vec_to_params(self, vec: np.ndarray) -> Dict:
        return {
            k: self._choices[k][int(np.clip(v, 0, len(self._choices[k]) - 1))]
            for k, v in zip(self.param_names, vec)
        }

    def _raw_query(self, vec: np.ndarray) -> float:
        params = self.vec_to_params(vec)
        result = self.bench(params)
        return float(result[self._primary_metric])


# =============================================================================
# Synthetic benchmark evaluator (IOHexperimenter/BBOB)
# =============================================================================

class SyntheticEvaluator(BaseEvaluator):
    """
    Official BBOB (Black-Box Optimization Benchmarking) functions via IOHexperimenter.

    Uses IOHexperimenter (Doerr et al., 2018), the successor to COCO, which provides:
      - Standardized benchmark suite used globally
      - Rigorously tested implementations  
      - 24 BBOB functions covering diverse optimization challenges
      - Official reference for algorithm comparison
      - Active maintenance and updates

    bench_name format: "bbob_{fid}_{inst}_{dim}"
      - fid: function ID (1-24)
      - inst: instance number (for different transformations)
      - dim: problem dimension

    Score = -f(x)  (higher = better, optimum value varies by function)
    """

    def __init__(self, bench_name: str, seed: int = 0):
        super().__init__()
        if not HAS_IOH:
            raise ImportError(
                "BBOB benchmarks require 'ioh' library (IOHexperimenter).\n"
                "Install with: pip install ioh"
            )
        
        # Parse bench_name: "bbob_15_1_20" -> function_id=15, instance=1, dimension=20
        parts = bench_name.split("_")
        if len(parts) != 4 or parts[0] != "bbob":
            raise ValueError(
                f"Invalid BBOB bench_name format: {bench_name!r}\n"
                f"Expected: 'bbob_{{fid}}_{{inst}}_{{dim}}'"
            )
        
        func_id = int(parts[1])    # "15" -> 15
        instance = int(parts[2])   # "1" -> 1
        dimension = int(parts[3])  # "20" -> 20
        
        # Create IOH BBOB problem
        # IOH uses ProblemClass.BBOB for the BBOB suite
        self._problem = get_problem(
            fid=func_id,
            instance=instance,
            dimension=dimension,
            problem_class=ProblemClass.BBOB
        )
        
        if self._problem is None:
            raise ValueError(f"Could not load IOH BBOB problem: f{func_id}_i{instance}_d{dimension}")
        
        # Set attributes
        self._dims = dimension
        self.param_names = [f"x{i}" for i in range(self._dims)]
        self.bounds_low = np.array(self._problem.bounds.lb, dtype=np.float64)
        self.bounds_high = np.array(self._problem.bounds.ub, dtype=np.float64)
        self.metric_name = "f_value"
        self.direction = "minimize"
        self._bench_name = bench_name

    def _raw_query(self, vec: np.ndarray) -> float:
        """Query the official IOH BBOB benchmark function."""
        x = np.asarray(vec, dtype=np.float64)
        # IOH returns the function value directly
        val = float(self._problem(x))
        return val


# =============================================================================
# Noisy synthetic benchmark evaluator (IOHexperimenter/BBOB with noise)
# =============================================================================

class NoisySyntheticEvaluator(BaseEvaluator):
    """
    BBOB benchmark functions with added Gaussian noise.

    Uses IOHexperimenter BBOB functions with controlled noise injection:
      - Same base functions as SyntheticEvaluator
      - Moderate Gaussian noise (σ = 0.01 * |f(x)|)
      - Deterministic seed for reproducibility
      - Simulates noisy real-world optimization

    bench_name format: "bbob_{fid}_{inst}_{dim}_noisy"
      - fid: function ID (1-24)
      - inst: instance number
      - dim: problem dimension
      - _noisy: indicator for noise injection

    Score = -f(x)  (higher = better, includes noise in evaluation)
    """

    def __init__(self, bench_name: str, seed: int = 0):
        super().__init__()
        if not HAS_IOH:
            raise ImportError(
                "BBOB benchmarks require 'ioh' library (IOHexperimenter).\n"
                "Install with: pip install ioh"
            )
        
        # Parse bench_name: "bbob_15_1_20_noisy" -> function_id=15, instance=1, dimension=20
        parts = bench_name.split("_")
        if len(parts) != 5 or parts[0] != "bbob" or parts[4] != "noisy":
            raise ValueError(
                f"Invalid BBOB-noisy bench_name format: {bench_name!r}\n"
                f"Expected: 'bbob_{{fid}}_{{inst}}_{{dim}}_noisy'"
            )
        
        func_id = int(parts[1])    # "15" -> 15
        instance = int(parts[2])   # "1" -> 1
        dimension = int(parts[3])  # "20" -> 20
        
        # Create IOH BBOB problem (base function)
        self._problem = get_problem(
            fid=func_id,
            instance=instance,
            dimension=dimension,
            problem_class=ProblemClass.BBOB
        )
        
        if self._problem is None:
            raise ValueError(f"Could not load IOH BBOB problem: f{func_id}_i{instance}_d{dimension}")
        
        # Set attributes
        self._dims = dimension
        self.param_names = [f"x{i}" for i in range(self._dims)]
        self.bounds_low = np.array(self._problem.bounds.lb, dtype=np.float64)
        self.bounds_high = np.array(self._problem.bounds.ub, dtype=np.float64)
        self.metric_name = "f_value"
        self.direction = "minimize"
        self._bench_name = bench_name
        
        # Initialize RNG for noise with fixed seed for reproducibility
        self._seed = seed
        self._rng = np.random.RandomState(seed)
        self._noise_scale = 0.01  # 1% of function value

    def reset_tracker(self) -> None:
        """Reset tracker **and** re-seed the noise RNG for reproducibility."""
        super().reset_tracker()
        self._rng = np.random.RandomState(self._seed)

    def _raw_query(self, vec: np.ndarray) -> float:
        """Query the IOH BBOB benchmark function with added Gaussian noise."""
        x = np.asarray(vec, dtype=np.float64)
        # Get base function value
        base_val = float(self._problem(x))
        # Add proportional Gaussian noise
        noise = self._rng.normal(0, self._noise_scale * (1.0 + abs(base_val)))
        return base_val + noise


# =============================================================================
# Evaluator factory
# =============================================================================

def make_evaluator(cfg: DatasetConfig, seed: int = 0) -> BaseEvaluator:
    """Instantiate the correct evaluator for a DatasetConfig."""
    if cfg.family == "nasbench201":
        return NASBench201Evaluator(cfg.bench_name, seed=seed)
    if cfg.family == "hpobench":
        return HPOBenchEvaluator(cfg.bench_name, seed=seed)
    if cfg.family == "hpolib":
        return HPOLibEvaluator(cfg.bench_name, seed=seed)
    if cfg.family == "synthetic":
        return SyntheticEvaluator(cfg.bench_name, seed=seed)
    if cfg.family == "noisy_synthetic":
        return NoisySyntheticEvaluator(cfg.bench_name, seed=seed)
    raise ValueError(f"Unknown benchmark family: {cfg.family!r}")


# =============================================================================
# Result container
# =============================================================================

@dataclass
class RunResult:
    method:         str
    dataset:        str
    family:         str
    seed:           int
    best_score:     float        # direction-corrected, higher = better
    best_fitness:   float        # = -best_score
    raw_metric:     float        # original metric value (for display)
    metric_name:    str
    direction:      str
    auc:            float
    time_s:         float
    iters_to_95pct: int          # –1 = never reached
    iters_to_99pct: int          # –1 = never reached
    score_curve:    List[float] = field(default_factory=list)


# =============================================================================
# Utility functions
# =============================================================================

def auc_from_curve(
    curve: List[float],
    ref_start: Optional[float] = None,
    ref_end: Optional[float] = None,
) -> float:
    """Normalised AUC of a best-so-far convergence curve.

    When *ref_start* and *ref_end* are supplied they define a **shared**
    normalisation window (e.g. cross-method global worst / best) so that
    AUC values are comparable across different algorithms on the same
    dataset and seed.
    """
    if not curve:
        return 0.0

    ys = np.asarray(curve, dtype=np.float64)
    finite = np.isfinite(ys)
    if not np.any(finite):
        return 0.0
    ys = ys[finite]

    # Convergence should be evaluated on best-so-far (monotone non-decreasing)
    ys = np.maximum.accumulate(ys)
    if ys.size == 1:
        return 1.0

    start = ref_start if ref_start is not None else float(ys[0])
    end   = ref_end   if ref_end   is not None else float(ys[-1])
    span  = end - start
    if abs(span) <= 1e-12:
        return 1.0

    norm = (ys - start) / span
    norm = np.clip(norm, 0.0, 1.0)

    # Integrate over normalised x-axis so the result stays in [0, 1].
    xs  = np.linspace(0.0, 1.0, num=len(norm), dtype=np.float64)
    _t  = getattr(np, "trapezoid", np.trapz)
    return float(np.clip(_t(norm, xs), 0.0, 1.0))


def iter_to_pct_of_peak(curve: List[float], fraction: float = 0.95) -> int:
    """First iteration where score reaches *fraction* of the improvement
    from the initial score to the peak score.

    Works correctly regardless of the sign of the scores (e.g. negative
    direction-corrected scores on minimisation benchmarks).
    Returns -1 if the threshold is never reached.
    """
    if not curve:
        return -1
    start = curve[0]
    peak  = max(curve)
    if abs(peak - start) < 1e-15:
        return 1   # already converged from the start
    threshold = start + fraction * (peak - start)
    for i, v in enumerate(curve):
        if v >= threshold:
            return i + 1
    return -1


# =============================================================================
# DPO runner  (generic – works for all three families)
# =============================================================================

_SILENT = logging.getLogger("DPO-Bench")
_SILENT.setLevel(logging.CRITICAL)
_SILENT.propagate = False
_SILENT.disabled  = True


def run_dpo(
    evaluator:    BaseEvaluator,
    cfg:          DatasetConfig,
    population:   int,
    iterations:   int,
    seed:         int,
    variant:      str = "full",  # "full", "nodebt", "noaccept", "norepay"
) -> RunResult:
    np.random.seed(seed)
    evaluator.reset_tracker()
    dim = len(evaluator.param_names)

    problem = ContinuousOptimizationProblem(
        objective_fn = evaluator.objective,
        param_bounds = list(zip(evaluator.bounds_low.tolist(),
                                evaluator.bounds_high.tolist())),
        param_names  = evaluator.param_names,
    )
    # Select the preset that matches the problem domain
    if cfg.family == "nasbench201":
        config = DPO_Presets.NAS_Config(
            population_size = population,
            max_iterations  = iterations,
            aggressive_mode = True,
        )
    elif cfg.family == "synthetic":
        # Analytic benchmark functions — dedicated preset: no overshoot, low alpha,
        # strong global pull, continuous_mode=True (skips pareto rank + lazy repayment)
        config = DPO_Presets.ContinuousAnalytic_Config(
            population_size = population,
            max_iterations  = iterations,
            num_params      = dim,
        )
    elif cfg.family == "noisy_synthetic":
        # Noisy / rugged continuous functions — allow debt-driven exploration
        config = DPO_Presets.HyperparameterTuning_Config(
            population_size = population,
            max_iterations  = iterations,
        )
    else:
        # hpobench / hpolib are continuous HPO – use the dedicated preset
        config = DPO_Presets.HyperparameterTuning_Config(
            population_size = population,
            max_iterations  = iterations,
        )

    config.verbose     = False
    config.num_islands = max(1, min(config.num_islands, population // 30))

    # Apply ablation settings based on variant
    # CRITICAL: Ablation variants must maintain logical consistency
    if variant == "nodebt":
        # Disable debt mechanism entirely (no accumulation, no repayment)
        # Keep worse acceptance to test SA-like behavior without debt tracking
        config.enable_debt_accumulation = False
        config.enable_worse_acceptance = True
        config.enable_debt_repayment = False  # No debt to repay
    elif variant == "noaccept":
        # Pure greedy: never accept worse, so debt mechanisms irrelevant
        config.enable_debt_accumulation = False  # Never triggered
        config.enable_worse_acceptance = False
        config.enable_debt_repayment = False     # Never triggered
    elif variant == "norepay":
        # Accumulate debt without repayment (no overshoot jumps)
        config.enable_debt_accumulation = True
        config.enable_worse_acceptance = True
        config.enable_debt_repayment = False
    # else: variant == "full", use defaults (all True)

    # For small budgets, start exploitation quickly.
    if iterations < 100:
        budget_scale = iterations / 100.0
        config.exploration_phase_ratio  = max(0.05, config.exploration_phase_ratio  * budget_scale)
        config.debt_accumulation_phase  = max(0.25, config.debt_accumulation_phase  * budget_scale)
        config.convergence_window       = max(5,    int(config.convergence_window    * budget_scale))
        config.patience                 = max(10,   int(config.patience              * budget_scale))

    optimizer = DPO_Universal(problem=problem, config=config, logger=_SILENT)

    t0      = time.perf_counter()
    result  = optimizer.optimize()
    elapsed = time.perf_counter() - t0

    # evaluator._tracker["score_history"] is the ground-truth convergence curve:
    # it is appended on every real objective call by BaseEvaluator._update_tracker,
    # recording the running best (direction-corrected, higher = better).
    # This is more reliable than DPO's internal history["best_fitness"], which
    # records the *internal fitness* (includes cost/penalty terms and may be
    # mis-scaled), and is often empty or starts well after iteration 0.
    score_curve: List[float] = [float(v) for v in
                                 evaluator._tracker.get("score_history", [])]

    # Subsample score_curve to `iterations` points so its x-axis (evaluation
    # call index) is comparable to baseline curves (one point per epoch).
    # Without this, ->95%/->99% for DPO would be in units of individual
    # evaluations (~pop*iter calls) while baselines are in units of epochs.
    if len(score_curve) > iterations:
        indices = [int(round(i * (len(score_curve) - 1) / (iterations - 1)))
                   for i in range(iterations)]
        score_curve = [score_curve[idx] for idx in indices]

    # Map variant to method name
    method_name = {
        "full": "DPO",
        "nodebt": "DPO-NoDebt",
        "noaccept": "DPO-NoAccept",
        "norepay": "DPO-NoRepay",
    }.get(variant, "DPO")

    # Prefer tracker's best_score (ground truth from real evaluations)
    best_score = float(evaluator._tracker.get("best_score", -1e18))
    # Sanity-check / fallback: take the best from the curve as well
    if score_curve:
        best_score = max(best_score, max(score_curve))
    # If tracker is still empty, nothing worked
    if best_score <= -1e17:
        best_score = 0.0

    # Reconstruct raw metric from direction-corrected score
    raw_metric = best_score if evaluator.direction == "maximize" else -best_score

    peak = best_score
    auc  = auc_from_curve(score_curve)

    return RunResult(
        method         = method_name,
        dataset        = cfg.key,
        family         = cfg.family,
        seed           = seed,
        best_score     = best_score,
        best_fitness   = -best_score,
        raw_metric     = raw_metric,
        metric_name    = evaluator.metric_name,
        direction      = evaluator.direction,
        auc            = auc,
        time_s         = elapsed,
        iters_to_95pct = iter_to_pct_of_peak(score_curve, 0.95),
        iters_to_99pct = iter_to_pct_of_peak(score_curve, 0.99),
        score_curve    = score_curve,
    )


# =============================================================================
# Baseline runner  (generic)
# =============================================================================

def run_baseline(
    name:         str,
    optimizer_obj,
    evaluator:    BaseEvaluator,
    cfg:          DatasetConfig,
    iterations:   int,
    seed:         int,
) -> RunResult:
    np.random.seed(seed)
    evaluator.reset_tracker()
    objective = evaluator.make_baseline_objective()
    bounds    = (evaluator.bounds_low, evaluator.bounds_high)
    dim       = len(evaluator.param_names)

    t0 = time.perf_counter()
    best_vec, best_fit, raw_conv = optimizer_obj.optimize(
        objective_fn   = objective,
        bounds         = bounds,
        dim            = dim,
        max_iterations = iterations,
    )
    elapsed = time.perf_counter() - t0

    # raw_conv is fitness curve (lower = better) -> score curve
    score_curve = [float(-f) for f in raw_conv]
    best_score  = float(-best_fit)
    if score_curve:
        best_score = max(best_score, max(score_curve))

    raw_metric = best_score if evaluator.direction == "maximize" else -best_score
    peak = best_score
    auc  = auc_from_curve(score_curve)

    return RunResult(
        method         = name,
        dataset        = cfg.key,
        family         = cfg.family,
        seed           = seed,
        best_score     = best_score,
        best_fitness   = -best_score,
        raw_metric     = raw_metric,
        metric_name    = evaluator.metric_name,
        direction      = evaluator.direction,
        auc            = auc,
        time_s         = elapsed,
        iters_to_95pct = iter_to_pct_of_peak(score_curve, 0.95),
        iters_to_99pct = iter_to_pct_of_peak(score_curve, 0.99),
        score_curve    = score_curve,
    )


# =============================================================================
# Statistical aggregation
# =============================================================================

def aggregate(results: List[RunResult]) -> Dict:
    if not results:
        return {}
    scores = [r.best_score  for r in results]
    fits   = [r.best_fitness for r in results]
    aucs   = [r.auc         for r in results]
    times  = [r.time_s      for r in results]
    t95s   = [r.iters_to_95pct for r in results if r.iters_to_95pct > 0]
    t99s   = [r.iters_to_99pct for r in results if r.iters_to_99pct > 0]

    curves = [r.score_curve for r in results if r.score_curve]
    if curves:
        max_len    = max(len(c) for c in curves)
        padded     = [np.pad(c, (0, max_len - len(c)), mode="edge") for c in curves]
        mean_curve = np.mean(padded, axis=0).tolist()
        std_curve  = np.std(padded,  axis=0).tolist()
    else:
        mean_curve, std_curve = [], []

    return {
        "score_mean":  float(np.mean(scores)),
        "score_std":   float(np.std(scores)),
        "score_best":  float(np.max(scores)),
        "score_worst": float(np.min(scores)),
        "fit_mean":    float(np.mean(fits)),
        "auc_mean":    float(np.mean(aucs)),
        "auc_std":     float(np.std(aucs)),
        "time_mean":   float(np.mean(times)),
        "time_std":    float(np.std(times)),
        "t95_mean":    float(np.mean(t95s)) if t95s else float("nan"),
        "t99_mean":    float(np.mean(t99s)) if t99s else float("nan"),
        "t95_rate":    len(t95s) / len(results),
        "t99_rate":    len(t99s) / len(results),
        "mean_curve":  mean_curve,
        "std_curve":   std_curve,
        "raw_scores":  scores,
        "metric_name": results[0].metric_name,
        "direction":   results[0].direction,
        "family":      results[0].family,
    }


def wilcoxon_dpo_vs(
    all_results: Dict[str, Dict[str, List[RunResult]]],
    datasets:    List[str],
    methods:     List[str],
) -> Dict[str, Dict[str, Dict]]:
    """Wilcoxon signed-rank test (DPO vs baseline, two-sided, α=0.05) per dataset."""
    if not HAS_SCIPY:
        return {}
    import warnings
    tests: Dict[str, Dict[str, Dict]] = {}
    for ds in datasets:
        tests[ds] = {}
        dpo_vals = [r.best_score for r in all_results[ds].get("DPO", [])]
        if not dpo_vals:
            continue
        for m in methods:
            if m == "DPO":
                continue
            base_vals = [r.best_score for r in all_results[ds].get(m, [])]
            if len(base_vals) != len(dpo_vals) or len(dpo_vals) < 2:
                continue
            # Skip if all paired differences are zero (identical scores)
            diffs = [a - b for a, b in zip(dpo_vals, base_vals)]
            if all(abs(d) < 1e-15 for d in diffs):
                tests[ds][m] = {
                    "statistic":   0.0,
                    "p_value":     1.0,
                    "significant": False,
                }
                continue
            try:
                with warnings.catch_warnings():
                    warnings.simplefilter("ignore", RuntimeWarning)
                    stat, pval = wilcoxon(dpo_vals, base_vals, alternative="two-sided")
                # Handle NaN from remaining edge cases (e.g. too few non-zero diffs)
                if math.isnan(pval):
                    pval = 1.0
                tests[ds][m] = {
                    "statistic":   float(stat),
                    "p_value":     float(pval),
                    "significant": bool(pval < 0.05),
                }
            except Exception:
                pass
    return tests


def friedman_mean_ranks(
    all_stats: Dict[str, Dict[str, Dict]],
    datasets:  List[str],
    methods:   List[str],
) -> Dict[str, float]:
    """Mean rank across datasets with proper fractional tie-breaking."""
    rank_sum: Dict[str, float] = defaultdict(float)
    n_datasets = 0
    for ds in datasets:
        valid_methods = [m for m in methods if all_stats[ds].get(m)]
        if not valid_methods:
            continue
        n_datasets += 1
        if HAS_SCIPY:
            # scipy.stats.rankdata with fractional ties (lower rank = better)
            scores = [all_stats[ds][m]["score_mean"] for m in valid_methods]
            # Negate: rankdata ranks ascending, but higher score = better
            ranks = rankdata([-s for s in scores], method='average')
            for m, r in zip(valid_methods, ranks):
                rank_sum[m] += r
        else:
            sorted_m = sorted(valid_methods, key=lambda m: -all_stats[ds][m]["score_mean"])
            for rank, m in enumerate(sorted_m, 1):
                rank_sum[m] += rank
    return {m: rank_sum[m] / max(1, n_datasets) for m in methods}


# =============================================================================
# Console output
# =============================================================================

DIVIDER_WIDE  = "=" * 100
DIVIDER_LIGHT = "-" * 100


def _fmt(val: float, fmt: str = ".4f", nan_str: str = "  N/A ") -> str:
    return nan_str if math.isnan(val) else f"{val:{fmt}}"


def _fmt_score(val: float, nan_str: str = "  N/A ") -> str:
    if math.isnan(val):
        return nan_str
    # Avoid visual noise from signed zero in near-perfect minimization tasks.
    v = 0.0 if abs(val) < 5e-8 else val
    # Keep tiny values readable instead of printing "-0.0000".
    if v != 0.0 and abs(v) < 1e-3:
        return f"{v:.2e}"
    return f"{v:.4f}"


def _fmt_score_latex(val: float, nan_str: str = "--") -> str:
    if math.isnan(val):
        return nan_str
    v = 0.0 if abs(val) < 5e-8 else val
    if v != 0.0 and abs(v) < 1e-3:
        return f"{v:.2e}"
    return f"{v:.4f}"


def _direction_arrow(direction: str) -> str:
    return "^" if direction == "maximize" else "v (-score)"


def print_dataset_table(
    cfg:          DatasetConfig,
    stats:        Dict[str, Dict],
    method_order: List[str],
    wilcoxon_res: Optional[Dict[str, Dict]] = None,
) -> None:
    # Infer direction from first available stat
    direction = next(
        (stats[m]["direction"] for m in method_order if stats.get(m)), "maximize"
    )
    metric_name = next(
        (stats[m]["metric_name"] for m in method_order if stats.get(m)), "score"
    )
    print(f"\n{DIVIDER_WIDE}")
    print(f"  [{cfg.family.upper()}]  Dataset : {cfg.display_name}  |  {cfg.description}")
    print(f"  Metric  : {metric_name}  {_direction_arrow(direction)}")
    print(f"  Cite    : {cfg.citation}")
    print(DIVIDER_WIDE)

    hdr = (
        f"{'Rank':<5}{'Method':<10}"
        f"{'Score(mean+-std)':>20}"
        f"{'Best Score':>12}"
        f"{'AUC mean':>10}"
        f"{'AUC std':>9}"
        f"{'Time(s)':>9}"
        f"{'->95%':>7}"
        f"{'->99%':>7}"
        f"{'95% hit':>9}"
        f"{'p-val*':>9}"
    )
    print(hdr)
    print(DIVIDER_LIGHT)

    valid_methods = [m for m in method_order if stats.get(m)]
    ranked = sorted(valid_methods, key=lambda m: -stats[m]["score_mean"])
    for rank, m in enumerate(ranked, 1):
        s       = stats[m]
        sc_str  = f"{_fmt_score(s['score_mean'])} +- {_fmt_score(s['score_std'])}"
        p_str   = ""
        if wilcoxon_res and m in wilcoxon_res:
            pv  = wilcoxon_res[m]["p_value"]
            sig = " *" if wilcoxon_res[m]["significant"] else "  "
            p_str = f"{pv:.3f}{sig}"
        marker = " <" if rank == 1 else "  "
        print(
            f"{rank:<5}{m:<10}"
            f"{sc_str:>20}"
            f"{_fmt_score(s['score_best']):>12}"
            f"{_fmt(s['auc_mean']):>10}"
            f"{_fmt(s['auc_std']):>9}"
            f"{_fmt(s['time_mean'], '.2f'):>9}"
            f"{_fmt(s['t95_mean'], '.1f'):>7}"
            f"{_fmt(s['t99_mean'], '.1f'):>7}"
            f"{s['t95_rate']:>8.0%}"
            f"{p_str:>9}"
            f"{marker}"
        )
    print(f"\n  * Wilcoxon signed-rank: DPO vs baseline, two-sided, alpha = 0.05")
    print(f"  Score = {metric_name} if maximize / -{metric_name} if minimize  (always higher-is-better)")


def print_family_summary(
    family:     str,
    all_stats:  Dict[str, Dict[str, Dict]],
    datasets:   List[str],
    methods:    List[str],
    ranks:      Dict[str, float],      # global ranks (used only for fallback)
) -> None:
    fam_datasets = [d for d in datasets if ALL_DATASETS[d].family == family]
    if not fam_datasets:
        return

    # Compute per-family ranks (not the global ones)
    fam_ranks = friedman_mean_ranks(all_stats, fam_datasets, methods)

    print(f"\n{DIVIDER_WIDE}")
    print(f"  FAMILY SUMMARY: {family.upper()}")
    print(DIVIDER_WIDE)
    wins: Dict[str, int] = defaultdict(int)
    for ds in fam_datasets:
        valid = {m: all_stats[ds][m] for m in methods if all_stats[ds].get(m)}
        if valid:
            best_m = max(valid.keys(), key=lambda m: valid[m]["score_mean"])
            wins[best_m] += 1
    print(
        f"{'Method':<12}  {'Mean Rank':>10}  {'#1 Wins':>8}"
        f"  {'Mean Score':>12}  {'Mean AUC':>12}  {'Mean Time(s)':>14}"
    )
    print(DIVIDER_LIGHT)
    valid_methods = [m for m in methods if any(all_stats[ds].get(m) for ds in fam_datasets)]
    for m in sorted(valid_methods, key=lambda x: fam_ranks.get(x, 999)):
        mean_score = float(np.mean([
            all_stats[ds][m]["score_mean"]
            for ds in fam_datasets if all_stats[ds].get(m)
        ]))
        mean_auc = float(np.mean([
            all_stats[ds][m]["auc_mean"]
            for ds in fam_datasets if all_stats[ds].get(m)
        ]))
        mean_time = float(np.mean([
            all_stats[ds][m]["time_mean"]
            for ds in fam_datasets if all_stats[ds].get(m)
        ]))
        mk = "  < BEST" if m == min(valid_methods, key=lambda x: fam_ranks.get(x, 999)) else ""
        print(
            f"{m:<12}  {fam_ranks.get(m, float('nan')):>10.2f}  {wins[m]:>8d}"
            f"  {_fmt_score(mean_score):>12}  {mean_auc:>12.4f}  {mean_time:>14.2f}{mk}"
        )


def print_overall_summary(
    all_stats:  Dict[str, Dict[str, Dict]],
    datasets:   List[str],
    methods:    List[str],
    mean_ranks: Dict[str, float],
) -> None:
    print(f"\n{DIVIDER_WIDE}")
    print("  OVERALL SUMMARY  --  Mean rank across ALL datasets (lower = better)")
    print(DIVIDER_WIDE)
    wins: Dict[str, int] = defaultdict(int)
    for ds in datasets:
        valid = {m: all_stats[ds][m] for m in methods if all_stats[ds].get(m)}
        if valid:
            best_m = max(valid.keys(), key=lambda m: valid[m]["score_mean"])
            wins[best_m] += 1
    valid_methods = [m for m in methods if mean_ranks.get(m) is not None]
    print(
        f"{'Method':<12}  {'Mean Rank':>10}  {'#1 Wins':>8}"
        f"  {'Mean Score':>12}  {'Mean AUC':>12}  {'Mean Time(s)':>14}"
    )
    print(DIVIDER_LIGHT)
    for m in sorted(valid_methods, key=lambda x: mean_ranks.get(x, 999)):
        sc_vals  = [all_stats[ds][m]["score_mean"] for ds in datasets if all_stats[ds].get(m)]
        auc_vals = [all_stats[ds][m]["auc_mean"]   for ds in datasets if all_stats[ds].get(m)]
        t_vals   = [all_stats[ds][m]["time_mean"]  for ds in datasets if all_stats[ds].get(m)]
        oms  = float(np.mean(sc_vals))  if sc_vals  else float("nan")
        oma  = float(np.mean(auc_vals)) if auc_vals else float("nan")
        omt  = float(np.mean(t_vals))   if t_vals   else float("nan")
        mk   = "  < BEST" if m == min(valid_methods, key=lambda x: mean_ranks.get(x, 999)) else ""
        print(
            f"{m:<12}  {mean_ranks[m]:>10.2f}  {wins[m]:>8d}"
            f"  {_fmt_score(oms):>12}  {oma:>12.4f}  {omt:>14.2f}{mk}"
        )


def print_latex_table(
    all_stats:  Dict[str, Dict[str, Dict]],
    datasets:   List[str],
    methods:    List[str],
    mean_ranks: Dict[str, float],
) -> None:
    print(f"\n% --- LaTeX Table -----------------------------------------------")
    ds_cfgs = [ALL_DATASETS[d] for d in datasets]
    col_spec = "l" + "r" * len(datasets) + "rr"
    print(r"\begin{table}[t]")
    print(r"\centering")
    print(r"\caption{Comprehensive HPO Benchmark: mean score $\pm$ std (all runs) and mean runtime in seconds}")
    print(f"\\begin{{tabular}}{{{col_spec}}}")
    print(r"\toprule")
    ds_headers = " & ".join(f"\\textbf{{{c.display_name}}}" for c in ds_cfgs)
    print(f"\\textbf{{Method}} & {ds_headers} & \\textbf{{Rank}} & \\textbf{{Time(s)}}" + r" \\")
    print(r"\midrule")
    valid = [m for m in methods if mean_ranks.get(m) is not None]
    for m in sorted(valid, key=lambda x: mean_ranks.get(x, 999)):
        cols = []
        for ds in datasets:
            s = all_stats[ds].get(m)
            if not s:
                cols.append("--")
                continue
            cell = f"{_fmt_score_latex(s['score_mean'])}$\\pm${_fmt_score_latex(s['score_std'])}"
            fam_methods = [mm for mm in valid if all_stats[ds].get(mm)]
            if fam_methods:
                best_m = max(fam_methods, key=lambda mm: all_stats[ds][mm]["score_mean"])
                if m == best_m:
                    cell = f"\\textbf{{{cell}}}"
            cols.append(cell)
        t_vals = [all_stats[ds][m]["time_mean"] for ds in datasets if all_stats[ds].get(m)]
        mean_time = float(np.mean(t_vals)) if t_vals else float("nan")
        print(f"{m} & {' & '.join(cols)} & {mean_ranks[m]:.2f} & {_fmt_score_latex(mean_time)}" + r" \\")
    print(r"\bottomrule")
    print(r"\end{tabular}")
    print(r"\end{table}")


# =============================================================================
# Plot palette
# =============================================================================

PALETTE = {
    "DPO":  "#e74c3c",           # Red - Full DPO
    "DPO-NoDebt": "#c0392b",     # Dark red - NoDebt variant
    "DPO-NoAccept": "#e67e22",   # Orange - NoAccept variant
    "DPO-NoRepay": "#f39c12",    # Light orange - NoRepay variant
    "JADE": "#2980b9",
    "DE":   "#16a085",
    "PSO":  "#8e44ad",
    "GWO":  "#f39c12",
    "ABC":  "#27ae60",
    "WOA":  "#d35400",
    "GA":   "#7f8c8d",
    "FA":   "#bdc3c7",
    "SA":   "#95a5a6",
    "ACO":  "#1abc9c",
}
DPO_STYLE  = {"lw": 2.6, "ls": "-",  "zorder": 10, "alpha": 1.0}
BASE_STYLE = {"lw": 1.2, "ls": "--", "zorder": 5,  "alpha": 0.75}

MPL_PARAMS = {
    "font.family":         "DejaVu Sans",
    "font.size":           9,
    "axes.titlesize":      10,
    "axes.labelsize":      9,
    "legend.fontsize":     7.5,
    "xtick.labelsize":     8,
    "ytick.labelsize":     8,
    "figure.dpi":          150,
    "axes.spines.top":     False,
    "axes.spines.right":   False,
    "axes.grid":           True,
    "grid.alpha":          0.30,
    "savefig.bbox":        "tight",
}

FAMILY_TITLE = {
    "nasbench201":    "NAS-Bench-201",
    "hpobench":       "HPOBench",
    "hpolib":         "HPOLib",
    "synthetic":      "Synthetic",
    "noisy_synthetic": "Noisy-Synthetic",
}


def _apply_mpl():
    if HAS_MPL:
        plt.rcParams.update(MPL_PARAMS)


# =============================================================================
# Plot helpers  (each function is self-contained, takes all_stats / all_results)
# =============================================================================

def _cols_rows(n: int) -> Tuple[int, int]:
    ncols = min(4, n)
    nrows = math.ceil(n / ncols)
    return ncols, nrows


def plot_convergence(
    all_stats: Dict[str, Dict[str, Dict]],
    datasets:  List[str],
    methods:   List[str],
    out_dir:   Path,
) -> None:
    if not HAS_MPL:
        return
    _apply_mpl()
    # Group by family for subplot titles
    ncols, nrows = _cols_rows(len(datasets))
    fig, axes = plt.subplots(nrows, ncols, figsize=(6 * ncols, 4 * nrows))
    axes = np.array(axes).flatten()

    for ax_i, ds in enumerate(datasets):
        ax  = axes[ax_i]
        cfg = ALL_DATASETS[ds]
        for m in methods:
            s = all_stats[ds].get(m, {})
            if not s.get("mean_curve"):
                continue
            xs    = np.arange(1, len(s["mean_curve"]) + 1)
            ys    = np.array(s["mean_curve"])
            se    = np.array(s["std_curve"]) / 2
            color = PALETTE.get(m, "#555")
            # Apply DPO style to all DPO variants
            is_dpo_variant = m.startswith("DPO")
            style = DPO_STYLE if is_dpo_variant else BASE_STYLE
            ax.plot(xs, ys, color=color, label=m, **style)
            ax.fill_between(xs, ys - se, ys + se, color=color, alpha=0.10)
        ax.set_title(f"[{FAMILY_TITLE[cfg.family]}] {cfg.display_name}", fontweight="bold")
        ax.set_xlabel("Iteration")
        ax.set_ylabel("Best Score (↑ better)")
        ax.xaxis.set_major_locator(MaxNLocator(integer=True, nbins=6))
        ax.legend(ncol=2, fontsize=7)

    for ax_i in range(len(datasets), len(axes)):
        axes[ax_i].set_visible(False)

    fig.suptitle(
        "Convergence Curves  (mean ± ½std over seeds)  —  All Benchmark Families",
        fontsize=12, fontweight="bold", y=1.01,
    )
    plt.tight_layout()
    out = out_dir / "convergence_curves.png"
    plt.savefig(out, dpi=150)
    plt.close(fig)
    print(f"  [Plot] Convergence curves      -> {out}")


def plot_score_bars(
    all_stats: Dict[str, Dict[str, Dict]],
    datasets:  List[str],
    methods:   List[str],
    out_dir:   Path,
) -> None:
    if not HAS_MPL:
        return
    _apply_mpl()
    ncols, nrows = _cols_rows(len(datasets))
    fig, axes = plt.subplots(nrows, ncols, figsize=(5 * ncols, 4.5 * nrows))
    axes = np.array(axes).flatten()

    for ax_i, ds in enumerate(datasets):
        ax  = axes[ax_i]
        cfg = ALL_DATASETS[ds]
        valid  = [m for m in methods if all_stats[ds].get(m)]
        sorted_m = sorted(valid, key=lambda m: -all_stats[ds][m]["score_mean"])
        means  = [all_stats[ds][m]["score_mean"] for m in sorted_m]
        stds   = [all_stats[ds][m]["score_std"]  for m in sorted_m]
        colors = [PALETTE.get(m, "#555") for m in sorted_m]
        bars   = ax.bar(range(len(sorted_m)), means, yerr=stds,
                        color=colors, capsize=4, error_kw={"lw": 1.2})
        ax.set_xticks(range(len(sorted_m)))
        ax.set_xticklabels(sorted_m, rotation=40, ha="right")
        ax.set_title(f"[{FAMILY_TITLE[cfg.family]}] {cfg.display_name}", fontweight="bold")
        ax.set_ylabel("Mean Best Score (↑)")
        if means:
            ax.set_ylim(max(min(means) - 0.05 * abs(min(means)), min(means) - 0.05),
                        max(means) + 0.06 * abs(max(means)))
        for i, m in enumerate(sorted_m):
            if m.startswith("DPO"):  # Highlight all DPO variants
                bars[i].set_edgecolor("black")
                bars[i].set_linewidth(1.8)

    for ax_i in range(len(datasets), len(axes)):
        axes[ax_i].set_visible(False)

    fig.suptitle("Mean Best Score per Dataset  (error bars = ±1 std)", fontsize=11, fontweight="bold")
    plt.tight_layout()
    out = out_dir / "score_bars.png"
    plt.savefig(out, dpi=150)
    plt.close(fig)
    print(f"  [Plot] Score bars              -> {out}")


def plot_boxplots(
    all_results: Dict[str, Dict[str, List[RunResult]]],
    datasets:    List[str],
    methods:     List[str],
    out_dir:     Path,
) -> None:
    if not HAS_MPL:
        return
    _apply_mpl()
    ncols, nrows = _cols_rows(len(datasets))
    fig, axes = plt.subplots(nrows, ncols, figsize=(5 * ncols, 5 * nrows))
    axes = np.array(axes).flatten()

    for ax_i, ds in enumerate(datasets):
        ax  = axes[ax_i]
        cfg = ALL_DATASETS[ds]
        valid    = [m for m in methods if all_results[ds].get(m)]
        sorted_m = sorted(
            valid,
            key=lambda m: -float(np.mean([r.best_score for r in all_results[ds][m]])),
        )
        data   = [[r.best_score for r in all_results[ds][m]] for m in sorted_m]
        colors = [PALETTE.get(m, "#555") for m in sorted_m]
        bp     = ax.boxplot(data, patch_artist=True, notch=False, vert=True,
                            medianprops={"color": "black", "lw": 1.5})
        for patch, c in zip(bp["boxes"], colors):
            patch.set_facecolor(c)
            patch.set_alpha(0.70)
        ax.set_xticks(range(1, len(sorted_m) + 1))
        ax.set_xticklabels(sorted_m, rotation=40, ha="right")
        ax.set_title(f"[{FAMILY_TITLE[cfg.family]}] {cfg.display_name}", fontweight="bold")
        ax.set_ylabel("Best Score (↑)")

    for ax_i in range(len(datasets), len(axes)):
        axes[ax_i].set_visible(False)

    fig.suptitle("Score Distribution (box = IQR, whiskers = 1.5×IQR)", fontsize=11, fontweight="bold")
    plt.tight_layout()
    out = out_dir / "boxplot_comparison.png"
    plt.savefig(out, dpi=150)
    plt.close(fig)
    print(f"  [Plot] Box plots               -> {out}")


def plot_rank_heatmap(
    all_stats: Dict[str, Dict[str, Dict]],
    datasets:  List[str],
    methods:   List[str],
    out_dir:   Path,
) -> None:
    if not HAS_MPL:
        return
    _apply_mpl()
    rank_mat = np.zeros((len(methods), len(datasets)), dtype=np.float32)
    for j, ds in enumerate(datasets):
        valid = [m for m in methods if all_stats[ds].get(m)]
        sorted_m = sorted(valid, key=lambda m: -all_stats[ds][m]["score_mean"])
        for rank, m in enumerate(sorted_m, 1):
            rank_mat[methods.index(m), j] = rank
    ds_labels = [ALL_DATASETS[d].display_name for d in datasets]

    fig, ax = plt.subplots(figsize=(max(8, len(datasets) * 1.6), max(5, len(methods) * 0.65)))
    im = ax.imshow(rank_mat, aspect="auto", cmap="RdYlGn_r", vmin=1, vmax=len(methods))
    ax.set_xticks(range(len(datasets)))
    ax.set_xticklabels(ds_labels, rotation=30, ha="right")
    ax.set_yticks(range(len(methods)))
    ax.set_yticklabels(methods)
    for i in range(len(methods)):
        for j in range(len(datasets)):
            val = rank_mat[i, j]
            if val == 0:
                continue
            ax.text(j, i, f"{val:.0f}", ha="center", va="center", fontsize=9,
                    color="white" if val > len(methods) * 0.6 else "black")
    plt.colorbar(im, ax=ax, label="Rank (1 = best)")
    ax.set_title("Rank Heatmap — All Benchmark Families", fontsize=12, fontweight="bold")
    plt.tight_layout()
    out = out_dir / "rank_heatmap.png"
    plt.savefig(out, dpi=150)
    plt.close(fig)
    print(f"  [Plot] Rank heatmap            -> {out}")


def plot_auc_lineplot(
    all_stats: Dict[str, Dict[str, Dict]],
    datasets:  List[str],
    methods:   List[str],
    out_dir:   Path,
) -> None:
    if not HAS_MPL:
        return
    _apply_mpl()
    ds_labels  = [ALL_DATASETS[d].display_name for d in datasets]
    families   = [ALL_DATASETS[d].family for d in datasets]
    # Colour x-axis labels by family
    fam_colors = {
        "nasbench201": "#2c3e50",
        "hpobench":    "#1a5276",
        "hpolib":      "#6e2f7a",
        "synthetic":   "#145a32",
    }
    fig, ax    = plt.subplots(figsize=(max(10, len(datasets) * 1.6), 5))

    sorted_m = sorted(
        [m for m in methods if any(all_stats[ds].get(m) for ds in datasets)],
        key=lambda m: -float(np.mean([
            all_stats[ds][m]["auc_mean"]
            for ds in datasets if all_stats[ds].get(m)
        ])),
    )
    for m in sorted_m:
        aucs   = [all_stats[ds][m]["auc_mean"] if all_stats[ds].get(m) else float("nan")
                  for ds in datasets]
        color  = PALETTE.get(m, "#555")
        # Apply DPO style to all DPO variants
        is_dpo_variant = m.startswith("DPO")
        style  = DPO_STYLE if is_dpo_variant else BASE_STYLE
        ax.plot(range(len(datasets)), aucs, color=color, marker="o", ms=5, label=m, **style)

    ax.set_xticks(range(len(datasets)))
    ax.set_xticklabels(ds_labels, rotation=30, ha="right")
    for lbl, fam in zip(ax.get_xticklabels(), families):
        lbl.set_color(fam_colors.get(fam, "#000"))
    ax.set_title("Mean AUC across All Datasets", fontsize=12, fontweight="bold")
    ax.set_ylabel("Mean AUC")
    ax.legend(ncol=2, fontsize=8)
    plt.tight_layout()
    out = out_dir / "auc_lineplot.png"
    plt.savefig(out, dpi=150)
    plt.close(fig)
    print(f"  [Plot] AUC line plot           -> {out}")


def plot_convergence_speed(
    all_stats: Dict[str, Dict[str, Dict]],
    datasets:  List[str],
    methods:   List[str],
    out_dir:   Path,
) -> None:
    if not HAS_MPL:
        return
    _apply_mpl()
    ncols, nrows = _cols_rows(len(datasets))
    fig, axes = plt.subplots(nrows, ncols, figsize=(4.5 * ncols, 4.5 * nrows))
    axes = np.array(axes).flatten()

    for ax_i, ds in enumerate(datasets):
        ax  = axes[ax_i]
        cfg = ALL_DATASETS[ds]
        valid    = [m for m in methods if all_stats[ds].get(m)]
        sorted_m = sorted(valid, key=lambda m: (
            all_stats[ds][m]["t95_mean"]
            if not math.isnan(all_stats[ds][m]["t95_mean"]) else 9999
        ))
        valid95  = [(m, all_stats[ds][m]["t95_mean"]) for m in sorted_m
                    if not math.isnan(all_stats[ds][m]["t95_mean"])]
        if not valid95:
            ax.set_title(f"{cfg.display_name} (no 95% data)")
            continue
        ms_, vs_ = zip(*valid95)
        colors   = [PALETTE.get(m, "#555") for m in ms_]
        ax.barh(range(len(ms_)), vs_, color=colors, edgecolor="white", height=0.6)
        ax.set_yticks(range(len(ms_)))
        ax.set_yticklabels(ms_)
        ax.invert_yaxis()
        ax.set_xlabel("Mean iters to 95% of peak")
        ax.set_title(f"[{FAMILY_TITLE[cfg.family]}] {cfg.display_name}", fontweight="bold")

    for ax_i in range(len(datasets), len(axes)):
        axes[ax_i].set_visible(False)

    fig.suptitle("Convergence Speed — Iterations to 95% of Best Score",
                 fontsize=11, fontweight="bold")
    plt.tight_layout()
    out = out_dir / "convergence_speed.png"
    plt.savefig(out, dpi=150)
    plt.close(fig)
    print(f"  [Plot] Convergence speed       -> {out}")


# =============================================================================
# Baseline factory
# =============================================================================

def build_baseline(name: str, pop: int):
    n = name.upper()
    if n == "PSO":  return PSO(pop_size=pop)
    if n == "GWO":  return GWO(pop_size=pop)
    if n == "GA":   return GA(pop_size=pop)
    if n == "DE":   return DE(pop_size=pop)
    if n == "JADE": return JADE(pop_size=pop)
    if n == "WOA":  return WOA(pop_size=pop)
    if n == "ABC":  return ABC(pop_size=pop)
    if n == "FA":   return FA(pop_size=pop)
    if n == "SA":   return SA()
    if n == "ACO":  return ACO(pop_size=pop)
    raise ValueError(f"Unknown baseline: {name}")


AVAILABLE_METHODS = ["DPO", "DPO-NoDebt", "DPO-NoAccept", "DPO-NoRepay", 
                     "JADE", "DE", "PSO", "GWO", "ABC", "WOA", "GA", "FA", "SA", "ACO"]


# =============================================================================
# Main
# =============================================================================

def main() -> int:
    import sys
    if hasattr(sys.stdout, "reconfigure"):
        sys.stdout.reconfigure(encoding="utf-8", errors="replace")

    parser = argparse.ArgumentParser(
        description="Comprehensive HPO Publication Benchmark  (NASBench201 + HPOBench + HPOLib)",
        formatter_class=argparse.ArgumentDefaultsHelpFormatter,
    )
    parser.add_argument("--population",  type=int,  default=40,
                        help="Population / swarm size")
    parser.add_argument("--iterations",  type=int,  default=60,
                        help="Max iterations per run")
    parser.add_argument("--seeds",       type=int,  default=10,
                        help="Independent seeds for statistics (>=10 recommended)")
    parser.add_argument(
        "--families", nargs="+",
        default=["nasbench201", "hpobench", "hpolib"],
        choices=list(FAMILY_DATASETS.keys()),
        help="Benchmark families to include (synthetic = analytic test functions)",
    )
    parser.add_argument(
        "--datasets", nargs="+", default=None,
        choices=list(ALL_DATASETS.keys()),
        help="Override: pick specific datasets (ignores --families)",
    )
    parser.add_argument(
        "--methods", nargs="+", default=AVAILABLE_METHODS,
        choices=AVAILABLE_METHODS,
        help="Algorithms to benchmark",
    )
    parser.add_argument("--out-dir",  type=str, default="hpo_benchmark_results",
                        help="Output directory for plots and JSON")
    parser.add_argument("--no-plots", action="store_true", help="Skip all plots")
    parser.add_argument("--no-latex", action="store_true", help="Skip LaTeX table")
    args = parser.parse_args()

    # Resolve datasets
    if args.datasets:
        datasets = args.datasets
    else:
        datasets = []
        for fam in args.families:
            datasets.extend(FAMILY_DATASETS[fam])

    seeds   = list(range(args.seeds))
    methods = args.methods
    out_dir = Path(args.out_dir)
    out_dir.mkdir(parents=True, exist_ok=True)

    n_total = len(datasets) * len(methods) * len(seeds)
    print(f"\n{DIVIDER_WIDE}")
    print(f"  DPO -- Comprehensive HPO Publication Benchmark")
    print(f"  Families : {args.families}")
    print(f"  Datasets : {datasets}")
    print(f"  Methods  : {methods}")
    print(f"  Seeds    : {seeds}")
    print(f"  Pop      : {args.population}   Iterations: {args.iterations}")
    eval_budget = args.population * (args.iterations + 1)
    print(f"  Eval budget per run: {eval_budget}  (= pop * (iters+1))")
    print(f"  Note: DPO uses domain-specific presets; baselines use recommended defaults")
    print(f"  Total runs: {n_total}")
    print(f"  Output   : {out_dir.resolve()}")
    print(DIVIDER_WIDE)

    # Storage
    all_results: Dict[str, Dict[str, List[RunResult]]] = {
        ds: {m: [] for m in methods} for ds in datasets
    }

    run_idx = 0
    for ds_key in datasets:
        cfg = ALL_DATASETS[ds_key]
        print(f"\n  >> [{cfg.family.upper()}] {cfg.display_name}  ({cfg.description})")
        for seed in seeds:
            for m in methods:
                # Fresh evaluator per (dataset, seed, method) to avoid
                # IOH internal state leaking across algorithms.
                evaluator = make_evaluator(cfg, seed=seed)
                evaluator.set_eval_budget(eval_budget)
                run_idx += 1
                label = f"[{run_idx}/{n_total}]  {m:<14}  {cfg.display_name}  seed={seed}"
                print(f"    {label} ...", end=" ", flush=True)
                t0 = time.perf_counter()
                try:
                    if m in ["DPO", "DPO-NoDebt", "DPO-NoAccept", "DPO-NoRepay"]:
                        # Map method name to variant
                        variant_map = {
                            "DPO": "full",
                            "DPO-NoDebt": "nodebt",
                            "DPO-NoAccept": "noaccept",
                            "DPO-NoRepay": "norepay",
                        }
                        result = run_dpo(evaluator, cfg, args.population,
                                        args.iterations, seed, 
                                        variant=variant_map[m])
                    else:
                        bl     = build_baseline(m, args.population)
                        result = run_baseline(m, bl, evaluator, cfg,
                                              args.iterations, seed)
                    all_results[ds_key][m].append(result)
                    elapsed = time.perf_counter() - t0
                    print(f"score={_fmt_score(result.best_score)}  auc={result.auc:.4f}  t={elapsed:.1f}s")
                except Exception as exc:
                    import traceback
                    print(f"ERROR: {exc}")
                    traceback.print_exc()

    # Cross-method AUC normalisation: for each (dataset, seed), find the
    # global score range across all methods and recompute AUC so that
    # values are directly comparable between algorithms.
    for ds in datasets:
        for seed_val in seeds:
            all_curves = [
                r.score_curve
                for m in methods
                for r in all_results[ds].get(m, [])
                if r.seed == seed_val and r.score_curve
            ]
            if not all_curves:
                continue
            ref_start = min(c[0] for c in all_curves)
            ref_end   = max(max(c) for c in all_curves)
            for m in methods:
                for r in all_results[ds].get(m, []):
                    if r.seed == seed_val and r.score_curve:
                        r.auc = auc_from_curve(r.score_curve, ref_start, ref_end)

    # Aggregate (with cross-method normalised AUC)
    all_stats: Dict[str, Dict[str, Dict]] = {}
    for ds in datasets:
        all_stats[ds] = {}
        for m in methods:
            rs = all_results[ds].get(m, [])
            if rs:
                all_stats[ds][m] = aggregate(rs)

    # Wilcoxon tests
    wilcoxon_res = wilcoxon_dpo_vs(all_results, datasets, methods)

    # Per-dataset tables
    for ds in datasets:
        print_dataset_table(
            ALL_DATASETS[ds],
            all_stats[ds],
            methods,
            wilcoxon_res.get(ds, {}),
        )

    # Family summaries
    mean_ranks = friedman_mean_ranks(all_stats, datasets, methods)
    for fam in FAMILY_DATASETS:
        print_family_summary(fam, all_stats, datasets, methods, mean_ranks)

    # Overall summary
    print_overall_summary(all_stats, datasets, methods, mean_ranks)

    # LaTeX table
    if not args.no_latex:
        print_latex_table(all_stats, datasets, methods, mean_ranks)

    # Save JSON
    json_path = out_dir / "results.json"
    serialisable: Dict = {}
    for ds in datasets:
        serialisable[ds] = {}
        for m in methods:
            serialisable[ds][m] = [
                {
                    "seed":           r.seed,
                    "best_score":     r.best_score,
                    "best_fitness":   r.best_fitness,
                    "raw_metric":     r.raw_metric,
                    "metric_name":    r.metric_name,
                    "direction":      r.direction,
                    "auc":            r.auc,
                    "time_s":         r.time_s,
                    "iters_to_95pct": r.iters_to_95pct,
                    "iters_to_99pct": r.iters_to_99pct,
                }
                for r in all_results[ds].get(m, [])
            ]
    with open(json_path, "w") as fh:
        json.dump(serialisable, fh, indent=2)
    print(f"\n  [JSON] Results saved           -> {json_path}")

    # Plots
    if not args.no_plots:
        print(f"\n  Generating plots ...")
        plot_convergence(      all_stats,   datasets, methods, out_dir)
        plot_score_bars(       all_stats,   datasets, methods, out_dir)
        plot_boxplots(         all_results, datasets, methods, out_dir)
        plot_rank_heatmap(     all_stats,   datasets, methods, out_dir)
        plot_auc_lineplot(     all_stats,   datasets, methods, out_dir)
        plot_convergence_speed(all_stats,   datasets, methods, out_dir)

    print(f"\n  Done.  All outputs in: {out_dir.resolve()}\n")
    return 0


if __name__ == "__main__":
    raise SystemExit(main())
