src/common/method_registry.py

from __future__ import annotations

import json
from pathlib import Path


KNOWN_METHODS = {
    "teki": {
        "abbreviation": "TEKI",
        "Method": "Tikhonov Regularized Ensemble Kalman Inversion",
        "family": "Kalman",
        "aliases": ["teki"],
    },
    "etki": {
        "abbreviation": "ETKI",
        "Method": "Ensemble Transform Kalman Inversion",
        "family": "Kalman",
        "aliases": ["etki"],
    },
    "iekf": {
        "abbreviation": "IEKF",
        "Method": "Iterative Ensemble Kalman Filter",
        "family": "Kalman",
        "aliases": ["iekf", "gnsl", "gnki"],
    },
    "uki": {
        "abbreviation": "UKI",
        "Method": "Unscented Kalman Inversion",
        "family": "Kalman",
        "aliases": ["uki"],
    },
    "abc": {
        "abbreviation": "ABC",
        "Method": "Approximate Bayesian Calibration",
        "family": "Bayesian",
        "aliases": ["abc"],
    },
    "hm": {
        "abbreviation": "HM",
        "Method": "History Matching",
        "family": "Bayesian",
        "aliases": ["hm"],
    },
    "ces-eki-dmc": {
        "abbreviation": "CES-EKI-DMC",
        "Method": "Calibrate Emulate Sample (EKI-DataMisfitController)",
        "family": "calibrate_then_emulate",
        "aliases": ["ces-eki-dmc"]
    },
    "adam": {
        "abbreviation": "ADAM",
        "Method": "Adaptive Moment Estimation",
        "family": "gradient",
        "aliases": ["adam"],
    },
    "lm": {
        "abbreviation": "LM",
        "Method": "Levenberg-Marquardt",
        "family": "gradient",
        "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")