tracking_repo/src/biohub_tracking/io.py

"""I/O utilities for tracking challenge datasets."""

from dataclasses import dataclass, field
from pathlib import Path

import dask.array as da
import numpy as np
import polars as pl
import torch
import tracksdata as td
import zarr

DEFAULT_SCALE: tuple[float, float, float] = (1.625, 0.40625, 0.40625)


@dataclass
class Dataset:
    path: Path
    image: da.Array | np.ndarray | torch.Tensor | None
    tracks: td.graph.IndexedRXGraph | None
    scale: tuple[float, float, float]
    original_scale: tuple[float, float, float] | None = None
    image_shape: tuple[int, ...] | None = None  # (T, Z, Y, X), always available
    zarr_path: Path | None = None
    quantiles: dict[str, float] = field(default_factory=dict)

    def napari_tracks(self) -> tuple[np.ndarray, dict[int, int]]:
        return td.functional.to_napari_format(
            self.tracks,
            self.image.shape,
            solution_key=None,
            output_tracklet_id_key="track_id",
            mask_key=None,
        )


def open_dataset(ds_path: Path | str,
                target_scale: tuple[float, float, float] | None = None,
                normalize: bool = True,
                gamma: float = 1.0,
                device: str = "cuda",
                require_tracks: bool = False,
                load_image: bool = True,
                downsample: tuple[int, ...] | None = None) -> Dataset:
    """Open a dataset from a zarr file and optionally a geff tracks file.

    Parameters
    ----------
    ds_path : Path or str
        Path to the dataset (without extension, or with .zarr/.geff).
    target_scale : tuple, optional
        Resample to this isotropic (Z, Y, X) voxel scale. If None, no resampling.
    normalize : bool
        Whether to normalize the image intensities.
    gamma : float
        Gamma correction value for normalization.
    device : str
        Device to use for GPU processing.
    require_tracks : bool
        Whether to require a .geff tracks file.
    load_image : bool
        If False, skip loading image data (metadata + tracks only).
    downsample : tuple, optional
        Spatial downsample strides (Z, Y, X). Applied via strided indexing.

    Returns
    -------
    Dataset
        Dataset with image shape (T, Z, Y, X). When ``load_image=False``,
        ``image`` is None but ``image_shape`` and ``zarr_path`` are populated.
    """
    ds_path = Path(ds_path)
    if ds_path.suffix in (".zarr", ".geff"):
        ds_path = ds_path.parent / ds_path.stem

    image_path = ds_path.parent / f"{ds_path.stem}.zarr"
    tracks_path = ds_path.parent / f"{ds_path.stem}.geff"

    if not image_path.exists():
        raise FileNotFoundError(f"Image file not found: {image_path}")
    if require_tracks and not tracks_path.exists():
        raise FileNotFoundError(f"Tracks file not found: {tracks_path}")

    img_ds = zarr.open_group(image_path, mode="r")
    tracks = None
    if require_tracks and tracks_path.exists():
        result = td.graph.IndexedRXGraph.from_geff(tracks_path)
        tracks = result[0] if isinstance(result, tuple) else result

    attrs = dict(img_ds.attrs)
    scale = _parse_scale(attrs)
    quantiles = attrs.get("image_statistics", {}).get("quantiles", {})

    raw_shape = tuple(img_ds["0"].shape)
    if downsample is not None:
        ds_shape = raw_shape[:1] + tuple(-(-s // d) for s, d in zip(raw_shape[1:], downsample))
    else:
        ds_shape = raw_shape

    if not load_image:
        return Dataset(
            path=ds_path, image=None, tracks=tracks, scale=scale,
            image_shape=ds_shape, zarr_path=image_path, quantiles=quantiles,
        )

    da_arr = da.from_zarr(img_ds["0"]).compute()

    resample = target_scale is not None
    original_scale = None
    if resample or normalize:
        if resample:
            original_scale = scale
        da_arr, tracks, scale = _process_on_gpu(
            da_arr, tracks, scale, device,
            resample=resample, target_scale=target_scale,
            normalize=normalize, gamma=gamma,
            precomputed_quantiles=quantiles,
        )

    assert da_arr.ndim == 4

    return Dataset(
        path=ds_path, image=da_arr, tracks=tracks, scale=scale,
        original_scale=original_scale, image_shape=tuple(da_arr.shape),
        zarr_path=image_path, quantiles=quantiles,
    )


def _parse_scale(attrs: dict) -> tuple[float, float, float]:
    """Extract (Z, Y, X) voxel scale from OME-NGFF zarr attrs."""
    if "multiscales" in attrs:
        transform = attrs["multiscales"][0]["datasets"][0]["coordinateTransformations"][0]
        if transform["type"] != "scale":
            raise ValueError(f"Transform type is not 'scale': {transform}")
        return tuple(transform["scale"][-3:])
    return DEFAULT_SCALE


def _lookup_precomputed_quantile(
    quantiles: dict[str, float] | None, q: float, tol: float = 1e-6,
) -> float | None:
    """Return the stored quantile value for *q* if present in *quantiles*."""
    if not quantiles:
        return None
    for key, val in quantiles.items():
        try:
            if abs(float(key) - q) <= tol:
                return float(val)
        except (TypeError, ValueError):
            continue
    return None


def _process_on_gpu(
    image: np.ndarray,
    tracks: td.graph.IndexedRXGraph | None,
    scale: tuple[float, float, float],
    device: str,
    resample: bool = False,
    target_scale: tuple[float, float, float] | None = None,
    normalize: bool = True,
    gamma: float = 1.0,
    q_min: float = 0.000,
    q_max: float = 1.000,
    subsample_factor: int = 50,
    precomputed_quantiles: dict[str, float] | None = None,
) -> tuple[torch.Tensor, td.graph.IndexedRXGraph | None, tuple[float, float, float]]:
    """Process image on GPU: normalization and isotropic resampling.

    All GPU processing is consolidated here to minimize CPU<->GPU transfers.

    Expects 4D input: (T, Z, Y, X)

    Processing order:
    1. Normalize (quantile + gamma) - on CPU/numpy
    2. Transfer to GPU
    3. Resample to isotropic (if enabled)
    4. Transfer back to CPU
    """
    torch_device = torch.device(device)

    # Convert to float32 for processing (zarr data is typically uint16)
    image = image.astype(np.float32)

    # 1. Quantile stats for normalization. Prefer precomputed values from zarr
    #    attrs (instant); fall back to a subsampled np.quantile if missing.
    q1, q2 = None, None
    if normalize:
        q1 = _lookup_precomputed_quantile(precomputed_quantiles, q_min)
        q2 = _lookup_precomputed_quantile(precomputed_quantiles, q_max)
        if q1 is None or q2 is None:
            flat = image.ravel()[::subsample_factor]
            q1, q2 = np.quantile(flat, [q_min, q_max]).astype(np.float32)
        else:
            q1 = np.float32(q1)
            q2 = np.float32(q2)

    # 2. Transfer raw float32 to GPU once (pinned memory for faster DMA)
    tensor = torch.from_numpy(image).pin_memory().to(torch_device, non_blocking=True)

    # 3. Apply normalization on GPU
    if normalize:
        tensor = (tensor - float(q1)) / (float(q2) - float(q1) + 1e-6)
        tensor = tensor.clamp(min=0.0)
        if gamma != 1.0:
            tensor = tensor.pow(gamma)
        tensor = tensor.clamp(0.0, 4.0)

    # 4. Resample to target scale (if enabled)
    if resample:
        scale_arr = np.array(scale)
        if target_scale is None:
            target_scale_val = scale_arr.min()
        else:
            target_scale_val = np.array(target_scale)

        zoom_factors = scale_arr / target_scale_val

        # For 4D input (T, Z, Y, X), reshape to (T, 1, Z, Y, X) for interpolate
        T = tensor.shape[0]
        new_spatial_shape = (np.array(tensor.shape[1:]) * zoom_factors).astype(int).tolist()
        tensor = tensor[:, None]  # (T, 1, Z, Y, X)
        tensor = torch.nn.functional.interpolate(
            tensor, size=new_spatial_shape, mode="trilinear", align_corners=False
        )
        tensor = tensor[:, 0]  # (T, Z, Y, X)

        # Update track coordinates if tracks are provided
        if tracks is not None:
            node_attrs = tracks.node_attrs()
            orig_dtypes = {col: node_attrs.schema[col] for col in ["z", "y", "x"]}
            node_attrs = node_attrs.with_columns(
                (pl.col("z") * zoom_factors[0]).round(0).cast(orig_dtypes["z"]),
                (pl.col("y") * zoom_factors[1]).round(0).cast(orig_dtypes["y"]),
                (pl.col("x") * zoom_factors[2]).round(0).cast(orig_dtypes["x"]),
            )
            tracks.update_node_attrs(
                attrs=node_attrs.select("z", "y", "x").to_dict(),
                node_ids=node_attrs[td.DEFAULT_ATTR_KEYS.NODE_ID].to_list(),
            )

        # Update scale to isotropic
        if target_scale is not None:
            scale = target_scale
        else:
            scale = (float(target_scale_val),) * 3

    return tensor, tracks, scale



def invert_time_graph(tracks: td.graph.IndexedRXGraph, max_t: int = 100) -> td.graph.IndexedRXGraph:
    """Invert the time axis of the image and tracks."""

    # Update node time attributes

    node_attrs = tracks.node_attrs()
    node_attrs = node_attrs.with_columns(
        ((max_t - 1) - pl.col("t")).alias("t"),
    )
    tracks.update_node_attrs(
        attrs=node_attrs.select("t").to_dict(),
        node_ids=node_attrs[td.DEFAULT_ATTR_KEYS.NODE_ID].to_list(),
    )

    # Collect all edge info first before modifying the graph
    edge_attrs = tracks.edge_attrs()
    edges_to_reverse = []
    for row in edge_attrs.iter_rows(named=True):
        source = row[td.DEFAULT_ATTR_KEYS.EDGE_SOURCE]
        target = row[td.DEFAULT_ATTR_KEYS.EDGE_TARGET]
        edge_id = row[td.DEFAULT_ATTR_KEYS.EDGE_ID]
        attrs = {k: v for k, v in row.items()
                    if k not in (td.DEFAULT_ATTR_KEYS.EDGE_SOURCE,
                                td.DEFAULT_ATTR_KEYS.EDGE_TARGET,
                                td.DEFAULT_ATTR_KEYS.EDGE_ID)}
        edges_to_reverse.append((edge_id, source, target, attrs))

    # Now modify the graph
    for edge_id, source, target, attrs in edges_to_reverse:
        tracks.remove_edge(edge_id=edge_id)
        tracks.add_edge(source_id=target, target_id=source, attrs=attrs)

    return tracks


def rescale_graph_to_original(
    graph: td.graph.BaseGraph,
    original_scale: tuple[float, float, float],
) -> td.graph.BaseGraph:
    """
    Rescale graph coordinates from isotropic space back to original anisotropic space.

    This inverts the isotropic resampling transformation applied during dataset loading.

    Parameters
    ----------
    graph : td.graph.BaseGraph
        Graph with coordinates in isotropic space.
    original_scale : tuple[float, float, float]
        Original (z, y, x) voxel scale before isotropic resampling.

    Returns
    -------
    td.graph.BaseGraph
        Graph with coordinates scaled back to original space.
    """
    scale_arr = np.array(original_scale)
    target_scale = scale_arr.min()
    zoom_factors = scale_arr / target_scale

    # Invert the zoom by dividing coordinates
    inverse_zoom = 1.0 / zoom_factors

    node_attrs = graph.node_attrs()
    node_attrs = node_attrs.with_columns(
        (pl.col("z") * inverse_zoom[0]),
        (pl.col("y") * inverse_zoom[1]),
        (pl.col("x") * inverse_zoom[2]),
    )
    graph.update_node_attrs(
        attrs=node_attrs.select("z", "y", "x").to_dict(),
        node_ids=node_attrs[td.DEFAULT_ATTR_KEYS.NODE_ID].to_list(),
    )

    return graph


def save_graph(graph: td.graph.BaseGraph, output_path: Path | str, overwrite: bool = True) -> None:
    """
    Save a tracksdata graph to a .geff file.

    Parameters
    ----------
    graph : td.graph.BaseGraph
        The graph to save.
    output_path : Path or str
        Output path for the .geff file.
    overwrite : bool
        Whether to overwrite existing files.
    """
    import shutil

    output_path = Path(output_path)
    output_path.parent.mkdir(parents=True, exist_ok=True)

    # Ensure path has .geff extension
    if not output_path.suffix == ".geff":
        output_path = output_path.with_suffix(".geff")

    # Remove existing file/directory if overwrite is enabled
    if overwrite and output_path.exists():
        if output_path.is_dir():
            shutil.rmtree(output_path)
        else:
            output_path.unlink()

    graph.to_geff(output_path)


def list_datasets(data_dir: Path | str, require_geff: bool = True) -> list[Path]:
    """
    List all valid datasets in a directory.

    A valid dataset has a .zarr file and optionally a corresponding .geff file.

    Parameters
    ----------
    data_dir : Path or str
        Directory containing datasets.
    require_geff : bool
        If True, only return datasets with both .zarr and .geff files.
        If False, return all .zarr files.

    Returns
    -------
    list[Path]
        List of dataset paths (without extension).
    """
    data_dir = Path(data_dir)

    if not data_dir.exists():
        raise FileNotFoundError(f"Directory not found: {data_dir}")

    zarr_files = sorted(data_dir.glob("*.zarr"))

    datasets = []
    for zarr_file in zarr_files:
        if require_geff:
            geff_file = data_dir / f"{zarr_file.stem}.geff"
            if geff_file.exists():
                datasets.append(data_dir / zarr_file.stem)
        else:
            datasets.append(data_dir / zarr_file.stem)

    return datasets