labeled_data_cache

def build_from_zarr(
    self,
    found_locations: dict[str, tuple[str, int]],
    label_df: pd.DataFrame,
    cfg: DotMap,
) -> None:
    """Extract raw labeled images from Zarr stores and write to cache.

    Reads directly from known ``(store_path, local_index)`` positions
    without building a full source index.

    Args:
        found_locations: Map of ``id -> (store_path, local_index)``
            from validation.
        label_df: Label DataFrame with ``id``, ``label`` columns.
        cfg: Configuration (used for extraction hash).
    """
    self.clear()
    self._cache_dir.mkdir(parents=True, exist_ok=True)

    if not found_locations:
        logger.warning("LabeledDataCache: no locations to build from")
        return

    images, metadata_rows = self._extract_zarr_images(found_locations, label_df)
    self._write_cache(images, metadata_rows, DataSourceType.ZARR, cfg, label_df=label_df)
    logger.info("LabeledDataCache: built Zarr cache with {} images", len(images))

def build_from_cutana(
    self,
    found_locations: dict[str, tuple[str, int]],
    label_df: pd.DataFrame,
    cfg: DotMap,
) -> None:
    """Extract raw labeled cutouts from Cutana catalogues and write to cache.

    Groups source IDs by catalogue file and streams cutouts via
    Cutana with normalisation *disabled* (``CONVERSION_ONLY`` + float32
    output) at :data:`LABELED_CACHE_RESOLUTION`.  That way the
    stored values survive any change to ``normalisation_method``,
    ``output_dtype`` or ``channel_combination`` — those are applied
    at read time by
    :func:`~anomaly_match.data_io.container_loaders.decode_cutana_raw_images`.

    Args:
        found_locations: Map of ``source_id -> (catalogue_path, row_index)``
            from validation.
        label_df: Label DataFrame with ``id``, ``label`` columns.
        cfg: Configuration.  Only raw-extraction fields
            (fits_extension, padding, flux conversion,
            interpolation_order) affect the cached pixels;
            normalisation / channel-combination fields are ignored
            at build time.
    """
    self.clear()
    self._cache_dir.mkdir(parents=True, exist_ok=True)

    if not found_locations:
        logger.warning("LabeledDataCache: no locations to build from")
        return

    # Work from a cfg override that disables normalisation and pins
    # the resolution to the cache size — so the rest of the build
    # pipeline (orchestrator config, hash, write_cache) all see the
    # same raw-extraction settings even if a UI widget callback
    # mutates the shared cfg mid-build on another thread.
    extraction_cfg = _build_raw_extraction_cfg(cfg)

    images, metadata_rows = self._extract_cutana_images(
        found_locations, label_df, extraction_cfg
    )
    self._write_cache(
        images, metadata_rows, DataSourceType.CUTANA, extraction_cfg, label_df=label_df
    )
    logger.info(
        "LabeledDataCache: built Cutana cache with {} image(s) at {} px (raw float32)",
        len(images),
        extraction_cfg.normalisation.image_size[0],
    )

def append_from_zarr(
    self,
    new_locations: dict[str, tuple[str, int]],
    new_label_df: pd.DataFrame,
    full_label_df: pd.DataFrame,
) -> None:
    """Append newly labeled Zarr images to an existing cache.

    Args:
        new_locations: New ``id -> (store_path, local_index)`` entries.
        new_label_df: Label DataFrame for the new entries.
        full_label_df: The complete label DataFrame (whole CSV) used to
            refresh ``label_csv_hash`` — see :meth:`_append_to_cache`.

    Raises:
        RuntimeError: If the cache has not been built yet.
    """
    if not self.is_populated:
        raise RuntimeError("Cannot append to an unpopulated cache")

    existing_meta = self.get_label_df()
    existing_names = set(existing_meta["id"])

    # Only extract truly new entries
    new_locs = {fn: loc for fn, loc in new_locations.items() if fn not in existing_names}
    if not new_locs:
        logger.debug("LabeledDataCache: no new entries to append")
        return

    images, metadata_rows = self._extract_zarr_images(new_locs, new_label_df)
    self._append_to_cache(images, metadata_rows, full_label_df)
    logger.info("LabeledDataCache: appended {} images", len(images))

