src/data_store.py

from __future__ import annotations

import importlib
from pathlib import Path
import warnings 
import pandas as pd
import streamlit as st

try:
    from common.method_registry import (
        dump_method_registry_snapshot,
        canonicalize_method_name,
        get_method_meta,
        normalize_method_name,
    )
except ModuleNotFoundError:
    from src.common.method_registry import (
        dump_method_registry_snapshot,
        canonicalize_method_name,
        get_method_meta,
        normalize_method_name,
    )


DATASET_FILES = {
    "L63": [
        "l63_ensemble_results.nc",
        "UKI_results/uki_l63_ensemble_results.nc",
        "bayesian/l63_abc.nc",
        "bayesian/l63_hm.nc"
    ],
    "L96": [
        "l96_ensemble_results.nc",
        "UKI_results/uki_l96_ensemble_results.nc",
        "bayesian/l96_abc.nc"
    ],
    "L96_NN_FORCING": [
        "l96_nn_forcing_ensemble_results.nc",
        "UKI_results/uki_l96_nn_forcing_ensemble_results.nc"
    ],
    "L96_SPATIAL_FORCING": [
        "l96_spatial_forcing_ensemble_results.nc",
        "UKI_results/uki_l96_spatial_forcing_ensemble_results.nc",
        "bayesian/l96_varying_abc.nc"
    ],
}



EXPECTED_DIMS = ("algorithm_type", "rmse_target", "ensemble_size", "random_seed")


@st.cache_data(show_spinner=False)
def load_metric_store() -> pd.DataFrame:
    xr = importlib.import_module("xarray")

    project_root = Path(__file__).resolve().parent.parent
    data_dir = project_root / "data"

    records: list[pd.DataFrame] = []

    for benchmark_name, filenames in DATASET_FILES.items():
        for filename in filenames:
            file_path = data_dir / filename
            if not file_path.exists():
                warnings.warn(f"Dataset {file_path} does not exist")
                continue

            with xr.open_dataset(file_path) as dataset:
                if "metric" not in dataset:
                    warnings.warn(f"Dataset {file_path} does not contain 'metric' variable")
                    continue

                metric = dataset["metric"]
                
                # Track failures (metric == -1) as a percentage
                failures_count = (metric == -1).sum(dim="random_seed")

                failure_rate = (failures_count / dataset.sizes["random_seed"]) * 100
                
                # Filter out failed calibrations (metric == -1)
                metric = metric.where(metric != -1)
                
                metric_mean = metric.mean(dim="random_seed", skipna=True)
                df = metric_mean.to_dataframe(name="metric").reset_index()
                
                df_failures = failure_rate.to_dataframe(name="failure_rate").reset_index()
                df["failure_rate"] = df_failures["failure_rate"]

                df["benchmark"] = benchmark_name
                
                if "algorithm_type" not in df.columns:
                    if "abc" in file_path.name:
                        df["algorithm_type"] = "abc"
                    elif "hm" in file_path.name:
                        df["algorithm_type"] = "hm"
                        
                if "ensemble_size" not in df.columns:
                    if "abc" in file_path.name:
                        df["ensemble_size"] = 1
                        
                missing_cols = [c for c in ["benchmark", "algorithm_type", "rmse_target", "ensemble_size", "metric", "failure_rate"] if c not in df.columns]
                if missing_cols:
                    st.warning(
                        f"Skipping {file_path.name}: metric dataframe is missing expected columns {missing_cols}"
                    )
                    continue

                records.append(df[["benchmark", "algorithm_type", "rmse_target", "ensemble_size", "metric", "failure_rate"]])

    if not records:
        return pd.DataFrame(columns=["benchmark", "algorithm_type", "rmse_target", "ensemble_size", "metric", "failure_rate"])

    merged = pd.concat(records, ignore_index=True)
    merged["algorithm_alias"] = merged["algorithm_type"].map(normalize_method_name)
    merged["algorithm_type"] = merged["algorithm_alias"].map(canonicalize_method_name)
    merged["ensemble_size"] = pd.to_numeric(merged["ensemble_size"], errors="coerce")
    merged["metric"] = pd.to_numeric(merged["metric"], errors="coerce")
    merged["failure_rate"] = pd.to_numeric(merged["failure_rate"], errors="coerce").fillna(0.0)
    merged = merged.dropna(subset=["ensemble_size", "metric"])
    merged["ensemble_size"] = merged["ensemble_size"].astype(int)
    merged["forward_model_runs"] = merged["metric"]

    merged["abbreviation"] = merged["algorithm_type"].map(lambda method: get_method_meta(method).get("abbreviation"))
    merged["Method"] = merged["algorithm_type"].map(lambda method: get_method_meta(method).get("Method"))
    merged["family"] = merged["algorithm_type"].map(lambda method: get_method_meta(method).get("family"))
    merged["abbreviation"] = merged["abbreviation"].fillna(merged["algorithm_type"].str.upper())
    merged["Method"] = merged["Method"].fillna(merged["abbreviation"])
    merged["family"] = merged["family"].fillna("Kalman")

    try:
        dump_method_registry_snapshot(project_root=project_root, observed_methods=set(merged["algorithm_type"].unique()))
    except OSError:
        st.warning("Unable to write method registry snapshot to .cache/known_methods_snapshot.json")

    return merged


def build_leaderboard(metric_store: pd.DataFrame) -> pd.DataFrame:
    if metric_store.empty:
        return pd.DataFrame(
            columns=[
                "family",
                "Method",
                "abbreviation",
                "benchmark",
                "Forward Model Runs",
                "Optimal Ensemble Size",
                "Target Level",
            ]
        )

    best_idx = metric_store.groupby(["benchmark", "algorithm_type"])["metric"].idxmin()
    best = metric_store.loc[best_idx].copy()
    best = best.rename(columns={"metric": "Forward Model Runs", "rmse_target": "Target Level", "ensemble_size": "Optimal Ensemble Size"})
    return best[["family", "Method", "abbreviation", "benchmark", "Forward Model Runs", "Optimal Ensemble Size", "Target Level"]].sort_values(
        ["benchmark", "Forward Model Runs"], ascending=[True, True]
    )