"""Authoritative 20-model registry used by notebooks, training, and reporting."""

from __future__ import annotations

from dataclasses import asdict, dataclass


@dataclass(frozen=True)
class ModelSpec:
    model_id: str
    display_name: str
    family: str
    input_kind: str
    framework: str


MODEL_CATALOG: tuple[ModelSpec, ...] = (
    ModelSpec("extra_trees", "Extra Trees", "classical", "tabular", "scikit-learn"),
    ModelSpec("gradient_boosting", "Gradient Boosting", "classical", "tabular", "scikit-learn"),
    ModelSpec("random_forest", "Random Forest", "classical", "tabular", "scikit-learn"),
    ModelSpec("xgboost", "XGBoost", "classical", "tabular", "xgboost"),
    ModelSpec("lightgbm", "LightGBM", "classical", "tabular", "lightgbm"),
    ModelSpec("svr", "Support Vector Regression", "classical", "tabular", "scikit-learn"),
    ModelSpec("ridge", "Ridge Regression", "classical", "tabular", "scikit-learn"),
    ModelSpec("knn", "K-Nearest Neighbours", "classical", "tabular", "scikit-learn"),
    ModelSpec("vanilla_lstm", "Vanilla LSTM", "recurrent", "sequence", "pytorch"),
    ModelSpec("bidirectional_lstm", "Bidirectional LSTM", "recurrent", "sequence", "pytorch"),
    ModelSpec("gru", "GRU", "recurrent", "sequence", "pytorch"),
    ModelSpec("attention_lstm", "Attention LSTM", "recurrent", "sequence", "pytorch"),
    ModelSpec("battery_gpt", "BatteryGPT", "transformer", "sequence", "pytorch"),
    ModelSpec("temporal_fusion_transformer", "Temporal Fusion Transformer", "transformer", "sequence", "pytorch"),
    ModelSpec("itransformer", "iTransformer", "transformer", "sequence", "tensorflow"),
    ModelSpec("physics_itransformer", "Physics-Informed iTransformer", "transformer", "sequence", "tensorflow"),
    ModelSpec("dynamic_graph_itransformer", "Dynamic-Graph iTransformer", "graph", "sequence", "tensorflow"),
    ModelSpec("vae_lstm", "VAE-LSTM", "variational", "sequence", "pytorch"),
    ModelSpec("stacking_ensemble", "Stacking Ensemble", "ensemble", "tabular", "scikit-learn"),
    ModelSpec("weighted_ensemble", "Weighted Average Ensemble", "ensemble", "tabular", "scipy"),
)


def model_catalog_records() -> list[dict[str, str]]:
    return [asdict(spec) for spec in MODEL_CATALOG]


def get_model_spec(model_id: str) -> ModelSpec:
    for spec in MODEL_CATALOG:
        if spec.model_id == model_id:
            return spec
    raise KeyError(f"Unknown model_id: {model_id}")