"""Loaders of a dataset package: ``trajectories_fine.npz`` holds the blocks ``vecs_train`` (conditions, realizations,
snapshots, D), ``vecs_eval`` and optionally ``vecs_new``, keyed by ``gids_train`` / ``gids_eval`` / ``gids_new``, and
the snapshot times ``t_values``; ``split.json`` names the training, validation and test conditions;
``normalization.npz`` the per-coordinate mean and std.

  * training pairs are read from ``vecs_train`` only, restricted to the (condition, realization) rows of a pair
    manifest (``bifurcation.data.manifests``); segments are ``boundary_idx`` snapshot pairs and the conditioning
    times are the physical times of the stored grid;
  * validation and test inputs are ``vecs_eval`` at the snapshot ``fine_horizon_idx``;
  * x_0 of a condition is member 0 at t = 0 (all members share it; asserted at 1e-9).
"""

from __future__ import annotations

import hashlib
import json
import os
import zipfile
from typing import List, Tuple

import numpy as np
import torch
from torch.utils.data import Dataset

FINE_DT = 0.025
FINE_FILE = "trajectories_fine.npz"
ENDPOINT_SCHEMA = "endpoint_pairs_v1"      # the endpoint manifests of bifurcation.data.manifests


# --------------------------------------------------------------------------- #
# npz member access (uncompressed members -> memmap; else plain load)
# --------------------------------------------------------------------------- #

def _npz_member_memmap(npz_path: str, name: str) -> np.ndarray:
    """Memory-map one array stored in an UNCOMPRESSED ``.npz``.

    The packages store every member with ``compress_type == 0``, so the raw ``.npy`` bytes sit contiguously
    inside the zip and can be memmapped: fancy-indexing then reads only the pages of the selected trajectories.

    Falls back to ``np.load`` (full read) if the member is compressed.
    """
    member = name if name.endswith(".npy") else name + ".npy"
    with zipfile.ZipFile(npz_path) as zf:
        info = zf.getinfo(member)
        if info.compress_type != zipfile.ZIP_STORED:
            return np.load(npz_path)[name]
        with open(npz_path, "rb") as f:
            f.seek(info.header_offset)
            local = f.read(30)
            if local[:4] != b"PK\x03\x04":
                raise IOError(f"bad local zip header for {member} in {npz_path}")
            name_len  = int.from_bytes(local[26:28], "little")
            extra_len = int.from_bytes(local[28:30], "little")
            f.seek(info.header_offset + 30 + name_len + extra_len)
            version = np.lib.format.read_magic(f)
            np.lib.format._check_version(version)
            shape, fortran, dtype = np.lib.format._read_array_header(f, version)
            offset = f.tell()
    order = "F" if fortran else "C"
    return np.memmap(npz_path, dtype=dtype, mode="r", offset=offset,
                     shape=shape, order=order)


def _sha256_file(path: str) -> str:
    h = hashlib.sha256()
    with open(path, "rb") as f:
        for chunk in iter(lambda: f.read(1 << 20), b""):
            h.update(chunk)
    return h.hexdigest()


def _resolve_path(path: str) -> str:
    """Absolute paths pass through; relative paths resolve against the repository root."""
    if os.path.isabs(path):
        return path
    return os.path.join(os.path.dirname(os.path.dirname(os.path.dirname(os.path.abspath(__file__)))), path)


def _load_manifest(slice_manifest: str) -> dict:
    path = _resolve_path(slice_manifest)
    with open(path) as f:
        man = json.load(f)
    man["_manifest_path"] = path
    # Provenance guard: the manifest must certify the master ordering it was
    # cut from, and that file must still hash to the recorded value.
    master = os.path.join(os.path.dirname(path), "master_order.npz")
    if not os.path.isfile(master):
        raise FileNotFoundError(
            f"master_order.npz not found next to slice manifest {path} — "
            "slice provenance cannot be verified")
    got = _sha256_file(master)
    want = man["master_order_sha256"]
    if got != want:
        raise AssertionError(
            f"master_order.npz hash mismatch for {path}: {got} != {want}")
    return man


