src/common/method_registry.py

from __future__ import annotations

import json
from pathlib import Path


# Taxonomy
# --------
# Every method carries four independent tags, each scoped to its own axis:
#
# ``parallelism`` — how the search explores parameter space:
#   "serial"               — a single point estimate advanced step by step (ADAM, LM).
#   "parallel-independent" — a population of candidates updated with no coupling
#                             between members (ABC's accepted samples, HM's per-wave
#                             resampling from the non-implausible region).
#   "parallel-interacting" — an ensemble whose members are coupled through a shared
#                             update each iteration (any Kalman-based method).
#
# ``update_type`` — the mechanism driving each update step:
#   "gradient" — follows the loss gradient or a Gauss-Newton approximation of it.
#   "kalman"   — a (possibly linearized/unscented) Kalman-style ensemble update.
#   "general"  — anything else (e.g. ABC's rejection sampling, HM's implausibility cuts).
#
# ``method_goal`` — what the method is built to report:
#   "optimization" — a single best-fit parameter estimate.
#   "uq"           — the full posterior / parameter uncertainty. UQ methods can still
#                    be scored on the Optimization leaderboard (usually less
#                    competitive there, since they're not optimizing for speed-to-target).
#
# ``emulator_use`` — when/whether a surrogate model of the forward model is used:
#   "none"            — samples/evaluates the true forward model throughout.
#   "within-optimize" — refits a surrogate at each iteration of the search (HM waves).
#   "after-optimize"  — fits a surrogate once, after calibration finishes (CES).
#
# Note: Kalman methods are Bayesian in spirit too (approximate Gaussian posterior
# updates) — ``update_type`` is about mechanism, not a "Bayesian vs not" philosophy.
KNOWN_METHODS = {
    "teki": {
        "abbreviation": "TEKI",
        "Method": "Tikhonov Regularized Ensemble Kalman Inversion",
        "parallelism": "parallel-interacting",
        "update_type": "kalman",
        "method_goal": "optimization",
        "emulator_use": "none",
        "aliases": ["teki"],
    },
    "etki": {
        "abbreviation": "ETKI",
        "Method": "Ensemble Transform Kalman Inversion",
        "parallelism": "parallel-interacting",
        "update_type": "kalman",
        "method_goal": "optimization",
        "emulator_use": "none",
        "aliases": ["etki"],
    },
    "iekf": {
        "abbreviation": "IEKF",
        "Method": "Iterative Ensemble Kalman Filter",
        "parallelism": "parallel-interacting",
        "update_type": "kalman",
        "method_goal": "uq",
        "emulator_use": "none",
        "aliases": ["iekf", "gnsl", "gnki"],
    },
    "uki": {
        "abbreviation": "UKI",
        "Method": "Unscented Kalman Inversion",
        "parallelism": "parallel-interacting",
        "update_type": "kalman",
        "method_goal": "optimization",
        "emulator_use": "none",
        "aliases": ["uki"],
    },
    "abc": {
        "abbreviation": "ABC",
        "Method": "Approximate Bayesian Calibration",
        "parallelism": "parallel-independent",
        "update_type": "general",
        "method_goal": "uq",
        "emulator_use": "none",
        "aliases": ["abc"],
    },
    "hm": {
        "abbreviation": "HM",
        "Method": "History Matching",
        "parallelism": "parallel-independent",
        "update_type": "general",
        "method_goal": "uq",
        "emulator_use": "within-optimize",
        "aliases": ["hm"],
    },
    "ces-eki-dmc": {
        "abbreviation": "CES-EKI-DMC",
        "Method": "Calibrate Emulate Sample (EKI-DataMisfitController)",
        "parallelism": "parallel-interacting",
        "update_type": "kalman",
        "method_goal": "uq",
        "emulator_use": "after-optimize",
        "aliases": ["ces-eki-dmc"]
    },
    "ces-eki-const": {
        "abbreviation": "CES-EKI-CONST",
        "Method": "Calibrate Emulate Sample (EKI-Constant Scheduler)",
        "parallelism": "parallel-interacting",
        "update_type": "kalman",
        "method_goal": "uq",
        "emulator_use": "after-optimize",
        "aliases": ["ces-eki-const"]
    },
    "ces-iekf-const": {
        "abbreviation": "CES-IEKF-CONST",
        "Method": "Calibrate Emulate Sample (IEKF-Constant Scheduler)",
        "parallelism": "parallel-interacting",
        "update_type": "kalman",
        "method_goal": "uq",
        "emulator_use": "after-optimize",
        "aliases": ["ces-iekf-const"]
    },
    "adam": {
        "abbreviation": "ADAM",
        "Method": "Adaptive Moment Estimation",
        "parallelism": "serial",
        "update_type": "gradient",
        "method_goal": "optimization",
        "emulator_use": "none",
        "aliases": ["adam"],
    },
    "lm": {
        "abbreviation": "LM",
        "Method": "Levenberg-Marquardt",
        "parallelism": "serial",
        "update_type": "gradient",
        "method_goal": "optimization",
        "emulator_use": "none",
        "aliases": ["lm", "levenberg_marquardt", "levenberg-marquardt", "gradient_descent"],
    },
}


# Vega tableau10 palette — one slot per method in KNOWN_METHODS declaration order.
# New methods appended to KNOWN_METHODS get the next slot; existing colors never shift.
_METHOD_PALETTE = [
    "#4c78a8", "#f58518", "#e45756", "#72b7b2", "#54a24b",
    "#eeca3b", "#b279a2", "#ff9da6", "#9d755d", "#bab0ac",
]

# Stable abbreviation → hex color mapping.  Import this wherever Altair charts are built
# so every plot in the app assigns the same color to each method.
METHOD_COLORS: dict[str, str] = {
    meta["abbreviation"]: _METHOD_PALETTE[i % len(_METHOD_PALETTE)]
    for i, meta in enumerate(KNOWN_METHODS.values())
}


def build_alias_lookup() -> dict[str, str]:
    lookup: dict[str, str] = {}
    for canonical_name, meta in KNOWN_METHODS.items():
        lookup[canonical_name] = canonical_name
        lookup[canonical_name.upper()] = canonical_name
        for alias in meta.get("aliases", []):
            lookup[alias.lower()] = canonical_name
            lookup[alias.upper()] = canonical_name
    return lookup


ALIAS_TO_CANONICAL = build_alias_lookup()


def normalize_method_name(name: object) -> str:
    text = str(name).strip()

    if text.startswith("b'") and text.endswith("'"):
        text = text[2:-1]
    elif text.startswith('b"') and text.endswith('"'):
        text = text[2:-1]

    text = text.strip("\"'").strip()
    return text.lower()


def canonicalize_method_name(name: object) -> str:
    normalized = normalize_method_name(name)
    return ALIAS_TO_CANONICAL.get(normalized, normalized)


def get_method_meta(canonical_name: str) -> dict[str, str]:
    return KNOWN_METHODS.get(canonical_name, {})


def dump_method_registry_snapshot(project_root: Path, observed_methods: set[str]) -> None:
    snapshot = {
        "known_methods": KNOWN_METHODS,
        "observed_methods": sorted(observed_methods),
        "unmapped_observed_methods": sorted([method for method in observed_methods if method not in KNOWN_METHODS]),
    }
    cache_dir = project_root / ".cache"
    cache_dir.mkdir(parents=True, exist_ok=True)
    target_file = cache_dir / "known_methods_snapshot.json"
    target_file.write_text(json.dumps(snapshot, indent=2), encoding="utf-8")