"""Configuration files of the repository: the dataset cards, the evaluation settings and the recorded runs.

* ``configs/<system>/dataset.yaml`` names the data of one system; ``${VAR}`` entries are read from the environment.
* ``configs/evaluation.yaml`` holds the csFID settings and the sampling seeds shared by all systems.
* ``records/<system>/units.json`` lists the models of the paper and their recorded scores.
"""
from __future__ import annotations

import json
import os
import re
from pathlib import Path

import yaml

REPO = Path(__file__).resolve().parents[1]
CONFIGS = REPO / "configs"
RECORDS = REPO / "records"
SYSTEMS = ("euler", "sqg", "waves", "nlkg", "passive_scalar")

_ENV_RE = re.compile(r"\$\{([A-Za-z_][A-Za-z0-9_]*)(?::-([^}]*))?\}")


def expand_env(value):
    """``${VAR}`` / ``${VAR:-default}`` expansion in strings (recursively in dicts and lists)."""
    if isinstance(value, str):
        def sub(m):
            var, default = m.group(1), m.group(2)
            v = os.environ.get(var)
            if v is None or v == "":
                if default is None:
                    raise KeyError(f"environment variable {var} is not set (used in a config)")
                return default
            return v
        return _ENV_RE.sub(sub, value)
    if isinstance(value, dict):
        return {k: expand_env(v) for k, v in value.items()}
    if isinstance(value, list):
        return [expand_env(v) for v in value]
    return value


def load_yaml(path: str | os.PathLike) -> dict:
    with open(path) as f:
        return yaml.safe_load(f) or {}


def dataset_card(system: str, expand: bool = False) -> dict:
    """The dataset card of ``system``; with ``expand`` the ``${VAR}`` paths are resolved from the environment."""
    if system not in SYSTEMS:
        raise KeyError(f"unknown system {system!r} (one of {SYSTEMS})")
    card = load_yaml(CONFIGS / system / "dataset.yaml")
    return expand_env(card) if expand else card


def data_paths(system: str, data: str | None = None, truth_bank: str | None = None) -> tuple[Path, Path]:
    """The package and truth-bank directories: the given values, else the environment variables of the card."""
    card = dataset_card(system)
    if data is None:
        data = expand_env(card["data"])
    if truth_bank is None:
        truth_bank = expand_env(card["truth_bank"])
    return Path(data).resolve(), Path(truth_bank).resolve()


def evaluation_settings() -> dict:
    return load_yaml(CONFIGS / "evaluation.yaml")


def records(system: str) -> dict:
    return json.loads((RECORDS / system / "units.json").read_text())


def unit(system: str, method: str, S: int | None = None) -> tuple[dict, dict]:
    """The recorded model of ``system`` for a paper method (``direct_si``, ``edm``, ``cascaded_si``); with ``S`` the
    Cascaded SI model of that depth (Table 6). Returns (records, unit)."""
    rec = records(system)
    if method == "edm":
        if S is not None:
            raise SystemExit("--S applies to cascaded_si only")
        key = "edm"
    elif method == "direct_si":
        if S not in (None, 1):
            raise SystemExit("direct_si is S = 1; use --method cascaded_si --S <S> for other depths")
        key = rec["table2"]["direct_si"]
    elif method == "cascaded_si":
        key = rec["table2"]["cascaded_si"] if S is None else f"cascaded_si_S{int(S):02d}"
    else:
        raise SystemExit(f"unknown method {method!r}")
    u = next((u for u in rec["units"] if u["key"] == key), None)
    if u is None:
        raise SystemExit(f"{system}: no recorded model {key} (depths: {[u['S'] for u in rec['units'] if u['key'] != 'edm']})")
    return rec, u


def run_config(system: str, u: dict, data: Path) -> dict:
    """The recorded run configuration of a model with the dataset path filled in."""
    cfg = load_yaml(REPO / u["config"])
    cfg["data"]["packaged_dataset_path"] = str(Path(data).resolve()) + "/"
    return cfg