# --------------------------------------------------------------------------- #
# training dataset
# --------------------------------------------------------------------------- #

class SegmentPairDataset(Dataset):
    """``(x_tk, x_tk1, t_k, t_k1)`` pairs from ``vecs_train`` per a pair manifest.

    Checks: exactly ``S`` pairs per selected trajectory, exactly ``B == N_traj * S`` pairs in total,
    boundaries on stored snapshots, physical times consistent with the manifest, no validation or test gid.
    """

    def __init__(self, packaged_dataset_path: str, slice_manifest: str,
                 seed: int = 0,
                 expected_t_values: list[float] | None = None):
        man = _load_manifest(slice_manifest)
        fine_path = os.path.join(packaged_dataset_path, FINE_FILE)

        boundary_idx = np.asarray(man["boundary_idx"], dtype=np.int64)
        t_vals_man   = np.asarray(man["t_values"], dtype=np.float64)
        S       = int(man["S"])
        n_traj  = int(man["N_traj"])
        B       = int(man["B"])
        rows    = np.asarray(man["trajectories"], dtype=np.int64)  # (N,3)

        # --- manifest self-consistency ---
        if rows.shape != (n_traj, 3):
            raise AssertionError(
                f"manifest trajectories shape {rows.shape} != ({n_traj}, 3)")
        if n_traj * S != B:
            raise AssertionError(f"N_traj*S = {n_traj * S} != B = {B}")
        if len(boundary_idx) != S + 1 or boundary_idx[0] != 0:
            raise AssertionError(f"boundary_idx must be S+1 long, start at 0; "
                                 f"got {boundary_idx.tolist()}")
        if np.any(np.diff(boundary_idx) != boundary_idx[1] - boundary_idx[0]):
            raise AssertionError("boundary_idx must be uniformly strided")
        uniq = np.unique(rows[:, :2], axis=0)
        if uniq.shape[0] != n_traj:
            raise AssertionError("duplicate (group_row, member_col) in slice")

        gids_train = _npz_member_memmap(fine_path, "gids_train")
        t_values   = np.asarray(_npz_member_memmap(fine_path, "t_values"))
        if boundary_idx[-1] > t_values.shape[0] - 1:
            raise AssertionError("boundary_idx exceeds fine snapshot grid")
        if not np.allclose(t_values[boundary_idx], t_vals_man, atol=1e-12):
            raise AssertionError(
                f"manifest t_values {t_vals_man.tolist()} do not match the "
                f"stored grid at boundary_idx {boundary_idx.tolist()}")
        if expected_t_values is not None:
            exp = np.asarray([float(t) for t in expected_t_values])
            if len(exp) != len(t_vals_man) or not np.allclose(exp, t_vals_man,
                                                              atol=1e-9):
                raise AssertionError(
                    f"sampling.t_values {exp.tolist()} != slice-manifest "
                    f"t_values {t_vals_man.tolist()} — the sampler chain and "
                    "the training segments must describe the same physics")
        # gid column must be exactly gids_train[group_row] and never an
        # eval/new gid (train block only contains canonical train gids, but
        # assert the mapping so a corrupted manifest fails loudly).
        if not np.array_equal(np.asarray(gids_train)[rows[:, 0]], rows[:, 2]):
            raise AssertionError("manifest gid column inconsistent with "
                                 "gids_train[group_row]")
        with open(os.path.join(packaged_dataset_path, "split.json")) as f:
            split = json.load(f)
        forbidden = set(split["val"]) | set(split["test"])
        if forbidden & set(rows[:, 2].tolist()):
            raise AssertionError("slice manifest contains val/test gids")

        # --- gather (N_traj, S+1, D) raw states, normalize, build pairs ---
        vecs = _npz_member_memmap(fine_path, "vecs_train")   # (80,200,41,D)
        seg = np.asarray(vecs[rows[:, 0], rows[:, 1]])[:, boundary_idx, :]
        if not np.all(np.isfinite(seg)):
            raise AssertionError("non-finite values in selected trajectories")
        D = seg.shape[-1]

        norm = np.load(os.path.join(packaged_dataset_path, "normalization.npz"))
        mean = norm["mean"].astype(np.float32)
        std  = norm["std"].astype(np.float32)

        starts = seg[:, :-1, :].reshape(-1, D)               # (N*S, D)
        ends   = seg[:,  1:, :].reshape(-1, D)
        # float32 normalization arithmetic
        x_tk  = ((starts.astype(np.float32) - mean) / std).astype(np.float32)
        x_tk1 = ((ends.astype(np.float32)   - mean) / std).astype(np.float32)

        t_k_grid  = np.broadcast_to(t_vals_man[:-1], (n_traj, S)).reshape(-1)
        t_k1_grid = np.broadcast_to(t_vals_man[1:],  (n_traj, S)).reshape(-1)

        self.x_tk  = torch.from_numpy(x_tk)
        self.x_tk1 = torch.from_numpy(x_tk1)
        self.t_k   = torch.from_numpy(t_k_grid.astype(np.float32))
        self.t_k1  = torch.from_numpy(t_k1_grid.astype(np.float32))
        if self.x_tk.shape[0] != B:
            raise AssertionError(f"built {self.x_tk.shape[0]} pairs, expected {B}")

        # Provenance of each pair (gid, member, segment).
        self.pair_gid     = np.repeat(rows[:, 2], S)
        self.pair_member  = np.repeat(rows[:, 1], S)
        self.pair_segment = np.tile(np.arange(S), n_traj)
        self.n_segments = S
        self.manifest = man
        self.seed = int(seed)

    def __len__(self) -> int:
        return self.x_tk.shape[0]

    def __getitem__(self, idx: int):
        return self.x_tk[idx], self.x_tk1[idx], self.t_k[idx], self.t_k1[idx]


