"""Training pairs and evaluation inputs of a dataset package (``trajectories_fine.npz``, ``split.json``,
``normalization.npz``).

* ``get_dataloaders(cfg, seed)`` -> ``(train_loader, val_loader)``: segment pairs ``(x_tk, x_tk1, t_k, t_k1)`` for
  Cascaded SI (``method.data_mode: fine_segment_pairs``) or endpoint pairs ``(x_0, x_T)`` for Direct SI and EDM
  (``fine_endpoint_pairs``), from the pair manifest named by ``method.fine_slice_manifest``.
* ``get_split_set(cfg, split, seed, x0_bank)`` -> ``(x0_per_gid, truth_clouds, gids)`` at the forecast horizon
  ``method.fine_horizon_idx``, the inputs of sampling and of validation.
"""
from __future__ import annotations

from typing import Tuple

import torch
from torch.utils.data import DataLoader

from .fine import (EndpointPairDataset, SegmentPairDataset, ValEndpointPairDataset,  # noqa: F401
                   ValSegmentPairDataset, get_fine_split_set)

DATA_MODES = ("fine_segment_pairs", "fine_endpoint_pairs")


def get_dataloaders(cfg: dict, seed: int = 0) -> Tuple[DataLoader, DataLoader]:
    packaged_dir = cfg["data"]["packaged_dataset_path"]
    batch_size = int(cfg["training"]["batch_size"])
    method_cfg = cfg.get("method", {})
    data_mode = method_cfg.get("data_mode")
    num_workers = int(cfg["training"].get("num_workers", 0))
    if data_mode == "fine_segment_pairs":
        train_ds = SegmentPairDataset(packaged_dir, method_cfg["fine_slice_manifest"], seed=seed,
                                      expected_t_values=method_cfg.get("sampling", {}).get("t_values"))
        val_ds = ValSegmentPairDataset(packaged_dir, method_cfg["fine_slice_manifest"],
                                       n_members=int(method_cfg.get("fine_val_members", 20)))
    elif data_mode == "fine_endpoint_pairs":
        train_ds = EndpointPairDataset(packaged_dir, method_cfg["fine_slice_manifest"], seed=seed,
                                       expected_t_values=method_cfg.get("sampling", {}).get("t_values"))
        val_ds = ValEndpointPairDataset(packaged_dir, method_cfg["fine_slice_manifest"],
                                        n_members=int(method_cfg.get("fine_val_members", 20)))
    else:
        raise ValueError(f"method.data_mode must be one of {DATA_MODES}, got {data_mode!r}")
    g = torch.Generator()
    g.manual_seed(seed)
    train_loader = DataLoader(train_ds, batch_size=batch_size, shuffle=True, num_workers=num_workers, generator=g,
                              pin_memory=num_workers > 0)
    val_loader = DataLoader(val_ds, batch_size=batch_size, shuffle=False, num_workers=num_workers,
                            pin_memory=num_workers > 0)
    return train_loader, val_loader


def get_split_set(cfg: dict, split_name: str, seed: int = 0, x0_bank: str | None = None):
    return get_fine_split_set(cfg, split_name, seed=seed, x0_bank=x0_bank)