scripts/evolve_meta_model.py

#!/usr/bin/env python3
"""Evolve a dataset-conditioned BERNN policy across whole biological datasets."""

from __future__ import annotations

import argparse
import csv
import hashlib
import importlib.metadata
import json
import os
import sys
import uuid
from pathlib import Path

import numpy as np

ROOT = Path(__file__).resolve().parent.parent
if str(ROOT) not in sys.path:
    sys.path.insert(0, str(ROOT))

from scripts import hp_search
from src.dataset_splits import load_dataset_partitions
from src.evolutionary_meta import (
    EvolutionConfig,
    PolicyShape,
    aggregate_dataset_scores,
    breed_next_population,
    config_dict,
    decode_config,
    genome_digest,
    initialize_population,
    normalize_meta_features,
    recommended_batch_size,
)
from src.zero_shot_recommender.meta_features import (
    META_FEATURE_NAMES,
    extract_meta_features,
)
from src.meta_policy import META_POLICY_REPO, load_policy_metadata, publish_policy_if_improved

SCORE_THRESHOLDS = {
    "massbench_adenocarcinoma": {"reference": 0.24, "acceptable": 0.19},
    "massbench_alzheimer": {"reference": 0.44, "acceptable": 0.39},
    "massbench_benchmark": {"reference": 0.75, "acceptable": 0.70},
}


def parse_args(argv=None):
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--split-manifest", type=Path, default=ROOT / "config" / "evolution_development_datasets.json")
    parser.add_argument("--output-dir", type=Path, default=ROOT / "results" / "meta_evolution")
    parser.add_argument("--population-size", type=int, default=12)
    parser.add_argument("--generations", type=int, default=25)
    parser.add_argument(
        "--candidate-budget", type=int, default=None,
        help="Exact total number of source-only candidate policies to evaluate; "
             "defaults to population-size * generations and may end mid-generation.",
    )
    parser.add_argument("--hidden-size", type=int, default=16)
    parser.add_argument("--elite-count", type=int, default=2)
    parser.add_argument("--tournament-size", type=int, default=3)
    parser.add_argument("--crossover-rate", type=float, default=0.9)
    parser.add_argument("--mutation-rate", type=float, default=0.05)
    parser.add_argument("--mutation-scale", type=float, default=0.1)
    parser.add_argument("--worst-dataset-weight", type=float, default=0.25)
    parser.add_argument("--validation-patience", type=int, default=8)
    parser.add_argument("--n-epochs", type=int, default=100)
    parser.add_argument("--n-repeats", type=int, default=3)
    parser.add_argument("--num-workers", type=int, default=4)
    parser.add_argument(
        "--validation-epoch-interval", type=int, default=10,
        help="Log held-out validation epoch telemetry at epoch 0 and this interval.",
    )
    parser.add_argument("--device", default="cuda", choices=["cpu", "cuda"])
    parser.add_argument("--seed", type=int, default=42)
    parser.add_argument("--resume", action="store_true")
    parser.add_argument("--smoke-evaluator", action="store_true", help="Exercise evolution without training BERNN")
    parser.add_argument("--no-wandb", action="store_true", help="Disable live Weights & Biases telemetry")
    parser.add_argument("--wandb-project", default="BE_leaderboard_meta_evolution")
    parser.add_argument("--wandb-run-name", default=None)
    parser.add_argument("--hf-meta-policy-repo", default=META_POLICY_REPO)
    parser.add_argument(
        "--no-hf-meta-policy-push", action="store_true",
        help="Do not compare/promote the validation champion on Hugging Face.",
    )
    return parser.parse_args(argv)


def _sha256(path: Path) -> str:
    return hashlib.sha256(path.read_bytes()).hexdigest()


def _atomic_json(path: Path, value: dict) -> None:
    temporary = path.with_suffix(path.suffix + ".tmp")
    temporary.write_text(json.dumps(value, indent=2, default=str) + "\n")
    os.replace(temporary, path)


def _load_datasets(dataset_ids):
    datasets = {}
    meta = {}
    for dataset_id in dataset_ids:
        X, y, batches = hp_search.load_dataset(dataset_id)
        datasets[dataset_id] = (X, y, batches)
        extracted = extract_meta_features(X, y, batches)
        meta[dataset_id] = np.asarray([extracted[name] for name in META_FEATURE_NAMES], dtype=np.float32)
    return datasets, meta