class ValSegmentPairDataset(Dataset):
    """Validation pairs of the training objective: the run's segmentation applied to the validation gids of
    ``vecs_eval``, first ``n_members`` members per gid (the ``val_loss`` column of the log)."""

    def __init__(self, packaged_dataset_path: str, slice_manifest: str,
                 n_members: int = 20):
        man = _load_manifest(slice_manifest)
        fine_path = os.path.join(packaged_dataset_path, FINE_FILE)
        boundary_idx = np.asarray(man["boundary_idx"], dtype=np.int64)
        t_vals_man   = np.asarray(man["t_values"], dtype=np.float64)
        S = int(man["S"])

        with open(os.path.join(packaged_dataset_path, "split.json")) as f:
            split = json.load(f)
        val_gids = sorted(split["val"])

        gids_eval = np.asarray(_npz_member_memmap(fine_path, "gids_eval"))
        vecs_eval = _npz_member_memmap(fine_path, "vecs_eval")
        rows = []
        for gid in val_gids:
            where = np.nonzero(gids_eval == gid)[0]
            if len(where) != 1:
                raise AssertionError(f"val gid {gid} not found once in gids_eval")
            rows.append(int(where[0]))
        n_members = int(n_members)
        if not (1 <= n_members <= vecs_eval.shape[1]):
            raise ValueError(f"n_members {n_members} out of range")

        seg = np.asarray(vecs_eval[rows][:, :n_members, boundary_idx, :])
        if not np.all(np.isfinite(seg)):
            raise AssertionError("non-finite values in vecs_eval slice")
        n_traj = seg.shape[0] * seg.shape[1]
        D = seg.shape[-1]
        seg = seg.reshape(n_traj, S + 1, D)

        norm = np.load(os.path.join(packaged_dataset_path, "normalization.npz"))
        mean = norm["mean"].astype(np.float32)
        std  = norm["std"].astype(np.float32)
        starts = seg[:, :-1, :].reshape(-1, D)
        ends   = seg[:,  1:, :].reshape(-1, D)
        self.x_tk  = torch.from_numpy(((starts.astype(np.float32) - mean) / std)
                                      .astype(np.float32))
        self.x_tk1 = torch.from_numpy(((ends.astype(np.float32) - mean) / std)
                                      .astype(np.float32))
        self.t_k   = torch.from_numpy(np.broadcast_to(t_vals_man[:-1],
                                      (n_traj, S)).reshape(-1).astype(np.float32))
        self.t_k1  = torch.from_numpy(np.broadcast_to(t_vals_man[1:],
                                      (n_traj, S)).reshape(-1).astype(np.float32))
        self.n_segments = S

    def __len__(self) -> int:
        return self.x_tk.shape[0]

    def __getitem__(self, idx: int):
        return self.x_tk[idx], self.x_tk1[idx], self.t_k[idx], self.t_k1[idx]


