loader.py

"""Lightweight loaders for the MARIS-Forecast release.

Only ``pandas`` and ``numpy`` are required.  The loaders operate on a local
copy of the dataset downloaded from Zenodo (DOI 10.5281/zenodo.21224009) or
the Hugging Face mirror; nothing here downloads data.

Typical use
-----------
>>> from maris_forecast import loader
>>> leaf = loader.leaf_dir("/data/maris", track="A", region="dma")
>>> df = loader.load_split(leaf, "test", parse_arrays=True, consistent_only=True)
>>> hist, fut = loader.stack_trajectories(df)          # (N,30,2), (N,30,2)
>>> env = loader.load_environment(leaf, "test", align_to=df)  # masks / SDF aligned
"""
from __future__ import annotations
import json
import os
import glob
import numpy as np
import pandas as pd

from . import schema

__all__ = [
    "leaf_dir", "load_split", "stack_trajectories", "filter_consistent",
    "load_environment", "load_social", "load_environment_descriptors",
]


def leaf_dir(root: str, track: str, region: str) -> str:
    """Absolute path to a jurisdictional leaf directory inside ``root``."""
    return os.path.join(root, schema.leaf_relpath(track, region))


# --------------------------------------------------------------------- splits
def _split_files(leaf: str, split: str):
    files = sorted(glob.glob(os.path.join(leaf, split, "part-*.csv.gz")))
    if not files:
        raise FileNotFoundError(
            f"no split files under {os.path.join(leaf, split)} "
            f"(expected part-*.csv.gz)")
    return files


def _to_bool(series: pd.Series) -> pd.Series:
    """Robustly coerce a bool-ish column (True/False or 'true'/'false')."""
    if series.dtype == bool:
        return series
    return series.astype(str).str.strip().str.lower().isin(("true", "1", "yes"))


def filter_consistent(df: pd.DataFrame) -> pd.DataFrame:
    """Keep only OSM-temporally consistent rows (the recommended default view)."""
    if "osm_temporal_consistent" not in df.columns:
        raise KeyError("column 'osm_temporal_consistent' not present")
    return df[_to_bool(df["osm_temporal_consistent"])].reset_index(drop=True)


def _merge_qc_sidecar(leaf: str, split: str, df: pd.DataFrame) -> pd.DataFrame:
    """Attach the four Stage-17 flags from the side-car if not already inline."""
    if "osm_temporal_consistent" in df.columns:
        return df
    path = os.path.join(leaf, schema.qc_flags(split))
    if not os.path.exists(path):
        raise KeyError(
            "column 'osm_temporal_consistent' is not inline and no side-car "
            f"found at {path}")
    flags = pd.read_csv(path, usecols=lambda c: c in
                        (("sample_id",) + schema.CONSISTENCY_COLS))
    return df.merge(flags, on="sample_id", how="left")


def load_split(leaf: str, split: str, *, parse_arrays: bool = False,
               consistent_only: bool = False, columns=None) -> pd.DataFrame:
    """Load one split (train/val/test) of a leaf as a DataFrame.

    Parameters
    ----------
    parse_arrays : if True, JSON-array columns (geometry + kinematics) are
        parsed into python lists of floats.
    consistent_only : if True, apply the Stage-17 consistency filter.  The
        four flags are read inline when present (the release layout) or from
        the ``osm_temporal_consistency/`` side-car otherwise.
    columns : optional subset of columns to read (speeds up large splits).
    """
    frames = [pd.read_csv(f, usecols=columns) for f in _split_files(leaf, split)]
    df = pd.concat(frames, ignore_index=True) if len(frames) > 1 else frames[0]
    if consistent_only:
        df = _merge_qc_sidecar(leaf, split, df)
        df = filter_consistent(df)
    if parse_arrays:
        for col in schema.GEOMETRY_COLS + schema.KINEMATIC_COLS:
            if col in df.columns:
                df[col] = df[col].apply(json.loads)
    return df


def stack_trajectories(df: pd.DataFrame):
    """Return (history, future) float arrays of shape (N, T, 2), metres.

    Accepts a DataFrame whose geometry columns are either JSON strings or
    already-parsed lists.
    """
    def _col(name):
        v = df[name]
        return [x if isinstance(x, (list, tuple)) else json.loads(x) for x in v]

    hx, hy = _col("hist_x_json"), _col("hist_y_json")
    fx, fy = _col("fut_x_json"), _col("fut_y_json")
    hist = np.stack([np.asarray(hx, float), np.asarray(hy, float)], axis=-1)
    fut = np.stack([np.asarray(fx, float), np.asarray(fy, float)], axis=-1)
    return hist, fut


# ----------------------------------------------------------------- context
def _load_masks(raster_dir: str):
    """Load the six-channel raster, supporting both .npz (release) and .npy."""
    npz = os.path.join(raster_dir, "masks.npz")
    npy = os.path.join(raster_dir, "masks.npy")
    if os.path.exists(npz):
        with np.load(npz) as z:
            return z["masks"]
    if os.path.exists(npy):
        return np.load(npy, mmap_mode="r")
    raise FileNotFoundError(f"no masks.npz or masks.npy under {raster_dir}")


def load_environment(leaf: str, split: str, *, align_to: pd.DataFrame | None = None,
                     which: str = "context_v1", memmap: bool = True):
    """Load environmental tensors for a split.

    Returns a dict with ``masks`` (N,6,128,128 uint8), ``signed_dist_shore``
    and ``signed_dist_nav`` (N,128,128 float16), and ``sample_ids``.  If
    ``align_to`` is given, the tensors are reordered/subset to match that
    DataFrame's ``sample_id`` order.
    """
    rdir = os.path.join(leaf, which, "environment", "rasters", split)
    sids = np.load(os.path.join(rdir, "sample_ids.npy"), allow_pickle=True)
    masks = _load_masks(rdir)
    mm = "r" if memmap else None
    shore = np.load(os.path.join(rdir, "signed_dist_shore.npy"), mmap_mode=mm)
    nav = np.load(os.path.join(rdir, "signed_dist_nav.npy"), mmap_mode=mm)
    out = {"masks": masks, "signed_dist_shore": shore,
           "signed_dist_nav": nav, "sample_ids": sids}
    if align_to is not None:
        pos = {s: i for i, s in enumerate(sids)}
        idx = np.array([pos[s] for s in align_to["sample_id"]], dtype=int)
        out = {"masks": np.asarray(masks)[idx],
               "signed_dist_shore": np.asarray(shore)[idx],
               "signed_dist_nav": np.asarray(nav)[idx],
               "sample_ids": sids[idx]}
    return out


def load_environment_descriptors(leaf: str, split: str,
                                 which: str = "context_v1") -> pd.DataFrame:
    """Load the fourteen-dimensional per-sample scene descriptors."""
    return pd.read_csv(os.path.join(leaf, which, "environment", "features",
                                    split, "environment_descriptors.csv"))


def load_social(leaf: str, split: str) -> pd.DataFrame:
    """Load the target-centric social-neighbour records for a split."""
    return pd.read_csv(os.path.join(leaf, "context_v1", "social", "features",
                                    split, "social_features.csv"))