src/data_store.py

from __future__ import annotations

import importlib
from pathlib import Path
import warnings
import numpy as np
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,
    )


# ---------------------------------------------------------------------------
# Optimization (RMSE) datasets
# ---------------------------------------------------------------------------

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_OPT = ("algorithm_type", "rmse_target", "ensemble_size", "random_seed")
EXPECTED_DIMS_UQ  = ("random_seed", "ensemble_size", "k_iter", "coverage_quantile")


# ---------------------------------------------------------------------------
# Uncertainty Quantification datasets
# ---------------------------------------------------------------------------
# Add entries here when UQ result NetCDF files are available.  Each file must
# carry a ``metric`` variable and a ``uq_target`` coordinate (in place of the
# ``rmse_target`` used by the optimization files).  The rest of the schema
# (algorithm_type, ensemble_size, random_seed) is identical.
#
# Example:
#   UQ_DATASET_FILES = {
#       "L63": ["uq/l63_uq_results.nc"],
#       "L96": ["uq/l96_uq_results.nc"],
#   }

UQ_DATASET_FILES: dict[str, list[str]] = {}

# Canonical UQ target-scaling levels always offered in the UQ leaderboard selector.
UQ_TARGET_LEVELS: list[float] = [1.0, 1.25, 1.5, 1.75, 2.0, 2.25, 2.5]

# Budget-for-coverage source files: maps benchmark → list of (algorithm_type, nc_path) pairs.
# Each NC must carry output_coverage (coverage_quantile, k_iter, ensemble_size, random_seed),
# target_scaling, and output_dim.
UQ_BUDGET_FILES: dict[str, list[tuple[str, str]]] = {
    "L63":                [("ces-eki-dmc", "ces-eki-dmc_results/ces-eki-dmc_l63_ensemble_results_2026-06-15_minimal.nc")],
    "L96":                [("ces-eki-dmc", "ces-eki-dmc_results/ces-eki-dmc_l96_ensemble_results_2026-06-15_minimal.nc")],
    "L96_NN_FORCING":     [("ces-eki-dmc", "ces-eki-dmc_results/ces-eki-dmc_l96_nn_forcing_ensemble_results_2026-06-15_minimal.nc")],
    "L96_SPATIAL_FORCING": [("ces-eki-dmc", "ces-eki-dmc_results/ces-eki-dmc_l96_spatial_forcing_ensemble_results_2026-06-15_minimal.nc")],
}

# Quantile levels used for the all-quantiles-satisfied budget-for-coverage condition.
# Budget = N_ens · k_iter where first k s.t. |S(q)−q| ≤ c·√(q(1−q)/N_y) for ALL q below.
UQ_COVERAGE_QUANTILES: list[float] = [0.15, 0.5, 0.85]



# ---------------------------------------------------------------------------
# Private loader — shared by both public load_* functions
# ---------------------------------------------------------------------------