def append_from_cutana(
    self,
    new_locations: dict[str, tuple[str, int]],
    new_label_df: pd.DataFrame,
    cfg: DotMap,
    full_label_df: pd.DataFrame,
) -> None:
    """Append newly labeled Cutana cutouts to an existing cache.

    Args:
        new_locations: New ``source_id -> (catalogue_path, row_index)`` entries.
        new_label_df: Label DataFrame for the new entries.
        cfg: Configuration for Cutana orchestrator.
        full_label_df: The complete label DataFrame (whole CSV) used to
            refresh ``label_csv_hash`` — see :meth:`_append_to_cache`.

    Raises:
        RuntimeError: If the cache has not been built yet.
    """
    if not self.is_populated:
        raise RuntimeError("Cannot append to an unpopulated cache")

    existing_meta = self.get_label_df()
    existing_ids = set(existing_meta["id"])

    new_locs = {sid: loc for sid, loc in new_locations.items() if sid not in existing_ids}
    if not new_locs:
        logger.debug("LabeledDataCache: no new entries to append")
        return

    # Build the same raw-extraction cfg as ``build_from_cutana`` so
    # the appended cutouts match the on-disk shape/dtype.
    extraction_cfg = _build_raw_extraction_cfg(cfg)
    # Pin the cache resolution to whatever the existing images use
    # (stored shape may exceed the default because an earlier build
    # saw a higher user image_size).  The hash guard rebuilds when
    # other extraction params drift, so this is safe.
    stored_shape = self.get_cache_info().get("image_shape")
    if stored_shape:
        extraction_cfg.normalisation.image_size = list(stored_shape[:2])
        extraction_cfg.fitsbolt_cfg = None
        extraction_cfg = get_fitsbolt_config(extraction_cfg)
    images, metadata_rows = self._extract_cutana_images(new_locs, new_label_df, extraction_cfg)
    self._append_to_cache(images, metadata_rows, full_label_df)
    logger.info("LabeledDataCache: appended {} images", len(images))

def get_raw_images(self) -> list[tuple[str, np.ndarray]]:
    """Read all cached raw images.

    Returns:
        List of ``(id, raw_image_array)`` tuples.
    """
    if not self.is_populated:
        return []

    zarr_path = self._cache_dir / self.IMAGES_ZARR
    root = zarr.open_group(str(zarr_path), mode="r")
    images_arr = root["images"]

    meta_df = self.get_label_df()
    ids = meta_df["id"].tolist()

    # One zarr chunk file per cutout over NFS, so this loop is the long
    # silent gap before training's "Decoding labeled cache" begins.  Route
    # a tqdm meter through ``tqdm_logging_file`` (same path as the decode
    # meter) so the training status label and progress bar advance during
    # the read instead of sitting on "Loading datasets...".
    results = []
    for i, item_id in enumerate(
        tqdm(
            ids,
            desc=_CACHE_READ_DESC,
            unit="img",
            file=tqdm_logging_file(),
            mininterval=1.0,
        )
    ):
        results.append((item_id, np.array(images_arr[i])))
    return results

def get_raw_images_by_id(self, ids: list[str]) -> list[tuple[str, np.ndarray]]:
    """Read a specific subset of cached raw images by id.

    Callers are expected to have already checked that the cache is
    populated and that every requested id is in it — use
    :meth:`is_populated` and :meth:`get_label_df` for that.  The
    method raises on both invariants so a mis-wired call site fails
    loudly instead of returning a silently shorter list.

    Args:
        ids: Item ids to read.  Must be non-empty and every id must
            exist in the cache.  Order is preserved in the output.

    Returns:
        List of ``(id, raw_image_array)`` tuples, one per input id,
        in the input order.

    Raises:
        ValueError: If *ids* is empty.
        RuntimeError: If the cache has not been built yet.
        KeyError: If any requested id is not in the cache.
    """
    if not ids:
        raise ValueError("get_raw_images_by_id called with empty ids list")
    if not self.is_populated:
        raise RuntimeError("Cannot read from an unpopulated labeled cache")

    meta_df = self.get_label_df()
    index_by_id = {str(row_id): idx for idx, row_id in enumerate(meta_df["id"])}

    missing = [str(item_id) for item_id in ids if str(item_id) not in index_by_id]
    if missing:
        preview = missing[:5]
        suffix = " ..." if len(missing) > 5 else ""
        raise KeyError(f"{len(missing)} id(s) not in labeled cache: {preview}{suffix}")

    zarr_path = self._cache_dir / self.IMAGES_ZARR
    root = zarr.open_group(str(zarr_path), mode="r")
    images_arr = root["images"]

    # Each cutout is a separate zarr chunk file, so reading the preview's
    # few hundred ids over NFS is a long silent stretch behind the
    # "Loading preview..." spinner.  This runs in the UI kernel (not the
    # training subprocess), where tqdm-to-stderr wouldn't reach the
    # loguru-based progress tap — so emit throttled INFO lines the setup
    # spinner's cutana-hint matcher recognises instead.
    total = len(ids)
    log_every = _cache_read_log_stride(total)
    results: list[tuple[str, np.ndarray]] = []
    for done, item_id in enumerate(ids, start=1):
        results.append((str(item_id), np.array(images_arr[index_by_id[str(item_id)]])))
        if done == total or done % log_every == 0:
            logger.info("{}: {}/{} cutouts", _CACHE_READ_DESC, done, total)
    return results

