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"]
    }
}


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")