cutana_source

def build_cutana_orchestrator_config(
    catalogue_path: str,
    cfg: DotMap,
    *,
    sub_catalogue_df: pd.DataFrame | None = None,
    do_only_cutout_extraction: bool = False,
    prune_unused_bands: bool = False,
) -> DotMap:
    """Build a Cutana orchestrator config from AnomalyMatch config.

    The orchestrator is configured with **identity channel weights** — one
    entry per resolved extension, each a 1.0-weighted one-hot vector of
    length ``n_extensions``.  This turns Cutana's internal
    ``combine_channels`` step into a no-op, so the orchestrator output is
    per-band normalised data with ``n_extensions`` channels.  The user's
    ``channel_combination`` matrix is then applied post-orchestrator via
    :func:`apply_channel_combination_to_cutana_image`, which keeps
    combination logic in one place shared with the read-from-cache and
    prediction-subprocess paths.

    This is the single source of truth for orchestrator config building.
    ``CutanaSource``, ``LabeledDataCache``, and ``prediction_process_cutana``
    all delegate here.

    Args:
        catalogue_path: Path to the sub-catalogue file.  Used for filter-name
            resolution (reads the first row's ``fits_file_paths``) unless
            *sub_catalogue_df* is provided; always assigned to
            ``cutana_cfg.source_catalogue`` for orchestrator callers.
        cfg: AnomalyMatch configuration.
        sub_catalogue_df: Optional in-memory catalogue override.  When given,
            filter names are read from its first row instead of hitting disk —
            the path to parquet roundtrip that the direct-cutout caller would
            otherwise need.
        do_only_cutout_extraction: When ``True``, Cutana skips channel
            combination, normalisation, and the resize to ``target_resolution``,
            returning each cutout at its true native pixel size (the detail
            view's native path).
        prune_unused_bands: When ``True``, skip reading FITS extensions whose
            ``channel_combination`` column is entirely zero — they contribute
            nothing to any output channel, so reading them is wasted I/O.  Only
            safe for ephemeral streaming (prediction, unlabeled training): the
            labeled cache and preview paths must keep **all** bands so the user
            can retune the combination without re-extracting, so they leave this
            ``False``.  The downstream combine matmul drops the matching all-zero
            matrix columns, making the pruned result identical to the full one.

    Returns:
        Cutana configuration object.

    Raises:
        ValueError: If ``cfg.fitsbolt_cfg`` is None — without it the user's
            normalisation method would not reach Cutana and it would silently
            apply a linear (CONVERSION_ONLY) stretch.
        RuntimeError: If the installed Cutana release does not expose the
            ``skip_catalogue_validation`` flag.
    """
    # Broken-invariant guard runs first, before any catalogue/disk work in
    # _resolve_extension_names.  Without an external fitsbolt config, Cutana
    # falls back to its own default normalisation (CONVERSION_ONLY / linear),
    # silently discarding the user's ASINH/LOG/ZSCALE choice — the model would
    # then be fed linearly-stretched cutouts it was never trained on.  Every
    # production caller populates cfg.fitsbolt_cfg first (training via
    # get_fitsbolt_config, prediction via sync_normalisation_from_checkpoint,
    # the labeled cache via _build_raw_extraction_cfg), so a None here is a bug:
    # fail hard rather than ship a linear-normalised run.
    if cfg.fitsbolt_cfg is None:
        raise ValueError(
            "cfg.fitsbolt_cfg is None when building the Cutana orchestrator config; "
            "the normalisation settings would not reach Cutana and it would fall back "
            "to a linear (CONVERSION_ONLY) stretch. Populate cfg.fitsbolt_cfg from the "
            "model checkpoint (prediction) or from cfg.normalisation via "
            "get_fitsbolt_config (training) before streaming cutouts."
        )

    # A non-None fitsbolt_cfg can still carry the wrong method: the stretch
    # reaching Cutana comes entirely from external_fitsbolt_cfg (built below
    # from cfg.fitsbolt_cfg), while cfg.normalisation carries the intended
    # method. Both production paths keep them in lock-step — training rebuilds
    # fitsbolt_cfg from cfg.normalisation via get_fitsbolt_config, prediction
    # writes both from the checkpoint via sync_normalisation_from_checkpoint —
    # so a disagreement means a caller populated one without the other. Fail
    # hard rather than silently stretch cutouts with a method the model was
    # never trained on (the None guard above only catches the missing case).
    fitsbolt_method = NormalisationMethod(cfg.fitsbolt_cfg["normalisation_method"])
    if fitsbolt_method != cfg.normalisation.normalisation_method:
        raise ValueError(
            f"cfg.fitsbolt_cfg normalisation_method ({fitsbolt_method.name}) disagrees with "
            f"cfg.normalisation.normalisation_method ({cfg.normalisation.normalisation_method.name}); "
            "the cutouts would be stretched with a method the model was not trained on. "
            "Rebuild cfg.fitsbolt_cfg from cfg.normalisation via get_fitsbolt_config (training) "
            "or re-sync from the checkpoint via sync_normalisation_from_checkpoint (prediction)."
        )

    cutana_cfg = cutana.get_default_config()
    cutana_cfg.source_catalogue = catalogue_path
    cutana_cfg.target_resolution = cfg.normalisation.image_size[0]
    cutana_cfg.padding_factor = cfg.normalisation.cutout_padding_factor
    # Cutana's own data_type controls the final dtype conversion —
    # external_fitsbolt_cfg.output_dtype is used only for fitsbolt's
    # per-channel normalisation and is not propagated to the batch output.
    cutana_cfg.data_type = np.dtype(cfg.normalisation.output_dtype).name

    extension_names, has_integer_indices = _resolve_extension_names(
        cfg, catalogue_path=catalogue_path, sub_catalogue_df=sub_catalogue_df
    )
    n_extensions_full = len(extension_names)
    band_keep_mask: np.ndarray | None = None
    if prune_unused_bands:
        extension_names, band_keep_mask = _prune_unused_extensions(cfg, extension_names)
    n_extensions = len(extension_names)

    # Multi-file catalogues (one FITS per band) store each file under its own
    # PRIMARY HDU, so fits_extensions must be ["PRIMARY"] when filter names
    # came from file-name inspection.  When the user supplied filter name
    # strings directly, trust those.
    if has_integer_indices or cfg.normalisation.fits_extension is None:
        cutana_cfg.fits_extensions = ["PRIMARY"]
    else:
        cutana_cfg.fits_extensions = extension_names

    cutana_cfg.selected_extensions = [{"name": name, "ext": "PRIMARY"} for name in extension_names]

    # Identity channel weights: each extension gets its own output channel,
    # fully isolated.  Cutana's combine_channels becomes a no-op; all
    # per-band data passes through to apply_normalisation, then to the cache.
    cutana_cfg.channel_weights = {
        name: [1.0 if k == j else 0.0 for k in range(n_extensions)]
        for j, name in enumerate(extension_names)
    }

    cutana_cfg.apply_flux_conversion = cfg.normalisation.apply_flux_conversion

    # Build a Cutana-specific fitsbolt config that matches the
    # per-band pass-through shape: n_expected=n_output=n_extensions,
    # no channel_combination.  Per-channel normalisation params must
    # be resized to n_extensions so ASINH scales/clips align with
    # bands; they live under ``normalisation.asinh_{scale,clip}`` in
    # the fitsbolt config (the top-level ``norm_asinh_*`` names only
    # exist on AnomalyMatch's ``cfg.normalisation``).  When the user
    # picked n_output_channels=3 for 4-band Euclid data, the
    # 3-element list silently slipped past a mismatched-key lookup
    # here and the Cutana worker produced zero cutouts.
    fitsbolt = DotMap(cfg.fitsbolt_cfg.toDict(), _dynamic=False)
    fitsbolt.channel_combination = None
    fitsbolt.n_output_channels = n_extensions
    fitsbolt.n_expected_channels = n_extensions
    for attr in ("asinh_scale", "asinh_clip"):
        val = fitsbolt.normalisation.get(attr)
        if isinstance(val, (list, tuple)):
            # Resize at the *full* band count so each surviving band keeps the
            # exact param it would get in an unpruned run, then drop the pruned
            # bands' entries by the same mask.  Truncating to the pruned count
            # instead would shift later bands onto earlier params for a
            # non-prefix prune, normalising them differently than the full run.
            resized = _resize_param_list(val, n_extensions_full)
            if band_keep_mask is not None:
                resized = [p for p, keep in zip(resized, band_keep_mask) if keep]
            fitsbolt.normalisation[attr] = resized
    # Cutana's apply_normalisation does ``if "crop_for_maximum_value"
    # in external_cfg.normalisation: crop_value = ...; crop_value[0]``
    # — an ``in`` check rather than a None check, so a key set to
    # ``None`` (the fitsbolt default when the user hasn't enabled
    # cropping) makes Cutana attempt ``None[0]`` and crash the whole
    # batch.  Drop the key explicitly when its value is None.
    if fitsbolt.normalisation.get("crop_for_maximum_value") is None:
        fitsbolt.normalisation.pop("crop_for_maximum_value", None)
    cutana_cfg.external_fitsbolt_cfg = fitsbolt

    # Catalogues reaching this point have already been validated by
    # AnomalyMatch (``source_validation._validate_cutana`` reads them
    # column-by-column, and ``_is_cutana_catalogue`` gate-keeps both
    # source selection and preview loading).  Cutana's per-row
    # re-validation is pure overhead on huge catalogues (29 M-row Q1
    # search ≈ minute), so we opt out.  The flag must be present in the
    # cutana release that ships with the project — fail loudly on an
    # older install instead of silently setting an ignored attribute.
    if "skip_catalogue_validation" not in cutana_cfg:
        raise RuntimeError(
            "Installed cutana release does not expose `skip_catalogue_validation`; "
            "upgrade cutana to a release that supports it."
        )
    cutana_cfg.skip_catalogue_validation = True

    # Native (unresized) extraction for the detail view: Cutana skips channel
    # combination, normalisation, and the resize to target_resolution, returning
    # the cutout at its true native pixel size.  The detail screen applies
    # normalisation + channel combination in-process afterwards, so the user
    # inspects real pixels rather than a fixed-resolution resample.
    if do_only_cutout_extraction:
        if "do_only_cutout_extraction" not in cutana_cfg:
            raise RuntimeError(
                "Installed cutana release does not expose `do_only_cutout_extraction`; "
                "upgrade cutana to a release that supports native cutout extraction."
            )
        cutana_cfg.do_only_cutout_extraction = True

    return cutana_cfg