def get_label_df(self) -> pd.DataFrame:
    """Read the cached label metadata.

    Returns:
        DataFrame with ``id``, ``label`` columns (plus source-specific
        location columns like ``store_path`` or ``catalogue_path``).
    """
    meta_path = self._cache_dir / self.METADATA_PARQUET
    if not meta_path.exists():
        return pd.DataFrame(columns=["id", "label"])
    return pd.read_parquet(meta_path)

def get_cache_info(self) -> dict:
    """Read cache metadata (source type, extraction hash, etc.).

    Returns:
        Dict with cache metadata, or empty dict if not populated.
    """
    info_path = self._cache_dir / self.CACHE_INFO_JSON
    if not info_path.exists():
        return {}
    with open(info_path) as f:
        return json.load(f)

def needs_rebuild(
    self,
    cfg: DotMap,
    label_df: pd.DataFrame | None = None,
    found_ids: set[str] | None = None,
) -> bool:
    """Check if extraction-affecting parameters changed since cache was built.

    For Zarr: never rebuilds on config changes (raw array data is
    fixed), but still rebuilds on label CSV changes when *label_df*
    is supplied.

    For Cutana: rebuild when any *raw*-extraction parameter changed
    (FITS extension list, padding, flux conversion, interpolation
    order — see :func:`_compute_extraction_hash`), when the requested
    ``image_size`` exceeds the stored resolution (resize down is fine;
    we can't invent detail that wasn't cached), **or** when the
    source's resolved band count no longer matches the cached channel
    count (a band was added to or removed from the catalogues'
    ``fits_file_paths``).  Normalisation method, output dtype, channel
    combination and ``n_output_channels`` are all applied at read time
    and so never force a rebuild.

    When *label_df* is supplied, also compare its hash against
    ``cache_info.json['label_csv_hash']`` and rebuild on mismatch —
    this lets a persistent cache (next to the label CSV) detect when
    the CSV has been edited between sessions.

    When *found_ids* is supplied, compare it against the cached id
    set and rebuild on mismatch.  This catches the case where the
    CSV is unchanged but the data source grew (e.g. the user dropped
    another Cutana catalogue into ``data_dir``) — previously the
    cache went stale silently until the CSV hash or extraction
    params happened to change.

    Every rebuild decision is logged with its reason so spurious
    rebuilds are diagnosable from the session log alone.

    Args:
        cfg: Current configuration.
        label_df: Label DataFrame with ``id``/``label`` columns.  If
            supplied, its hash is compared against the stored one to
            detect CSV edits.
        found_ids: IDs the current validation pass matched against
            the source.  When supplied, mismatch against the cached
            ids triggers a rebuild.

    Returns:
        True if the cache should be rebuilt.
    """
    if not self.is_populated:
        logger.info("Labeled cache rebuild: cache not populated")
        return True

    info = self.get_cache_info()

    # Upgrade caches written in an older on-disk format (e.g. the pre-v2
    # one-file-per-cutout layout, which has no version field) by rebuilding
    # once into the current sharded+compressed format.
    stored_version = info.get("cache_format_version", 1)
    if stored_version != _CACHE_FORMAT_VERSION:
        logger.info(
            "Labeled cache rebuild: format version changed (stored={} current={})",
            stored_version,
            _CACHE_FORMAT_VERSION,
        )
        return True

    source_type = info.get("source_type", "")
    stored_hash = info.get("extraction_hash", "")

    current_hash = _compute_extraction_hash(source_type, cfg)
    if stored_hash != current_hash:
        logger.info(
            "Labeled cache rebuild: extraction hash changed "
            "(stored={} current={}) — current params: {}",
            stored_hash,
            current_hash,
            _extraction_params_for_log(source_type, cfg),
        )
        return True

    if label_df is not None:
        stored_csv_hash = info.get("label_csv_hash", "")
        current_csv_hash = _compute_label_csv_hash(label_df)
        if stored_csv_hash != current_csv_hash:
            logger.info(
                "Labeled cache rebuild: label CSV hash changed "
                "(stored={} current={}) — labels were edited since build",
                stored_csv_hash,
                current_csv_hash,
            )
            return True

    if found_ids is not None:
        cached_ids = set(self.get_label_df()["id"].astype(str))
        if cached_ids != found_ids:
            added = len(found_ids - cached_ids)
            removed = len(cached_ids - found_ids)
            logger.info(
                "Labeled cache rebuild: source-matched id set changed "
                "(cached={}, matched now={}, +{} -{})",
                len(cached_ids),
                len(found_ids),
                added,
                removed,
            )
            return True

    if source_type == DataSourceType.CUTANA:
        stored_shape = info.get("image_shape")
        # image_shape is whatever the writer stored, so treat its length as
        # untrusted: a cache written without a channel axis (or by a foreign
        # writer) must degrade to "can't tell" rather than IndexError here.
        if stored_shape and len(stored_shape) >= 2:
            requested = list(cfg.normalisation.image_size)
            if stored_shape[0] < requested[0] or stored_shape[1] < requested[1]:
                logger.info(
                    "Labeled cache rebuild: requested image_size {} larger "
                    "than cached resolution {} — can't upsample, rebuilding",
                    requested,
                    stored_shape,
                )
                return True

            # The cache stores one channel per resolved FITS band, and the
            # read path passes those bands straight through (channel
            # combination is applied on top, never used to invent a missing
            # band).  So if the source now resolves to a different band count
            # than the cache was built with — e.g. a NIR-J band was added to
            # the catalogues' ``fits_file_paths`` — the cached cutouts no
            # longer match the fresh unlabeled stream, and training would
            # stack 3-band labeled cutouts against a 4-band model.
            # ``fits_extension`` is already in the extraction hash, so this
            # only fires when the *catalogue's* per-source band list changed
            # while the config did not.
            #
            # Comparing a count resolved from the user *cfg* against one
            # stored from ``_build_raw_extraction_cfg(cfg)`` is only sound
            # because that helper rewrites normalisation_method /
            # output_dtype / image_size / fitsbolt_cfg but leaves
            # ``fits_extension`` alone, and band resolution keys only off
            # ``fits_extension``.  If that ever changes, the two counts
            # diverge permanently and every call rebuilds — pinned by
            # test_raw_extraction_cfg_preserves_fits_extension.
            #
            # Checked before resolving: a shape without a channel axis has
            # nothing to compare against, so skip the catalogue read too.
            if len(stored_shape) >= 3:
                current_bands = _current_source_band_count(cfg)
                if current_bands is not None and stored_shape[2] != current_bands:
                    logger.info(
                        "Labeled cache rebuild: source band count changed "
                        "(cached={} current={}) — the catalogues' fits_file_paths "
                        "now resolve to a different number of bands",
                        stored_shape[2],
                        current_bands,
                    )
                    return True
    return False