# --------------------------------------------------------------------------- #
# endpoint datasets
# --------------------------------------------------------------------------- #

class EndpointPairDataset(Dataset):
    """``(x_0, x_T)`` endpoint pairs from ``vecs_train`` (block 0) and/or ``vecs_new`` (block 1) per an endpoint
    manifest (rows ``[block, group_row, member_col, gid]``), what the Direct SI and EDM losses consume.
    Normalization: float64 arithmetic, then cast to float32 (as ``get_fine_split_set``).

    Checks: exactly ``B == N_traj`` pairs (S = 1), boundaries exactly ``[0, idx_T]`` on stored snapshots, the gid
    mapping against the package's own key arrays, every gid in the training split, no duplicate row.
    """

    _BLOCKS = {0: ("vecs_train", "gids_train"), 1: ("vecs_new", "gids_new")}

    def __init__(self, packaged_dataset_path: str, slice_manifest: str,
                 seed: int = 0,
                 expected_t_values: list[float] | None = None):
        man = _load_manifest(slice_manifest)
        if man.get("schema") != ENDPOINT_SCHEMA:
            raise AssertionError(
                f"expected schema {ENDPOINT_SCHEMA}, got "
                f"{man.get('schema')!r} in {slice_manifest}")
        fine_path = os.path.join(packaged_dataset_path, FINE_FILE)

        boundary_idx = np.asarray(man["boundary_idx"], dtype=np.int64)
        t_vals_man   = np.asarray(man["t_values"], dtype=np.float64)
        S      = int(man["S"])
        n_traj = int(man["N_traj"])
        B      = int(man["B"])
        rows   = np.asarray(man["trajectories"], dtype=np.int64)   # (N,4)

        # --- manifest self-consistency ---
        if S != 1 or n_traj != B:
            raise AssertionError(f"endpoint slice must have S=1, N_traj=B; "
                                 f"got S={S} N_traj={n_traj} B={B}")
        if rows.shape != (n_traj, 4):
            raise AssertionError(
                f"manifest trajectories shape {rows.shape} != ({n_traj}, 4)")
        if len(boundary_idx) != 2 or boundary_idx[0] != 0:
            raise AssertionError(
                f"endpoint boundary_idx must be [0, idx_T]; got "
                f"{boundary_idx.tolist()}")
        if not set(np.unique(rows[:, 0]).tolist()) <= set(self._BLOCKS):
            raise AssertionError(
                f"unknown block ids {np.unique(rows[:, 0]).tolist()}")
        if np.unique(rows[:, :3], axis=0).shape[0] != n_traj:
            raise AssertionError(
                "duplicate (block, group_row, member_col) in slice")

        t_values = np.asarray(_npz_member_memmap(fine_path, "t_values"))
        if boundary_idx[-1] > t_values.shape[0] - 1:
            raise AssertionError("boundary_idx exceeds fine snapshot grid")
        if not np.allclose(t_values[boundary_idx], t_vals_man, atol=1e-12):
            raise AssertionError(
                f"manifest t_values {t_vals_man.tolist()} do not match the "
                f"stored grid at boundary_idx {boundary_idx.tolist()}")
        if expected_t_values is not None:
            exp = np.asarray([float(t) for t in expected_t_values])
            if len(exp) != len(t_vals_man) or not np.allclose(exp, t_vals_man,
                                                              atol=1e-9):
                raise AssertionError(
                    f"sampling.t_values {exp.tolist()} != slice-manifest "
                    f"t_values {t_vals_man.tolist()}")

        # --- gid guards against the dataset's own key arrays + split.json ---
        with open(os.path.join(packaged_dataset_path, "split.json")) as f:
            split = json.load(f)
        train_set = set(split["train"])
        forbidden = set(split["val"]) | set(split["test"])
        for blk, (_, gid_key) in self._BLOCKS.items():
            sub = rows[rows[:, 0] == blk]
            if sub.shape[0] == 0:
                continue
            block_gids = np.asarray(_npz_member_memmap(fine_path, gid_key))
            if not np.array_equal(block_gids[sub[:, 1]], sub[:, 3]):
                raise AssertionError(
                    f"manifest gid column inconsistent with "
                    f"{gid_key}[group_row] (block {blk})")
        man_gids = set(rows[:, 3].tolist())
        if man_gids & forbidden:
            raise AssertionError("slice manifest contains val/test gids")
        if not man_gids <= train_set:
            raise AssertionError("slice manifest contains gids outside "
                                 "split.train")

        # --- gather (N, 2, D) raw endpoints in manifest row order ---
        first_blk = int(rows[0, 0])
        probe = _npz_member_memmap(fine_path, self._BLOCKS[first_blk][0])
        D = probe.shape[-1]
        seg = np.empty((n_traj, 2, D), dtype=np.float64)
        for blk, (vec_key, _) in self._BLOCKS.items():
            sel = np.nonzero(rows[:, 0] == blk)[0]
            if sel.size == 0:
                continue
            vecs = _npz_member_memmap(fine_path, vec_key)
            seg[sel] = np.asarray(
                vecs[rows[sel, 1], rows[sel, 2]])[:, boundary_idx, :]
        if not np.all(np.isfinite(seg)):
            raise AssertionError("non-finite values in selected trajectories")

        norm = np.load(os.path.join(packaged_dataset_path, "normalization.npz"))
        mean, std = norm["mean"], norm["std"]
        # Shared endpoint convention: float64 arithmetic, then float32 cast.
        self.x0 = torch.from_numpy(((seg[:, 0] - mean) / std).astype(np.float32))
        self.xT = torch.from_numpy(((seg[:, 1] - mean) / std).astype(np.float32))
        if self.x0.shape[0] != B:
            raise AssertionError(f"built {self.x0.shape[0]} pairs, expected {B}")

        # Provenance of each pair (block, member, gid).
        self.pair_block  = rows[:, 0].copy()
        self.pair_member = rows[:, 2].copy()
        self.pair_gid    = rows[:, 3].copy()
        self.n_segments = 1
        self.manifest = man
        self.seed = int(seed)

    def __len__(self) -> int:
        return self.x0.shape[0]

    def __getitem__(self, idx: int):
        return self.x0[idx], self.xT[idx]