def extract_cutouts_direct(
    sub_catalogue_df: pd.DataFrame, cfg: DotMap, *, do_only_cutout_extraction: bool = False
) -> list[dict]:
    """Run Cutana's in-process ``create_cutouts_direct`` on a filtered catalogue.

    Intended for small batches (< :data:`DIRECT_CUTOUT_MAX_SOURCES` rows).
    Avoids the subprocess spawn + IPC overhead of
    ``StreamingOrchestrator.next_batch`` so label-cache builds and UI previews
    return in seconds instead of tens of seconds.

    Passes *sub_catalogue_df* directly into
    :func:`build_cutana_orchestrator_config` so filter-name resolution stays
    in-memory; ``create_cutouts_direct`` itself ignores
    ``cutana_cfg.source_catalogue``.

    Args:
        sub_catalogue_df: Filtered catalogue with Cutana columns.
        cfg: AnomalyMatch configuration.
        do_only_cutout_extraction: When ``True``, Cutana returns each cutout at
            its native pixel size with no channel combination / normalisation /
            resize — the detail view's true-native path. ``batch["cutouts"]`` is
            then a list of per-source native arrays (``_iter_batch_items`` handles
            both that and the usual stacked ndarray).

    Returns:
        List of Cutana batch result dicts (one per FITS file set), each with
        ``"cutouts"``, ``"metadata"``, ``"wcs"``, ``"channel_names"`` keys.
        The list matches the contract of iterating
        ``StreamingOrchestrator.next_batch``.
    """
    cutana_cfg = build_cutana_orchestrator_config(
        catalogue_path="",
        cfg=cfg,
        sub_catalogue_df=sub_catalogue_df,
        do_only_cutout_extraction=do_only_cutout_extraction,
    )
    # This wrapper is invoked once per catalogue (label-cache build and UI
    # preview both loop over catalogues), so Cutana's per-FITS-set heartbeat —
    # always a small, repeated count — is pure noise here; opt out when the
    # installed Cutana exposes the flag.  The guard keeps us compatible with
    # older Cutana builds (CI installs from Cutana's develop branch, which may
    # lag this flag); progress is conveyed by our own per-catalogue logging.
    if _CREATE_CUTOUTS_SUPPORTS_LOG_FLAG:
        return cutana.create_cutouts_direct(sub_catalogue_df, cutana_cfg, log_set_progress=False)
    return cutana.create_cutouts_direct(sub_catalogue_df, cutana_cfg)

