

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)

# -- 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"
    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) --------------------
_SYN_CITE = "Molga & Smutnicki 2005; Ackley 1987; Rastrigin 1974; Schwefel 1977"

SYNTHETIC_DATASETS: Dict[str, DatasetConfig] = {
    "ackley_20d": DatasetConfig(
        key="ackley_20d",     family="synthetic",  bench_name="ackley_20d",
        display_name="Ackley-20D",
        description="Ackley (20D) – highly multimodal, global min = 0 at origin",
        citation=_SYN_CITE,
    ),
    "rastrigin_20d": DatasetConfig(
        key="rastrigin_20d",  family="synthetic",  bench_name="rastrigin_20d",
        display_name="Rastrigin-20D",
        description="Rastrigin (20D) – ~400 local minima, global min = 0 at origin",
        citation=_SYN_CITE,
    ),
    "schwefel_20d": DatasetConfig(
        key="schwefel_20d",   family="synthetic",  bench_name="schwefel_20d",
        display_name="Schwefel-20D",
        description="Schwefel (20D) – deceptive landscape, global min ≈ 0 near boundary",
        citation=_SYN_CITE,
    ),
    "levy_20d": DatasetConfig(
        key="levy_20d",       family="synthetic",  bench_name="levy_20d",
        display_name="Lévy-20D",
        description="Lévy (20D) – several local minima, global min = 0 at (1,…,1)",
        citation=_SYN_CITE,
    ),
    "rosenbrock_20d": DatasetConfig(
        key="rosenbrock_20d", family="synthetic",  bench_name="rosenbrock_20d",
        display_name="Rosenbrock-20D",
        description="Rosenbrock (20D) – narrow curved valley, global min = 0 at (1,…,1)",
        citation=_SYN_CITE,
    ),
    "griewank_20d": DatasetConfig(
        key="griewank_20d",   family="synthetic",  bench_name="griewank_20d",
        display_name="Griewank-20D",
        description="Griewank (20D) – dense local-optima grid from interaction terms",
        citation=_SYN_CITE,
    ),
}

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

