cutana_loader

def load_single_cutout(cfg: DotMap, source_id: str) -> np.ndarray | None:
    """Create a single normalised cutout for a streaming source using Cutana.

    Locates *source_id* in the catalogues under ``cfg.prediction_search_dir``
    (via the #506 catalogue index when available) and runs Cutana to produce a
    display-ready cutout under the config's normalisation.

    Args:
        cfg: Full prediction config with normalisation settings.
        source_id: The source identifier stored in the prediction DB.

    Returns:
        HWC uint8 numpy array, or ``None`` if loading fails.
    """
    row_df, cat_file = _find_source_in_catalogues(cfg, cfg.prediction_search_dir, source_id)
    if row_df is None:
        return None
    image = _create_cutout(cfg, row_df)
    if image is None:
        # Found but no cutout (distinct from "not found"): surface source id +
        # method so a detail-view recurrence is traceable.
        logger.warning(
            "Source {} found in {} but produced no cutout under normalisation {} — "
            "tile unreachable, no data, or decode error.",
            source_id,
            os.path.basename(cat_file),
            cfg.normalisation.normalisation_method,
        )
    return image

def load_single_native_cutout(cfg: DotMap, source_id: str) -> np.ndarray | None:
    """Extract one cutout at its *native* pixel size for the detail view.

    Returns a normalisation-independent per-band float cutout the detail view
    re-normalises in-process (so retuning normalisation never re-reads the
    catalogue/tile).  Cutana is run with ``do_only_cutout_extraction`` so it
    skips the resize to a fixed target resolution — the cutout comes back at its
    true native pixel size (e.g. 14x14 for a small Euclid source), letting the
    user inspect real pixels rather than an interpolated resample.

    Args:
        cfg: Full prediction config (locates the catalogues / index DB).
        source_id: The source identifier stored in the prediction DB.

    Returns:
        HWC float32 array (one channel per FITS band) at native resolution, or
        ``None`` if the source can't be located or Cutana produces no cutout.
    """
    row_df, cat_file = _find_source_in_catalogues(cfg, cfg.prediction_search_dir, source_id)
    if row_df is None:
        return None

    native_cfg = copy.deepcopy(cfg)
    native_cfg.fitsbolt_cfg = None
    native_cfg = get_fitsbolt_config(native_cfg)
    try:
        batches = extract_cutouts_direct(row_df, native_cfg, do_only_cutout_extraction=True)
        items = (item for batch in batches for item in _iter_batch_items(batch))
        first = next(items, None)
        if first is None:
            logger.warning(
                "Source {} found in {} but Cutana produced no native cutout — "
                "tile unreachable or no data at that position.",
                source_id,
                os.path.basename(cat_file),
            )
            return None
        _sid, native_img = first
        logger.debug("Native cutout for {}: shape={}", source_id, native_img.shape)
        return native_img
    except Exception as exc:
        logger.warning("Native cutout extraction failed for source {}: {}", source_id, exc)
        return None

def load_sample_cutouts(
    cfg: DotMap, search_dir: str, n_samples: int = 9
) -> list[tuple[str, np.ndarray]]:
    """Create a few sample cutouts from the catalogue for preview.

    Args:
        cfg: Full prediction config with normalisation settings.
        search_dir: Path to the directory containing catalogue files.
        n_samples: Maximum number of sample cutouts to create.

    Returns:
        List of ``(source_id, image_array)`` tuples. May be shorter
        than *n_samples* if some cutouts fail to load.
    """
    cat_files = [
        os.path.join(search_dir, f)
        for f in os.listdir(search_dir)
        if f.lower().endswith((".csv", ".parquet"))
        and _is_cutana_catalogue(Path(os.path.join(search_dir, f)))
    ]
    if not cat_files:
        logger.warning("No Cutana catalogue files found in {}", search_dir)
        return []

    try:
        cat_file = cat_files[0]
        if cat_file.endswith(".parquet"):
            df = pd.read_parquet(cat_file)
        else:
            df = pd.read_csv(cat_file)

        n = min(n_samples, len(df))
        if n == 0:
            return []

        sample_df = df.sample(n=n, random_state=42)

        id_col = _find_id_column(sample_df)

        # Use batch mode for efficiency (single orchestrator call)
        batch_results = _create_batch_cutouts(cfg, sample_df, id_col)

        results = []
        for _, row in sample_df.iterrows():
            sid = str(row[id_col])
            if sid in batch_results:
                results.append((sid, batch_results[sid]))

        return results

    except Exception:
        logger.opt(exception=True).warning("Cutana sample cutouts failed")
        return []