def resolve_catalogue_filter_names(catalogue_path: str) -> list[str]:
    """Resolve filter/band names from a single Cutana catalogue file or folder.

    Reads ``fits_file_paths`` from the first source and extracts filter
    names using Cutana's naming conventions.  When names are unrecognised,
    falls back to positional ``Band_0``, ``Band_1``, ... labels.

    Args:
        catalogue_path: Path to a catalogue file or directory containing
            catalogues.

    Returns:
        List of filter name strings, e.g. ``["VIS", "NIR-H", "NIR-J", "NIR-Y"]``.

    Raises:
        FileNotFoundError: If *catalogue_path* is a directory with no
            catalogue files.
    """
    if os.path.isfile(catalogue_path):
        first_file = catalogue_path
    else:
        entries = sorted(os.listdir(catalogue_path))
        candidates = [
            os.path.join(catalogue_path, fname)
            for fname in entries
            if fname.endswith((".parquet", ".csv"))
        ]
        if not candidates:
            raise FileNotFoundError(f"No catalogue files found in {catalogue_path}")
        first_file = candidates[0]

    if first_file.endswith(".parquet"):
        df = pd.read_parquet(first_file, columns=["fits_file_paths"]).head(1)
    else:
        df = pd.read_csv(first_file, usecols=["fits_file_paths"], nrows=1)

    return _filter_names_from_fits_paths_entry(df["fits_file_paths"].iloc[0])