def predict_config_for_dataset(
    X,
    y,
    batches,
    *,
    model_path=None,
    meta_model_path=None,
    checkpoint_path=None,
    max_warmup=50,
    **kwargs,
):
    """Predict exactly one BERNN config for one unseen dataset.

    The frozen evolutionary meta-policy is loaded from ``best_policy.npz``.
    Meta-features for the unseen dataset are normalized with the mean/scale
    learned from meta-training datasets and stored in that artifact.
    """
    policy_path = model_path or meta_model_path or checkpoint_path
    if policy_path is None:
        policy_path = ROOT / "results" / "meta_evolution" / "best_policy.npz"

    policy_path = Path(policy_path)
    if policy_path.is_dir():
        policy_path = policy_path / "best_policy.npz"
    if not policy_path.exists():
        raise FileNotFoundError(
            f"Meta-policy not found: {policy_path}. "
            "Pass --meta-model-path /path/to/best_policy.npz."
        )

    saved = np.load(policy_path)
    required = {"genome", "meta_mean", "meta_scale", "n_inputs", "hidden_size"}
    missing = required.difference(saved.files)
    if missing:
        raise ValueError(
            f"Meta-policy {policy_path} is missing required arrays: {sorted(missing)}"
        )

    genome = saved["genome"]
    meta_mean = np.asarray(saved["meta_mean"], dtype=np.float32)
    meta_scale = np.asarray(saved["meta_scale"], dtype=np.float32)
    n_inputs = int(np.asarray(saved["n_inputs"]).reshape(-1)[0])
    hidden_size = int(np.asarray(saved["hidden_size"]).reshape(-1)[0])

    extracted = extract_meta_features(X, y, batches)
    raw_meta = np.asarray(
        [extracted[name] for name in META_FEATURE_NAMES],
        dtype=np.float32,
    )

    if raw_meta.shape[0] != n_inputs:
        raise ValueError(
            f"Meta-feature count mismatch: policy expects {n_inputs}, "
            f"current extractor produced {raw_meta.shape[0]}."
        )
    if meta_mean.shape != raw_meta.shape or meta_scale.shape != raw_meta.shape:
        raise ValueError(
            "Saved meta normalization shape does not match current meta-features: "
            f"mean={meta_mean.shape}, scale={meta_scale.shape}, current={raw_meta.shape}."
        )

    safe_scale = np.where(np.abs(meta_scale) < 1e-12, 1.0, meta_scale)
    normalized_meta = (raw_meta - meta_mean) / safe_scale

    shape = PolicyShape(n_inputs=n_inputs, hidden_size=hidden_size)
    return decode_config(
        genome,
        normalized_meta,
        shape,
        max_warmup=int(max_warmup),
    )


def _append_solution_record(output_dir: Path, record: dict, train_datasets: list[str]) -> None:
    jsonl_path = output_dir / "solutions.jsonl"
    with jsonl_path.open("a", encoding="utf-8") as stream:
        stream.write(json.dumps(record, default=str) + "\n")

    csv_path = output_dir / "solutions.csv"
    columns = [
        "solution_step",
        "generation",
        "population_index",
        "genome_digest",
        "train_fitness",
        "mean_mcc",
        "worst_mcc",
        *[f"mcc_{name}" for name in train_datasets],
        *[f"test_mcc_{name}" for name in train_datasets],
    ]
    row = {name: record.get(name) for name in columns}
    for name in train_datasets:
        row[f"mcc_{name}"] = record["train_scores"][name]
        row[f"test_mcc_{name}"] = record["test_scores"].get(name)
    write_header = not csv_path.exists() or csv_path.stat().st_size == 0
    with csv_path.open("a", newline="", encoding="utf-8") as stream:
        writer = csv.DictWriter(stream, fieldnames=columns)
        if write_header:
            writer.writeheader()
        writer.writerow(row)


def _mean_epoch_metric(metrics: dict, metric_name: str) -> dict[int, float]:
    """Average one BERNN epoch metric over grouped outer-CV folds/repeats."""
    values: dict[int, list[float]] = {}
    for rows in metrics.get("_epoch_traces_by_fold", {}).values():
        for row in rows:
            epoch = int(row["epoch"])
            matching = [
                float(value) for key, value in row.items()
                if key == metric_name or key.endswith(f"/{metric_name}")
            ]
            if matching:
                values.setdefault(epoch, []).extend(matching)
    return {epoch: float(np.mean(epoch_values)) for epoch, epoch_values in values.items()}