def load_batch_cutouts(
    cfg: DotMap,
    source_ids: list[str],
    *,
    search_dir: str | None = None,
    stale_check=None,
) -> BatchCutoutResult:
    """Create cutouts for multiple sources by scanning catalogues in parallel.

    Dispatches one worker per catalogue file across a thread pool
    (:func:`_scan_one_catalogue`).  Each worker reads its catalogue,
    filters to the requested ids, and builds cutouts for the matching
    rows in a single batch orchestrator call.  The parent thread merges
    results as workers complete and cancels queued workers once every
    requested id has resolved, so fast catalogues that hold the wanted
    ids return before a slow read finishes.  Far more efficient than
    calling ``load_single_cutout`` per source, and the parallel scan
    keeps one slow parquet read from freezing the whole walk (#444).

    Args:
        cfg: Full prediction config with normalisation settings.
        source_ids: List of source identifiers from the prediction DB.
        search_dir: Directory containing catalogue files.  Falls back to
            ``cfg.prediction_search_dir`` when not provided.
        stale_check: Optional zero-arg callable returning ``True`` when
            this request has been superseded by a newer one.  Checked
            once per catalogue so a stale preview bails out without
            driving through all N parquets.  Without this, five rapid
            norm changes caused five overlapping ``load_batch_cutouts``
            runs across the same 54 catalogues, deadlocking pyarrow on
            contested memmap'd reads.

    Returns:
        :class:`BatchCutoutResult` with the loaded cutouts dict plus a
        ``partial`` flag.  ``partial=True`` means the scan did not run to
        completion over every catalogue — either ``stale_check`` fired
        and the remaining workers were cancelled, or one or more
        catalogue reads raised.  In both cases the cutouts dict may be
        missing entries, so consumers that require complete results must
        inspect this flag rather than treat the dict as authoritative.
    """
    if not source_ids:
        return BatchCutoutResult()

    if search_dir is None:
        search_dir = cfg.prediction_search_dir
    if not search_dir or not os.path.isdir(search_dir):
        logger.warning("prediction_search_dir is not set or does not exist: {}", search_dir)
        return BatchCutoutResult()

    # Filter to actual Cutana catalogues (must have SourceID + fits_file_paths)
    cat_files = [
        os.path.join(search_dir, f)
        for f in os.listdir(search_dir)
        if f.lower().endswith((".csv", ".parquet"))
        and _is_cutana_catalogue(Path(os.path.join(search_dir, f)))
    ]
    if not cat_files:
        logger.warning("No catalogue files found in {} for batch cutouts", search_dir)
        return BatchCutoutResult()

    # Walk every discovered catalogue — a single source_ids list may span
    # multiple catalogue files, and picking only the first one would drop
    # matches in the rest (previously caused the setup-screen preview to
    # never show anomalies when anomaly IDs happened to live in a later
    # catalogue file).  Reads run in parallel because the bottleneck is
    # IO, not CPU: when a concurrent prediction subprocess saturates disk
    # bandwidth on one catalogue, that file alone could take minutes —
    # serialising the scan extended UI freezes to ~5 minutes (#444).
    # Parallelism turns worst-case ``n_cat × slow_read`` into
    # ``ceil(n_cat / workers) × slow_read``, and fast catalogues that
    # hold the requested ids return immediately so the user sees results
    # before the slow read finishes.
    wanted = {str(s) for s in source_ids}

    # Catalogue-origin index (#506): when the predictions DB already records
    # which catalogue holds each requested source, scan only those files
    # instead of every catalogue on the page.  The index is keyed to the
    # prediction run, so only the prediction gallery (search_dir ==
    # prediction_search_dir) uses it; other callers — e.g. the training-setup
    # preview, whose search_dir is the training folder — fall through to the
    # full scan.
    cat_files_to_scan = cat_files
    # Origins already recorded for the requested ids; used both to narrow the
    # scan and, at the end, to write back only genuinely-new mappings.
    recorded: dict[str, str] = {}
    use_index = _is_prediction_search_dir(cfg, search_dir)
    if use_index:
        cat_files_set = set(cat_files)
        index_map = _run_on_index_db(
            cfg,
            write=False,
            operation=lambda db: db.get_catalogue_map(list(wanted)),
            default={},
        )
        # Drop ids whose recorded catalogue no longer exists (folder changed):
        # they must be relocated by a full scan, not pinned to a stale path.
        recorded = {sid: cat for sid, cat in index_map.items() if cat in cat_files_set}
        if recorded and recorded.keys() >= wanted:
            known = set(recorded.values())
            cat_files_to_scan = [c for c in cat_files if c in known]
            logger.debug(
                "Catalogue index hit: scanning {}/{} catalogue(s) for {} id(s)",
                len(cat_files_to_scan),
                len(cat_files),
                len(wanted),
            )

    results: dict[str, np.ndarray] = {}
    # source_id -> catalogue path it was found in, recorded back into the DB so
    # the next page load can skip straight to the right catalogue.
    origin: dict[str, str] = {}
    n_cat = len(cat_files_to_scan)
    workers = min(_CATALOGUE_SCAN_WORKERS, n_cat)
    partial = False
    try:
        with ThreadPoolExecutor(max_workers=workers, thread_name_prefix="cutana-cat") as ex:
            # Each worker gets its own immutable snapshot of `wanted` taken
            # at submit time.  The parent thread is the sole writer to
            # `wanted` and `results` (via the `as_completed` loop below),
            # so no locks are needed — workers only read their snapshot
            # and return a result dict by value.  In-flight workers may do
            # redundant matching for ids a peer has already produced, but
            # that's harmless and avoids the lock overhead.
            futures = {
                ex.submit(
                    _scan_one_catalogue,
                    cat_file,
                    i,
                    n_cat,
                    cfg,
                    set(wanted),
                    stale_check,
                ): cat_file
                for i, cat_file in enumerate(cat_files_to_scan, start=1)
            }
            for fut in as_completed(futures):
                if stale_check is not None and stale_check():
                    logger.info("Preview load superseded — cancelling catalogue scan")
                    partial = True
                    for f in futures:
                        f.cancel()
                    break
                try:
                    cat_results = fut.result()
                except CancelledError:
                    # A future cancelled by a previous iteration of this
                    # loop still surfaces here via `as_completed`; its
                    # `.result()` raises CancelledError.  Skip it so we
                    # don't spuriously flip `partial` or log a misleading
                    # "Catalogue read failed".
                    continue
                except Exception as exc:
                    logger.warning("Catalogue read failed: {}", exc)
                    # A skipped catalogue means the cutouts dict cannot be
                    # treated as authoritative — cache builds and warmup must
                    # see partial=True so they don't lock in an empty result.
                    partial = True
                    continue
                if not cat_results:
                    continue
                results.update(cat_results)
                for sid in cat_results:
                    origin[sid] = futures[fut]
                wanted.difference_update(cat_results.keys())
                if not wanted:
                    # All requested ids found — cancel queued workers so
                    # we don't keep paying for unrelated parquet reads.
                    # Workers already running will finish (their reads
                    # may even contribute extra hits, which is harmless).
                    for f in futures:
                        f.cancel()
                    break

        # Best-effort: remember where each source lived so later pages skip
        # the full scan.  Only write origins the index doesn't already hold —
        # on an index hit every origin is already recorded, so re-writing them
        # would open a contended write transaction on every page turn for no
        # gain.  Never let an index write break the preview.
        if use_index:
            new_origins = {sid: cat for sid, cat in origin.items() if recorded.get(sid) != cat}
            if new_origins:
                _run_on_index_db(
                    cfg,
                    write=True,
                    operation=lambda db: db.record_catalogues(new_origins),
                    default=None,
                )
        return BatchCutoutResult(cutouts=results, partial=partial)

    except Exception as exc:
        logger.warning("Batch cutout loading failed: {}", exc)
        # An exception mid-walk leaves the same partial-truth situation as a
        # stale_check bail — flag it so callers don't treat the dict as
        # complete.
        return BatchCutoutResult(cutouts=results, partial=True)