class ValEndpointPairDataset(Dataset):
    """Validation endpoint pairs ``(x_0, x_T)`` at the manifest's ``[0, idx_T]`` boundaries from the validation gids
    of ``vecs_eval``, first ``n_members`` members per gid (the ``val_loss`` column of the log)."""

    def __init__(self, packaged_dataset_path: str, slice_manifest: str,
                 n_members: int = 20):
        man = _load_manifest(slice_manifest)
        if man.get("schema") != ENDPOINT_SCHEMA:
            raise AssertionError(
                f"expected schema {ENDPOINT_SCHEMA}, got "
                f"{man.get('schema')!r} in {slice_manifest}")
        fine_path = os.path.join(packaged_dataset_path, FINE_FILE)
        boundary_idx = np.asarray(man["boundary_idx"], dtype=np.int64)

        with open(os.path.join(packaged_dataset_path, "split.json")) as f:
            split = json.load(f)
        val_gids = sorted(split["val"])

        gids_eval = np.asarray(_npz_member_memmap(fine_path, "gids_eval"))
        vecs_eval = _npz_member_memmap(fine_path, "vecs_eval")
        rows = []
        for gid in val_gids:
            where = np.nonzero(gids_eval == gid)[0]
            if len(where) != 1:
                raise AssertionError(f"val gid {gid} not found once in gids_eval")
            rows.append(int(where[0]))
        n_members = int(n_members)
        if not (1 <= n_members <= vecs_eval.shape[1]):
            raise ValueError(f"n_members {n_members} out of range")

        seg = np.asarray(vecs_eval[rows][:, :n_members][:, :, boundary_idx, :])
        if not np.all(np.isfinite(seg)):
            raise AssertionError("non-finite values in vecs_eval slice")
        D = seg.shape[-1]
        seg = seg.reshape(-1, 2, D).astype(np.float64)

        norm = np.load(os.path.join(packaged_dataset_path, "normalization.npz"))
        mean, std = norm["mean"], norm["std"]
        self.x0 = torch.from_numpy(((seg[:, 0] - mean) / std).astype(np.float32))
        self.xT = torch.from_numpy(((seg[:, 1] - mean) / std).astype(np.float32))
        self.n_segments = 1

    def __len__(self) -> int:
        return self.x0.shape[0]

    def __getitem__(self, idx: int):
        return self.x0[idx], self.xT[idx]