def _load_store(dataset_files: dict[str, list[str]], target_col: str, *, drop_nan_metric: bool = True) -> pd.DataFrame:
    """Load and merge NetCDF metric files into a long-format DataFrame.

    Parameters
    ----------
    dataset_files:
        Mapping of benchmark name → list of relative NetCDF paths under
        ``data/``.
    target_col:
        Name of the target-coordinate dimension in the NetCDF files, e.g.
        ``"rmse_target"`` or ``"uq_target"``.
    """
    if not dataset_files:
        return pd.DataFrame(
            columns=["benchmark", "algorithm_type", target_col, "ensemble_size", "metric", "failure_rate"]
        )

    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"]
                if not set(metric.dims).issubset(EXPECTED_DIMS_OPT):
                    warnings.warn(
                        f"{file_path.name}: metric dims {metric.dims} contain unexpected "
                        f"dimensions not in EXPECTED_DIMS_OPT {EXPECTED_DIMS_OPT} — skipping"
                    )
                    continue

                # 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 runs (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

                required_cols = ["benchmark", "algorithm_type", target_col, "ensemble_size", "metric", "failure_rate"]
                missing_cols = [c for c in required_cols 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[required_cols])

    if not records:
        return pd.DataFrame(
            columns=["benchmark", "algorithm_type", target_col, "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)
    if drop_nan_metric:
        merged = merged.dropna(subset=["ensemble_size", "metric"])
    else:
        merged = merged.dropna(subset=["ensemble_size"])
    merged["ensemble_size"] = merged["ensemble_size"].astype(int)
    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


# ---------------------------------------------------------------------------
# Public cached loaders
# ---------------------------------------------------------------------------

@st.cache_data(show_spinner=False)
def load_metric_store() -> pd.DataFrame:
    """Load the optimization (RMSE) metric store from ``DATASET_FILES``."""
    return _load_store(DATASET_FILES, "rmse_target")


@st.cache_data(show_spinner=False)
def load_uq_store() -> pd.DataFrame:
    """Load the UQ metric store derived from budget-for-coverage computations.

    Returns a DataFrame compatible with ``render_leaderboard`` where ``metric``
    is the mean budget (N_ens·k_iter) to satisfy the coverage condition across
    all three quantiles in ``UQ_COVERAGE_QUANTILES``.  ``uq_target`` is the
    target-scaling level c.  Failure means the coverage target was never
    reached within the k_iter range.
    """
    budget_df = load_uq_budget_store()
    if budget_df.empty:
        return pd.DataFrame(
            columns=[
                "benchmark", "algorithm_type", "algorithm_alias", "abbreviation",
                "Method", "family", "uq_target", "ensemble_size", "metric", "failure_rate",
            ]
        )
    df = budget_df.copy()
    df["metric"] = df["mean_budget"]
    df["algorithm_alias"] = df["algorithm_type"].map(normalize_method_name)
    df["Method"] = df["algorithm_type"].map(
        lambda m: get_method_meta(canonicalize_method_name(normalize_method_name(m))).get(
            "Method", m.upper()
        )
    )
    return df[
        [
            "benchmark", "algorithm_type", "algorithm_alias", "abbreviation",
            "Method", "family", "uq_target", "ensemble_size", "metric", "failure_rate",
        ]
    ]


@st.cache_data(show_spinner=False)
def load_uq_budget_store() -> pd.DataFrame:
    """Compute budget-for-coverage and iterations-for-coverage from ensemble result NC files.

    For each (target_scaling c, ensemble_size N, random_seed):
      Find the minimum k_iter s.t. |S(q) − q| ≤ c·√(q(1−q)/N_y) for ALL q in
      UQ_COVERAGE_QUANTILES.  Budget = N · k_iter; NaN when target never reached.

    Returns a DataFrame with columns:
        benchmark, algorithm_type, abbreviation, family, uq_target,
        ensemble_size, mean_budget, mean_iters, failure_count, failure_rate, n_seeds
    """
    xr = importlib.import_module("xarray")

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

    bq_vals = np.array(UQ_COVERAGE_QUANTILES)
    records: list[dict] = []

    for benchmark_name, method_files in UQ_BUDGET_FILES.items():
        for algorithm_type, filename in method_files:
            file_path = data_dir / filename
            if not file_path.exists():
                warnings.warn(f"UQ budget file {file_path} does not exist")
                continue

            with xr.open_dataset(file_path) as ds:
                if "output_coverage" not in ds:
                    warnings.warn(f"{file_path.name}: missing output_coverage")
                    continue

                if set(ds["output_coverage"].dims) != set(EXPECTED_DIMS_UQ):
                    warnings.warn(
                        f"{file_path.name}: output_coverage dims {ds['output_coverage'].dims} "
                        f"do not match EXPECTED_DIMS_UQ {EXPECTED_DIMS_UQ} — skipping"
                    )
                    continue

                cov_q = ds["coverage_quantile"].values  # (n_cov_q,)
                k_vals = ds["k_iter"].values            # (n_k,)  1-indexed
                ens_vals = ds["ensemble_size"].values   # (n_ens,)
                ts_vals = ds["target_scaling"].values   # (n_ts,)
                # xarray drops dimensions that no variable uses; read output_dim
                # directly from the underlying netCDF4 file to handle minimal files.
                nc4 = importlib.import_module("netCDF4")
                with nc4.Dataset(file_path) as _nc:
                    if "output_dim" not in _nc.dimensions:
                        warnings.warn(f"{file_path.name}: missing output_dim dimension")
                        continue
                    n_y = len(_nc.dimensions["output_dim"])
                # dims: (coverage_quantile, k_iter, ensemble_size, random_seed)
                cov_np = ds["output_coverage"].values

                n_rng = cov_np.shape[3]

                # Indices into coverage_quantile for [0.15, 0.5, 0.85]
                bq_idx = [int(np.argmin(np.abs(cov_q - q))) for q in bq_vals]

                for si, c in enumerate(ts_vals):
                    tol = c * np.sqrt(bq_vals * (1.0 - bq_vals) / n_y)

                    for ei, ens_f in enumerate(ens_vals):
                        N_ens = int(ens_f)
                        budgets = np.full(n_rng, np.nan)
                        kiters = np.full(n_rng, np.nan)

                        for ri in range(n_rng):
                            for ki, kv in enumerate(k_vals):
                                s_q = cov_np[bq_idx, ki, ei, ri]
                                if np.any(np.isnan(s_q)):
                                    continue
                                if np.all(np.abs(s_q - bq_vals) <= tol):
                                    budgets[ri] = N_ens * float(kv)
                                    kiters[ri] = float(kv)
                                    break

                        valid_b = budgets[~np.isnan(budgets)]
                        valid_k = kiters[~np.isnan(kiters)]
                        n_fail = n_rng - len(valid_b)
                        records.append(
                            {
                                "benchmark": benchmark_name,
                                "algorithm_type": algorithm_type,
                                "uq_target": float(c),
                                "ensemble_size": N_ens,
                                "mean_budget": float(np.mean(valid_b)) if len(valid_b) > 0 else np.nan,
                                "mean_iters": float(np.mean(valid_k)) if len(valid_k) > 0 else np.nan,
                                "failure_count": n_fail,
                                "failure_rate": 100.0 * n_fail / n_rng,
                                "n_seeds": n_rng,
                            }
                        )

    if not records:
        return pd.DataFrame(
            columns=[
                "benchmark", "algorithm_type", "abbreviation", "family",
                "uq_target", "ensemble_size", "mean_budget", "mean_iters",
                "failure_count", "failure_rate", "n_seeds",
            ]
        )

    df = pd.DataFrame(records)

    def _meta(m: str, key: str, fallback: str) -> str:
        return get_method_meta(canonicalize_method_name(normalize_method_name(m))).get(key, fallback)

    df["abbreviation"] = df["algorithm_type"].map(lambda m: _meta(m, "abbreviation", m.upper()))
    df["family"] = df["algorithm_type"].map(lambda m: _meta(m, "family", "Kalman"))
    df["Method"] = df["algorithm_type"].map(lambda m: _meta(m, "Method", m.upper()))
    return df