def clear(self) -> None:
    """Delete all cache files."""
    for name in (self.IMAGES_ZARR, self.METADATA_PARQUET, self.CACHE_INFO_JSON):
        path = self._cache_dir / name
        if path.is_dir():
            shutil.rmtree(path)
        elif path.exists():
            path.unlink()

def build(
    self,
    found_locations: dict[str, tuple],
    label_df: pd.DataFrame,
    source_type: str,
    cfg: DotMap,
) -> None:
    """Build the cache from validated locations, dispatching by source type.

    Args:
        found_locations: Map from ID to source-specific location tuple.
        label_df: Label DataFrame with ``id``, ``label`` columns.
        source_type: A :class:`DataSourceType` value.
        cfg: Configuration for extraction.

    Raises:
        ValueError: If *source_type* is not a container type.
    """
    if source_type == DataSourceType.ZARR:
        self.build_from_zarr(found_locations, label_df, cfg)
    elif source_type == DataSourceType.CUTANA:
        self.build_from_cutana(found_locations, label_df, cfg)
    else:
        raise ValueError(f"Cannot build cache for source type: {source_type!r}")

def update(
    self,
    new_labels: dict[str, str],
    cfg: DotMap,
) -> None:
    """Update the cache with newly labeled items (rebuild or append).

    Decides whether to rebuild from scratch (if extraction parameters
    changed) or incrementally append new entries.

    Args:
        new_labels: Dict mapping ``id -> csv_label_string``.
        cfg: Current configuration.

    Raises:
        RuntimeError: If the cache is not populated.
    """
    if not self.is_populated:
        raise RuntimeError("Cannot update an unpopulated cache")

    source_dir = cfg.data_dir
    source_type = _auto_detect_source_type(source_dir)
    if source_type not in (DataSourceType.ZARR, DataSourceType.CUTANA):
        return

    new_label_df = pd.DataFrame(
        [{"id": item_id, "label": label} for item_id, label in new_labels.items()]
    )
    all_labels = pd.read_csv(cfg.label_file)

    if self.needs_rebuild(cfg, label_df=all_labels):
        logger.info("Labeled cache needs rebuild (extraction params or CSV changed)")
        result = validate_labeled_files(source_dir, all_labels, source_type)
        self.clear()
        self.build(result.found_locations, all_labels, source_type, cfg)
    else:
        result = validate_labeled_files(source_dir, new_label_df, source_type)
        if result.found_locations:
            if source_type == DataSourceType.ZARR:
                self.append_from_zarr(result.found_locations, new_label_df, all_labels)
            else:
                self.append_from_cutana(result.found_locations, new_label_df, cfg, all_labels)

    logger.info("Labeled cache updated ({} images)", self.get_cache_info().get("num_images", 0))