# --------------------------------------------------------------------------- #
# sampling and validation inputs at the forecast horizon
# --------------------------------------------------------------------------- #

def _x0_and_cloud_from_bank(bank_dir: str, gid: int, idx_T: int, cloud_n: int = 200):
    """x_0 and a truth cloud of a test condition that the package does not carry (waves, NLKG), read from the
    truth bank (group_<gid>.npz); used only for the test split when ``x0_bank`` is given. Fails loudly on any
    inconsistency."""
    p = os.path.join(bank_dir, f"group_{int(gid):04d}.npz")
    if not os.path.isfile(p):
        raise FileNotFoundError(
            f"test-x0-bank fallback: no group_{int(gid):04d}.npz under {bank_dir} for gid {gid}")
    with np.load(p) as z:
        if int(z["gid"]) != int(gid):
            raise AssertionError(
                f"test-x0-bank fallback: {p} declares gid {int(z['gid'])}, expected {gid} "
                "(bank file must be looked up and verified by gid, never by row position)")
        vecs = np.asarray(z["vecs"])                        # (N, n_snapshots, D)
    if not (0 <= idx_T < vecs.shape[1]):
        raise ValueError(f"test-x0-bank fallback: fine_horizon_idx {idx_T} outside bank grid for gid {gid}")
    x0_members = vecs[:, 0, :]
    spread = float(np.max(np.abs(x0_members - x0_members[0])))
    if spread > 1e-9:
        raise AssertionError(
            f"test-x0-bank fallback: gid {gid} x0 differs across bank members by {spread:.3e} — "
            "shared-IC assumption violated")
    n = min(int(cloud_n), vecs.shape[0])
    return x0_members[0], np.asarray(vecs[:n, idx_T, :])