FAMILY_DATASETS = {
    "nasbench201": list(NASBENCH201_DATASETS.keys()),
    "hpobench":   list(HPOBENCH_DATASETS.keys()),
    "hpolib":     list(HPOLIB_DATASETS.keys()),
    "synthetic":  list(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": []}

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

    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 = self._raw_query(vec)
        score   = self.raw_to_score(raw)
        fitness = self.score_to_fitness(score)
        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   = self._raw_query(vec)
            score = self.raw_to_score(raw)
            self._update_tracker(score, vec)
            return self.score_to_fitness(score)
        return _fn

    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
# =============================================================================

class SyntheticEvaluator(BaseEvaluator):
    """
    Classic analytic benchmark functions (minimisation).

    Chosen because they are:
      - Multimodal or deceptive  -> separate explorers from exploiters
      - Independent of external data files  -> always available
      - 20-dimensional           -> realistic HPO dimensionality

    Score = -f(x)  (higher = better, 0 = global optimum)
    """

    _REGISTRY: Dict[str, Dict] = {
        "ackley_20d": {
            "dims": 20, "lo": -32.768, "hi": 32.768,
            "description": "Ackley (20D) – highly multimodal, global min = 0 at origin",
        },
        "rastrigin_20d": {
            "dims": 20, "lo": -5.12, "hi": 5.12,
            "description": "Rastrigin (20D) – ~400 local minima, global min = 0 at origin",
        },
        "schwefel_20d": {
            "dims": 20, "lo": -500.0, "hi": 500.0,
            "description": "Schwefel (20D) – deceptive, global min ≈ 0 near boundary",
        },
        "levy_20d": {
            "dims": 20, "lo": -10.0, "hi": 10.0,
            "description": "Lévy (20D) – several local minima, global min = 0 at (1,…,1)",
        },
        "rosenbrock_20d": {
            "dims": 20, "lo": -5.0, "hi": 10.0,
            "description": "Rosenbrock (20D) – narrow curved valley, global min = 0 at (1,…,1)",
        },
        "griewank_20d": {
            "dims": 20, "lo": -600.0, "hi": 600.0,
            "description": "Griewank (20D) – interaction terms make local optima grid dense",
        },
    }

    def __init__(self, bench_name: str, seed: int = 0):
        super().__init__()
        np.random.seed(seed)
        cfg              = self._REGISTRY[bench_name]
        self._fn_name    = bench_name
        self._dims       = cfg["dims"]
        lo, hi           = cfg["lo"], cfg["hi"]
        self.param_names = [f"x{i}" for i in range(self._dims)]
        self.bounds_low  = np.full(self._dims, lo,  dtype=np.float64)
        self.bounds_high = np.full(self._dims, hi,  dtype=np.float64)
        self.metric_name = "f_value"
        self.direction   = "minimize"

    def _raw_query(self, vec: np.ndarray) -> float:
        x  = np.asarray(vec, dtype=np.float64)
        fn = self._fn_name

        if fn == "ackley_20d":
            a, b, c = 20.0, 0.2, 2.0 * np.pi
            d       = len(x)
            val = (-a * np.exp(-b * np.sqrt(np.mean(x ** 2)))
                   - np.exp(np.mean(np.cos(c * x)))
                   + a + np.e)
            return max(0.0, float(val))

        elif fn == "rastrigin_20d":
            d   = len(x)
            val = 10 * d + float(np.sum(x ** 2 - 10.0 * np.cos(2.0 * np.pi * x)))
            return max(0.0, val)

        elif fn == "schwefel_20d":
            d   = len(x)
            val = 418.9829 * d - float(np.sum(x * np.sin(np.sqrt(np.abs(x)))))
            return float(val)        # can be slightly negative near optimum

        elif fn == "levy_20d":
            w     = 1.0 + (x - 1.0) / 4.0
            term1 = float(np.sin(np.pi * w[0]) ** 2)
            term2 = float(np.sum(
                (w[:-1] - 1.0) ** 2 * (1.0 + 10.0 * np.sin(np.pi * w[:-1] + 1.0) ** 2)
            ))
            term3 = float((w[-1] - 1.0) ** 2 * (1.0 + np.sin(2.0 * np.pi * w[-1]) ** 2))
            return max(0.0, term1 + term2 + term3)

        elif fn == "rosenbrock_20d":
            val = float(np.sum(100.0 * (x[1:] - x[:-1] ** 2) ** 2 + (x[:-1] - 1.0) ** 2))
            return max(0.0, val)

        elif fn == "griewank_20d":
            d      = len(x)
            i_arr  = np.arange(1, d + 1, dtype=np.float64)
            val    = float(np.sum(x ** 2) / 4000.0
                           - float(np.prod(np.cos(x / np.sqrt(i_arr)))) + 1.0)
            return max(0.0, val)

        raise ValueError(f"Unknown synthetic function: {fn}")


# =============================================================================
# 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)
    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]) -> float:
    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)
    # and normalized to remove scale/sign effects across datasets.
    ys = np.maximum.accumulate(ys)
    if ys.size == 1:
        return 1.0

    start = float(ys[0])
    end   = float(ys[-1])
    span  = end - start
    if abs(span) <= 1e-12:
        # Already converged (or no measurable progress): treat as maximal speed.
        return 1.0

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

    # Integrate over normalized 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_threshold(curve: List[float], threshold: float) -> int:
    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,
) -> 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,
        )
    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))

    # 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]

    # 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         = "DPO",
        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_threshold(score_curve, 0.95 * peak) if peak > 0 else -1,
        iters_to_99pct = iter_to_threshold(score_curve, 0.99 * peak) if peak > 0 else -1,
        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_threshold(score_curve, 0.95 * peak) if peak > 0 else -1,
        iters_to_99pct = iter_to_threshold(score_curve, 0.99 * peak) if peak > 0 else -1,
        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 (DPO > baseline, one-sided, α=0.05) per dataset."""
    if not HAS_SCIPY:
        return {}
    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
            try:
                stat, pval = wilcoxon(dpo_vals, base_vals, alternative="greater")
                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]:
    rank_sum: Dict[str, float] = defaultdict(float)
    for ds in datasets:
        valid = {m: all_stats[ds][m] for m in methods if all_stats[ds].get(m)}
        if not valid:
            continue
        sorted_m = sorted(valid.keys(), key=lambda m: -valid[m]["score_mean"])
        for rank, m in enumerate(sorted_m, 1):
            rank_sum[m] += rank
    return {m: rank_sum[m] / max(1, len(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 "↓ (–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 > baseline, one-sided, α = 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],
) -> None:
    fam_datasets = [d for d in datasets if ALL_DATASETS[d].family == family]
    if not fam_datasets:
        return
    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}  {'Mean Score':>12}  {'Mean AUC':>12}")
    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: 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)
        ]))
        mk = "  ◀ BEST" if m == min(valid_methods, key=lambda x: ranks.get(x, 999)) else ""
        print(
            f"{m:<12}  {ranks.get(m, float('nan')):>10.2f}  {wins[m]:>8d}"
            f"  {_fmt_score(mean_score):>12}  {mean_auc:>12.4f}{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}  {'Mean Score':>12}  {'Mean AUC':>12}")
    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)]
        oms  = float(np.mean(sc_vals))  if sc_vals  else float("nan")
        oma  = float(np.mean(auc_vals)) if auc_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}{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) + "r"
    print(r"\begin{table}[t]")
    print(r"\centering")
    print(r"\caption{Comprehensive HPO Benchmark: mean score $\pm$ std (all runs)}")
    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}} \\\\")
    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)
        print(f"{m} & {' & '.join(cols)} & {mean_ranks[m]:.2f} \\\\")
    print(r"\bottomrule")
    print(r"\end{tabular}")
    print(r"\end{table}")


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

PALETTE = {
    "DPO":  "#e74c3c",
    "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",
}


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")
            style = DPO_STYLE if m == "DPO" 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 == "DPO":
                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")
        style  = DPO_STYLE if m == "DPO" 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", "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=5,
                        help="Independent seeds for statistics")
    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}")
    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:
            evaluator = make_evaluator(cfg, seed=seed)
            for m in methods:
                run_idx += 1
                label = f"[{run_idx}/{n_total}]  {m:<6}  {cfg.display_name}  seed={seed}"
                print(f"    {label} …", end=" ", flush=True)
                t0 = time.perf_counter()
                try:
                    if m == "DPO":
                        result = run_dpo(evaluator, cfg, args.population,
                                         args.iterations, seed)
                    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()

    # Aggregate
    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())
