"""Model and method factories shared by training, sampling and evaluation.

Methods: ``direct_si``, ``cascaded_si`` and ``edm``. Models: ``unet2_5ch`` (Euler, SQG), ``unet2_5ch_coord``
(passive scalar), ``unet2_5ch_coord_aug`` (heterogeneous waves) and ``unet2_2comp`` (NLKG); each run configuration
under records/<system>/runs/ names its model and method."""

from __future__ import annotations

from dataclasses import dataclass
from typing import Callable

from .methods.edm import EDMMethod
from .methods.direct_si import DirectSIMethod
from .methods.cascaded_si import CascadedSIMethod
from .models.unet2 import Conditional2D, LocalDrift2D

SEGMENT_DATA_MODES = ("segment_pairs", "fine_segment_pairs")

# reserved transient key: the method's segment boundaries in physical time, injected by
# build_model for backbones that declare physical times in their own YAML (see module docstring)
T_VALUES_KEY = "_method_t_values"


@dataclass(frozen=True)
class ModelSpec:
    name: str
    endpoint_cls: type
    segment_cls: type
    needs_dataset_norm: bool = False
    description: str = ""


MODELS: dict[str, ModelSpec] = {}
METHODS: dict[str, Callable[[dict], object]] = {}


def register_model(name: str, endpoint_cls: type, segment_cls: type,
                   needs_dataset_norm: bool = False, description: str = "",
                   overwrite: bool = False) -> ModelSpec:
    if name in MODELS and not overwrite:
        raise KeyError(f"model {name!r} already registered")
    spec = ModelSpec(name, endpoint_cls, segment_cls, needs_dataset_norm, description)
    MODELS[name] = spec
    return spec


def register_method(name: str, factory: Callable[[dict], object],
                    overwrite: bool = False) -> None:
    if name in METHODS and not overwrite:
        raise KeyError(f"method {name!r} already registered")
    METHODS[name] = factory


def head_for_data_mode(data_mode: str | None) -> str:
    return "segment" if data_mode in SEGMENT_DATA_MODES else "endpoint"


def build_model(model_cfg: dict, method_cfg: dict | None = None,
                data_cfg: dict | None = None):
    """Construct the backbone named by ``model_cfg['name']``.

    ``method_cfg['data_mode']`` selects the head. For models with ``needs_dataset_norm`` the dataset path is
    taken from ``data_cfg`` (``packaged_dataset_path``); a model configuration that carries its own copy of the
    key is accepted only if it agrees with the data config.
    """
    name = model_cfg["name"]
    if name not in MODELS:
        raise ValueError(f"unknown model: {name!r} (registered: {sorted(MODELS)})")
    spec = MODELS[name]
    cfg = dict(model_cfg)
    if spec.needs_dataset_norm:
        data_path = (data_cfg or {}).get("packaged_dataset_path")
        model_path = cfg.get("packaged_dataset_path")
        if data_path is None and model_path is None:
            raise KeyError(
                f"model {name!r} needs the dataset normalization: pass data_cfg with "
                "'packaged_dataset_path' (model configurations do not carry dataset paths)")
        if data_path is not None and model_path is not None:
            if _norm_path(data_path) != _norm_path(model_path):
                raise ValueError(
                    f"model.packaged_dataset_path ({model_path}) != data.packaged_dataset_path "
                    f"({data_path}) — the FFT wrapper must use the training dataset's statistics")
        cfg["packaged_dataset_path"] = str(data_path if data_path is not None else model_path)
    t_values = ((method_cfg or {}).get("sampling") or {}).get("t_values")
    if t_values:
        cfg[T_VALUES_KEY] = list(t_values)
    data_mode = (method_cfg or {}).get("data_mode")
    # a model configuration that declares the horizon planes must agree with the method's time grid:
    # t_start = t_values[0], t_end = t_values[-1]
    tv = ((method_cfg or {}).get("sampling") or {}).get("t_values")
    if tv is not None and ("t_start" in cfg or "t_end" in cfg):
        t0, t1 = float(tv[0]), float(tv[-1])
        if "t_start" in cfg and abs(float(cfg["t_start"]) - t0) > 1e-9:
            raise ValueError(f"model {name!r}: t_start {cfg['t_start']} != method t_values[0] = {t0}")
        if "t_end" in cfg and abs(float(cfg["t_end"]) - t1) > 1e-9:
            raise ValueError(f"model {name!r}: t_end {cfg['t_end']} != method horizon t_values[-1] = {t1} "
                             f"(the model constant must equal the forecast horizon)")
    cls = spec.segment_cls if head_for_data_mode(data_mode) == "segment" else spec.endpoint_cls
    return cls.from_config(cfg)


def build_method(method_cfg: dict):
    name = method_cfg["name"]
    if name not in METHODS:
        raise ValueError(f"unknown method: {name!r} (registered: {sorted(METHODS)})")
    return METHODS[name](method_cfg)


def _norm_path(p: str) -> str:
    import os
    return os.path.normpath(os.path.abspath(str(p)))


# ── built-ins ────────────────────────────────────────────────────────────────
register_model("unet2", Conditional2D, LocalDrift2D, needs_dataset_norm=True,
               description="FFT-wrapped 2-D circular-conv U-Net on the 32x32 coarse field "
                           "(Conditional2D 3-ch / LocalDrift2D 5-ch)")
register_method("edm", EDMMethod.from_config)
register_method("direct_si", DirectSIMethod.from_config)
register_method("cascaded_si", CascadedSIMethod.from_config)

# the backbones register themselves on import; unet2_coord provides the coordinate planes of unet2_5ch_coord(_aug)
from .models import unet2_coord as _unet2_coord  # noqa: F401,E402
from .models import unet2_5ch as _unet2_5ch  # noqa: F401,E402
from .models import unet2_5ch_coord as _unet2_5ch_coord  # noqa: F401,E402
from .models import unet2_5ch_coord_aug as _unet2_5ch_coord_aug  # noqa: F401,E402
from .models import unet2_2comp as _unet2_2comp  # noqa: F401,E402