def get_fine_split_set(cfg: dict, split_name: str, seed: int = 0, x0_bank: str | None = None
                       ) -> Tuple[torch.Tensor, torch.Tensor, List[int]]:
    """Per-gid x0 + truth cloud at the snapshot ``cfg.method.fine_horizon_idx``:
    ``(x0_per_gid (n,D) f32, truth_clouds (n,V,D) f32, gids list)``.

    x0 convention: member 0 at t=0. Cross-member spread at t=0 is asserted < 1e-9.

    ``x0_bank`` (test split only): a truth-bank directory (group_<gid>.npz) used for a test gid absent from the
    package's ``vecs_eval`` block; it changes nothing for the other splits or for a gid the package carries.
    """
    packaged_dir = cfg["data"]["packaged_dataset_path"]
    method_cfg = cfg.get("method", {})
    idx_T = int(method_cfg["fine_horizon_idx"])
    fine_path = os.path.join(packaged_dir, FINE_FILE)

    t_values = np.asarray(_npz_member_memmap(fine_path, "t_values"))
    if not (1 <= idx_T <= t_values.shape[0] - 1):
        raise ValueError(f"fine_horizon_idx {idx_T} outside grid")
    # The horizon index must agree with the sampler chain's terminal time.
    samp_t = method_cfg.get("sampling", {}).get("t_values")
    if samp_t is not None:
        if abs(float(samp_t[-1]) - float(t_values[idx_T])) > 1e-9:
            raise AssertionError(
                f"fine_horizon_idx={idx_T} (t={t_values[idx_T]}) != "
                f"sampling.t_values[-1]={samp_t[-1]}")

    with open(os.path.join(packaged_dir, "split.json")) as f:
        split = json.load(f)
    gids_override = cfg["data"].get("gids_override")
    if gids_override is not None:
        gids = sorted(int(g) for g in gids_override)
    else:
        if split_name not in split:
            raise KeyError(f"split '{split_name}' not in split.json")
        gids = sorted(split[split_name])
        if split_name == "train":
            # the training block carries the conditions of vecs_train only (not vecs_new)
            canonical = set(np.asarray(
                _npz_member_memmap(fine_path, "gids_train")).tolist())
            gids = [g for g in gids if g in canonical]

    if split_name in ("val", "test") and gids_override is None:
        block, gid_key = "vecs_eval", "gids_eval"
    else:
        # the training conditions of vecs_train; any other gid fails loudly below
        block, gid_key = "vecs_train", "gids_train"
    block_gids = np.asarray(_npz_member_memmap(fine_path, gid_key))
    vecs = _npz_member_memmap(fine_path, block)

    norm = np.load(os.path.join(packaged_dir, "normalization.npz"))
    mean, std = norm["mean"], norm["std"]

    x0_list, cloud_list = [], []
    for gid in gids:
        where = np.nonzero(block_gids == gid)[0]
        if len(where) != 1:
            if split_name == "test" and x0_bank is not None:
                x0_g, cloud_g = _x0_and_cloud_from_bank(x0_bank, gid, idx_T)
                x0_list.append(x0_g); cloud_list.append(cloud_g)
                continue
            raise KeyError(
                f"gid {gid} not found in {gid_key}: the package serves the conditions of "
                "vecs_train and the validation and test conditions of vecs_eval")
        r = int(where[0])
        x0_members = np.asarray(vecs[r, :, 0, :])          # (V, D)
        spread = float(np.max(np.abs(x0_members - x0_members[0])))
        if spread > 1e-9:
            raise AssertionError(
                f"gid {gid}: x0 differs across members by {spread:.3e} — "
                "shared-IC assumption violated")
        x0_list.append(x0_members[0])
        cloud_list.append(np.asarray(vecs[r, :, idx_T, :]))
    x0_raw = np.stack(x0_list, axis=0)
    clouds_raw = np.stack(cloud_list, axis=0)
    if not (np.all(np.isfinite(x0_raw)) and np.all(np.isfinite(clouds_raw))):
        raise AssertionError("non-finite values in fine split set")

    # Same normalization arithmetic as get_split_set (float64 then cast).
    x0 = torch.from_numpy(((x0_raw - mean) / std).astype(np.float32))
    clouds = torch.from_numpy(((clouds_raw - mean) / std).astype(np.float32))
    return x0, clouds, gids