def _classifier_widths(config: dict) -> list[int]:
    widths = [max(16, int(config["layer1"]))]
    for _ in range(1, max(1, int(config["n_layers"]))):
        widths.append(max(widths[-1] // 2, 16))
    return widths


def _numeric_config_telemetry(dataset_id: str, config: dict) -> dict[str, float]:
    telemetry = {}
    for name, value in config.items():
        if isinstance(value, (bool, np.bool_)):
            telemetry[f"solutions/config/{dataset_id}/{name}"] = float(value)
        elif isinstance(value, (int, float, np.integer, np.floating)):
            telemetry[f"solutions/config/{dataset_id}/{name}"] = float(value)
    widths = _classifier_widths(config)
    telemetry[f"solutions/config/{dataset_id}/n_neurons"] = float(sum(widths))
    for layer_index, width in enumerate(widths, 1):
        telemetry[f"solutions/config/{dataset_id}/layer{layer_index}_width"] = float(width)
    return telemetry


def _score_figure(
    dataset_id: str,
    rows: list[tuple[int, float, float]],
    thresholds: dict[str, float] | None = None,
):
    """Plot validation/test MCC plus BERNN-best and acceptable-goal traces."""
    import plotly.graph_objects as go

    x = [row[0] for row in rows]
    valid = [row[1] for row in rows]
    test = [row[2] for row in rows]
    figure = go.Figure()
    figure.add_trace(go.Scatter(x=x, y=valid, mode="lines+markers", name="valid_mcc"))
    figure.add_trace(go.Scatter(x=x, y=test, mode="lines+markers", name="test_mcc"))
    thresholds = dict(SCORE_THRESHOLDS.get(dataset_id, {}) if thresholds is None else thresholds)
    if "reference" in thresholds:
        figure.add_trace(go.Scatter(
            x=x, y=[thresholds["reference"]] * len(x), mode="lines",
            line={"dash": "solid", "width": 2}, name="Best valid MCC",
        ))
    if "acceptable" in thresholds:
        figure.add_trace(go.Scatter(
            x=x, y=[thresholds["acceptable"]] * len(x), mode="lines",
            line={"dash": "dash", "width": 2}, name="goal",
        ))
    figure.update_layout(
        title=f"{dataset_id}: MCC versus BERNN thresholds",
        xaxis_title="solution step", yaxis_title="grouped-CV MCC",
        yaxis_range=[-1.0, 1.0],
    )
    return figure


def _solution_records(output_dir: Path) -> list[dict]:
    path = output_dir / "solutions.jsonl"
    if not path.exists():
        return []
    return [json.loads(line) for line in path.read_text().splitlines() if line.strip()]


def _save_source_champion(
    output_dir: Path,
    genome: np.ndarray,
    *,
    mean: np.ndarray,
    scale: np.ndarray,
    shape: PolicyShape,
    record: dict,
) -> None:
    np.savez_compressed(
        output_dir / "best_policy.npz",
        genome=np.asarray(genome, dtype=np.float32),
        meta_mean=mean,
        meta_scale=scale,
        n_inputs=np.asarray([shape.n_inputs]),
        hidden_size=np.asarray([shape.hidden_size]),
    )
    _atomic_json(output_dir / "best_policy.json", record)


def _replay_source_champion(
    output_dir: Path,
    shape: PolicyShape,
    evolution: EvolutionConfig,
    seed: int,
) -> tuple[np.ndarray, dict] | None:
    """Reconstruct evaluated genomes without rerunning any BERNN fit.

    Evolution and breeding use source fitness only. Replaying the deterministic
    initializer and breeder therefore recovers the exact genome corresponding
    to every durable solution record, including legacy runs whose saved policy
    was selected with target validation telemetry.
    """
    records = _solution_records(output_dir)
    if not records:
        return None
    replay_rng = np.random.default_rng(seed)
    population = initialize_population(shape, evolution, replay_rng)
    best_genome = None
    best_record = None
    generations = sorted({int(row["generation"]) for row in records})
    for generation in generations:
        rows = sorted(
            (row for row in records if int(row["generation"]) == generation),
            key=lambda row: int(row["population_index"]),
        )
        fitness = np.full(evolution.population_size, np.nan, dtype=float)
        for row in rows:
            index = int(row["population_index"])
            if not 0 <= index < evolution.population_size:
                raise ValueError(f"Invalid population index in generation {generation}: {index}")
            observed_digest = str(row["genome_digest"])
            recovered_digest = genome_digest(population[index])
            if observed_digest != recovered_digest:
                raise ValueError(
                    f"Cannot replay generation {generation} index {index}: "
                    f"digest {recovered_digest} != {observed_digest}"
                )
            value = float(row["train_fitness"])
            fitness[index] = value
            if best_record is None or value > float(best_record["train_fitness"]):
                best_genome = population[index].copy()
                best_record = {
                    "selection_protocol": "source_only_best_fitness",
                    "generation": generation,
                    "population_index": index,
                    "solution_step": int(row["solution_step"]),
                    "train_fitness": value,
                    "train_scores": row["train_scores"],
                    "genome_digest": observed_digest,
                    "target_scores_observed_for_selection": False,
                }
        if len(rows) == evolution.population_size:
            if not np.isfinite(fitness).all():
                raise ValueError(f"Incomplete fitness vector for generation {generation}")
            population = breed_next_population(population, fitness, evolution, replay_rng)
        elif generation != generations[-1]:
            raise ValueError(f"Only the final generation may be partial; generation {generation} is incomplete")
    return best_genome, best_record


def main(argv=None) -> int:
    args = parse_args(argv)
    if args.n_repeats != 3:
        raise ValueError("Meta-evolution requires grouped CV=3")
    if importlib.metadata.version("bernn") != "1.0.6":
        raise RuntimeError("This experiment requires bernn==1.0.6")
    if args.generations < 1:
        raise ValueError("generations must be positive")
    candidate_budget = int(
        args.candidate_budget
        if args.candidate_budget is not None
        else args.population_size * args.generations
    )
    if candidate_budget < 1:
        raise ValueError("candidate_budget must be positive")
    partitions = load_dataset_partitions(args.split_manifest)

    # Target datasets provide descriptors for a frozen recommendation only.
    # Their BERNN scores and fixed tests are never opened by source search.
    heldout_ids = tuple(getattr(partitions, "test", ()))
    datasets, raw_meta = _load_datasets(partitions.train)
    fixed_tests = {
        dataset_id: hp_search.load_fixed_test_dataset(dataset_id)
        for dataset_id in partitions.train
    }

    train_matrix = np.stack([raw_meta[name] for name in partitions.train])
    normalized_train, mean, scale = normalize_meta_features(train_matrix)
    normalized_meta = {
        name: row for name, row in zip(partitions.train, normalized_train)
    }

    evolution = EvolutionConfig(
        population_size=args.population_size,
        elite_count=args.elite_count,
        tournament_size=args.tournament_size,
        crossover_rate=args.crossover_rate,
        mutation_rate=args.mutation_rate,
        mutation_scale=args.mutation_scale,
        worst_dataset_weight=args.worst_dataset_weight,
    )
    evolution.validate()
    shape = PolicyShape(len(META_FEATURE_NAMES), args.hidden_size)
    args.output_dir.mkdir(parents=True, exist_ok=True)
    manifest_hash = _sha256(args.split_manifest)
    checkpoint_path = args.output_dir / "checkpoint.npz"
    state_path = args.output_dir / "state.json"
    rng = np.random.default_rng(args.seed)
    generation_start = 0
    best_train_fitness = -np.inf
    solution_step = 0

    if args.resume:
        if not checkpoint_path.exists() or not state_path.exists():
            raise FileNotFoundError("--resume requires checkpoint.npz and state.json")
        state = json.loads(state_path.read_text())
        if state["development_manifest_sha256"] != manifest_hash:
            raise ValueError("Development dataset manifest changed since checkpoint creation")
        expected_shape = {"n_inputs": shape.n_inputs, "hidden_size": shape.hidden_size}
        if state.get("policy_shape") != expected_shape:
            raise ValueError("Policy shape differs from the saved checkpoint")
        if state.get("evolution_config") != config_dict(evolution):
            raise ValueError("Evolution configuration differs from the saved checkpoint")
        saved = np.load(checkpoint_path)
        population = saved["population"]
        generation_start = int(state["next_generation"])
        best_train_fitness = float(state.get("best_train_fitness", -np.inf))
        solution_step = int(state.get("solution_step", generation_start * args.population_size))
        rng.bit_generator.state = state["rng_state"]
    else:
        population = initialize_population(shape, evolution, rng)

    hp_args = hp_search.parse_args([])
    hp_args.n_epochs = args.n_epochs
    hp_args.n_repeats = args.n_repeats
    hp_args.num_workers = args.num_workers
    hp_args.log1p = True
    hp_args.device = args.device
    hp_args.no_wandb = True
    hp_args.combine_test = False
    hp_args.max_warmup = max(1, min(50, args.n_epochs))
    score_cache = {}

    run_metadata_path = args.output_dir / "run_metadata.json"
    if run_metadata_path.exists():
        run_metadata = json.loads(run_metadata_path.read_text())
    else:
        run_metadata = {"wandb_run_id": uuid.uuid4().hex[:8]}
    run_metadata.update({
        "arm": "meta_evolution",
        "candidate_budget": candidate_budget,
        "budget_unit": "candidate_policy",
        "population_size": int(args.population_size),
        "generations": int(args.generations),
        "bernn_version": importlib.metadata.version("bernn"),
        "log1p": True,
        "num_workers": int(args.num_workers),
        "train_datasets": list(partitions.train),
        "validation_datasets": list(partitions.validation),
        "selection_protocol": "source_only_best_fitness",
        "target_scores_observed": False,
        "cross_test": 1,
        "fixed_test_labels": "monitoring_only_excluded_from_fitness_and_selection",
    })
    _atomic_json(run_metadata_path, run_metadata)

    wandb_run = None
    if not args.no_wandb:
        import wandb

        wandb_run = wandb.init(
            project=args.wandb_project,
            name=args.wandb_run_name,
            id=run_metadata["wandb_run_id"],
            resume="allow",
            config={
                **vars(args),
                "split_manifest": str(args.split_manifest),
                "output_dir": str(args.output_dir),
                "train_datasets": list(partitions.train),
                "validation_datasets": list(partitions.validation),
                "heldout_datasets": list(heldout_ids),
            },
        )
        wandb.define_metric("solution_step")
        wandb.define_metric("solutions/*", step_metric="solution_step")
        wandb.define_metric("generation")
        wandb.define_metric("champion/*", step_metric="generation")
        for dataset_id, thresholds in SCORE_THRESHOLDS.items():
            wandb_run.summary[f"reference_mcc/{dataset_id}"] = thresholds["reference"]
            wandb_run.summary[f"acceptable_mcc/{dataset_id}"] = thresholds["acceptable"]
        wandb_run.summary["selection_protocol"] = "source_only_best_fitness"
        wandb_run.summary["target_scores_observed_after_contract_fix"] = False

    def evaluate(genome, dataset_id):
        key = f"{genome_digest(genome)}:{dataset_id}"
        if key in score_cache:
            return score_cache[key]
        metrics = {}
        try:
            if args.smoke_evaluator:
                target = raw_meta[dataset_id]
                score = float(np.tanh(np.mean(genome[: min(len(genome), len(target))]) + np.mean(target) * 1e-4))
            else:
                X, y, batches = datasets[dataset_id]
                config = decode_config(genome, normalized_meta[dataset_id], shape, max_warmup=hp_args.max_warmup)
                run_args = argparse.Namespace(**vars(hp_args))
                run_args.dataset = dataset_id
                run_args.seed = args.seed
                run_args.bs = recommended_batch_size(batches, cap=hp_args.bs)
                run_args.results_dir = str(args.output_dir / dataset_id)
                run_args.cv_split_cache = str(args.output_dir / "cv_splits" / f"{dataset_id}.npz")
                run_args.resolved_n_repeats = hp_search.resolve_n_repeats(run_args.n_repeats, batches)
                exp_id = f"meta_evolution_{dataset_id}_{genome_digest(genome)[:12]}"
                score, metrics = hp_search.run_trial(
                    config, run_args, (X, y, batches), exp_id,
                    fixed_test_data=fixed_tests[dataset_id],
                )
        except Exception as exc:  # noqa: BLE001 - invalid candidates receive sentinel fitness
            score = -1.0
            print(
                f"[evolution] failed genome={genome_digest(genome)[:12]} "
                f"dataset={dataset_id}: {type(exc).__name__}: {exc}",
                flush=True,
            )
        score_cache[key] = (float(score), metrics)
        return score_cache[key]

    replayed = _replay_source_champion(args.output_dir, shape, evolution, args.seed)
    if replayed is not None:
        replayed_genome, replayed_record = replayed
        best_train_fitness = float(replayed_record["train_fitness"])
        _save_source_champion(
            args.output_dir, replayed_genome, mean=mean, scale=scale,
            shape=shape, record=replayed_record,
        )

    ledger = args.output_dir / "source_generations.jsonl"
    existing_solutions = _solution_records(args.output_dir)
    score_history = {dataset_id: [] for dataset_id in partitions.train}
    for row in existing_solutions:
        for dataset_id in partitions.train:
            test_score = float(row.get("test_scores", {}).get(dataset_id, np.nan))
            score_history[dataset_id].append((
                int(row["solution_step"]), float(row["train_scores"][dataset_id]), test_score,
            ))

    def save_state(next_generation: int, checkpoint_population: np.ndarray) -> None:
        np.savez_compressed(
            checkpoint_path, population=checkpoint_population,
            meta_mean=mean, meta_scale=scale,
        )
        _atomic_json(state_path, {
            "next_generation": next_generation,
            "best_train_fitness": best_train_fitness,
            "solution_step": solution_step,
            "rng_state": rng.bit_generator.state,
            "selection_protocol": "source_only_best_fitness",
            "development_manifest_sha256": manifest_hash,
            "policy_shape": {"n_inputs": shape.n_inputs, "hidden_size": shape.hidden_size},
            "evolution_config": config_dict(evolution),
            "train_datasets": list(partitions.train),
            "validation_datasets": list(partitions.validation),
            "heldout_datasets": list(heldout_ids),
        })

    for generation in range(generation_start, args.generations):
        prior_rows = sorted(
            (row for row in _solution_records(args.output_dir)
             if int(row["generation"]) == generation),
            key=lambda row: int(row["population_index"]),
        )
        expected_indices = list(range(len(prior_rows)))
        observed_indices = [int(row["population_index"]) for row in prior_rows]
        if observed_indices != expected_indices:
            raise ValueError(
                f"Generation {generation} durable records are not a contiguous prefix: "
                f"{observed_indices}"
            )
        train_scores = np.full((len(population), len(partitions.train)), np.nan, dtype=np.float64)
        train_fitness = np.full(len(population), np.nan, dtype=np.float64)
        for row in prior_rows:
            index = int(row["population_index"])
            if genome_digest(population[index]) != row["genome_digest"]:
                raise ValueError(f"Checkpoint/ledger digest mismatch at generation {generation} index {index}")
            train_fitness[index] = float(row["train_fitness"])
            train_scores[index] = [float(row["train_scores"][name]) for name in partitions.train]

        for population_index in range(len(prior_rows), len(population)):
            if solution_step >= candidate_budget:
                break
            genome = population[population_index]
            scores = np.asarray([evaluate(genome, name)[0] for name in partitions.fitness_datasets()])
            fitness = aggregate_dataset_scores(scores, evolution.worst_dataset_weight)
            train_scores[population_index] = scores
            train_fitness[population_index] = fitness
            decoded = {
                name: decode_config(genome, normalized_meta[name], shape, max_warmup=hp_args.max_warmup)
                for name in partitions.train
            }
            solution_record = {
                "solution_step": solution_step,
                "generation": generation,
                "population_index": population_index,
                "genome_digest": genome_digest(genome),
                "train_fitness": float(fitness),
                "mean_mcc": float(np.mean(scores)),
                "worst_mcc": float(np.min(scores)),
                "train_scores": dict(zip(partitions.train, scores.tolist())),
                "test_scores": {
                    name: float(evaluate(genome, name)[1].get("test_mcc", np.nan))
                    for name in partitions.train
                },
                "config_by_dataset": decoded,
                "target_scores_observed": False,
            }
            solution_record["total_valid_mcc"] = float(fitness)
            solution_record["total_test_mcc"] = float(aggregate_dataset_scores(
                solution_record["test_scores"].values(),
                evolution.worst_dataset_weight,
            ))
            for dataset_id, score in zip(partitions.train, scores):
                score_history[dataset_id].append((
                    solution_step, float(score), solution_record["test_scores"][dataset_id],
                ))
            _append_solution_record(args.output_dir, solution_record, list(partitions.train))

            if fitness > best_train_fitness:
                best_train_fitness = float(fitness)
                champion_record = {
                    "selection_protocol": "source_only_best_fitness",
                    "generation": generation,
                    "population_index": population_index,
                    "solution_step": solution_step,
                    "train_fitness": float(fitness),
                    "train_scores": solution_record["train_scores"],
                    "genome_digest": solution_record["genome_digest"],
                    "target_scores_observed_for_selection": False,
                }
                _save_source_champion(
                    args.output_dir, genome, mean=mean, scale=scale,
                    shape=shape, record=champion_record,
                )
            if wandb_run is not None:
                config_telemetry = {}
                for dataset_id, dataset_config in decoded.items():
                    config_telemetry.update(_numeric_config_telemetry(dataset_id, dataset_config))
                figure_telemetry = {
                    f"solutions/mcc/{dataset_id}": _score_figure(
                        dataset_id, score_history[dataset_id]
                    )
                    for dataset_id in partitions.train
                }
                wandb_run.log({
                    "solution_step": solution_step,
                    "solutions/train_fitness": float(fitness),
                    "solutions/mean_mcc": float(np.mean(scores)),
                    "solutions/worst_mcc": float(np.min(scores)),
                    "solutions/total_valid_mcc": solution_record["total_valid_mcc"],
                    **{
                        f"solutions/metrics/{name}/valid_mcc": float(score)
                        for name, score in zip(partitions.train, scores)
                    },
                    **{
                        f"solutions/metrics/{name}/test_mcc": solution_record["test_scores"][name]
                        for name in partitions.train
                    },
                    "solutions/total_test_mcc": solution_record["total_test_mcc"],
                    **config_telemetry,
                    **figure_telemetry,
                })
            print(f"[evolution] source solution {solution_step} complete", flush=True)
            solution_step += 1

        evaluated = int(np.isfinite(train_fitness).sum())
        if evaluated < len(population):
            save_state(generation, population)
            print(
                f"[evolution] exact candidate budget reached at {solution_step}; "
                f"generation {generation} is {evaluated}/{len(population)} complete",
                flush=True,
            )
            break

        champion_index = int(np.argmax(train_fitness))
        source_record = {
            "generation": generation,
            "selection_protocol": "source_only_best_fitness",
            "train_champion_fitness": float(train_fitness[champion_index]),
            "train_scores": dict(zip(partitions.train, train_scores[champion_index].tolist())),
            "champion_digest": genome_digest(population[champion_index]),
            "target_scores_observed": False,
        }
        with ledger.open("a", encoding="utf-8") as stream:
            stream.write(json.dumps(source_record, default=str) + "\n")
        if wandb_run is not None:
            wandb_run.log({
                "generation": generation,
                "champion/train_fitness": float(train_fitness[champion_index]),
            })
        next_population = breed_next_population(population, train_fitness, evolution, rng)
        save_state(generation + 1, next_population)
        population = next_population

    best_policy_path = args.output_dir / "best_policy.npz"
    if not best_policy_path.exists():
        raise FileNotFoundError("Source-only evolution did not produce best_policy.npz")
    best_policy = np.load(best_policy_path)
    best_genome = best_policy["genome"]
    _, target_raw_meta = _load_datasets(partitions.validation)
    target_matrix = np.stack([target_raw_meta[name] for name in partitions.validation])
    normalized_targets, _, _ = normalize_meta_features(train_matrix, target_matrix)
    target_meta = {
        name: row for name, row in zip(partitions.validation, normalized_targets)
    }
    target_configs = {
        name: decode_config(
            best_genome, target_meta[name], shape,
            max_warmup=hp_args.max_warmup,
        )
        for name in partitions.validation
    }
    _atomic_json(args.output_dir / "final_recommendations.json", {
        "selection_protocol": "source_only_best_fitness",
        "source_candidates": int(solution_step),
        "source_best_train_fitness": float(best_train_fitness),
        "target_scores_observed": False,
        "policy": best_policy_path.name,
        "policy_digest": genome_digest(best_genome),
        "target_configs": target_configs,
    })
    if wandb_run is not None:
        wandb_run.summary["source_best_train_fitness"] = float(best_train_fitness)
        wandb_run.summary["source_candidates"] = int(solution_step)

    if wandb_run is not None:
        wandb_run.finish()
    return 0


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