def first_cutana_catalogue(folder: str) -> str | None:
    """Return the path of the first validated Cutana catalogue in *folder*.

    Reading the first ``.csv``/``.parquet`` blindly would trip over a stray
    non-catalogue file that sorts first (e.g. a ``labeled_data.csv`` label file
    living next to the catalogues).

    Args:
        folder: Directory to search.

    Returns:
        Path to the first file that passes ``_is_cutana_catalogue``, or ``None``
        when *folder* isn't a readable directory or holds no catalogue.
    """
    # Guard falsy paths explicitly: os.listdir(None) lists the cwd rather than
    # raising, which would then walk the wrong directory.
    if not folder:
        return None
    try:
        entries = sorted(os.listdir(folder))
    except OSError as exc:
        logger.warning("Cannot list {} to look for Cutana catalogues: {}", folder, exc)
        return None
    for fname in entries:
        fpath = Path(folder) / fname
        if fname.endswith((".parquet", ".csv")) and _is_cutana_catalogue(fpath):
            return str(fpath)
    return None

def detect_cutana_filter_names(folder: str, cfg: DotMap) -> list[str]:
    """Return the bands a Cutana source in *folder* will be extracted with.

    Resolves through :func:`_resolve_extension_names`, the resolver the
    orchestrator config and the labeled cache use, so the channel-combination
    editor is sized to exactly the bands extraction delivers.  A separate
    catalogue-only parse used to ignore an explicit ``fits_extension`` and
    returned nothing on a parse error, leaving the editor at its old width while
    extraction produced a different band count.

    Args:
        folder: Directory containing Cutana catalogue files.
        cfg: Configuration whose ``normalisation.fits_extension`` selects bands.

    Returns:
        Ordered band names, or an empty list when *folder* holds no Cutana
        catalogue.
    """
    catalogue_path = first_cutana_catalogue(folder)
    if catalogue_path is None:
        logger.warning("No Cutana catalogue found in {}; cannot resolve its bands", folder)
        return []
    extension_names, _ = _resolve_extension_names(cfg, catalogue_path=catalogue_path)
    return extension_names