Source code for spindoctor.nav_model.stars.nav_model_stars

"""Real-scene star NavModel.

Orchestrates the catalog-reduction, conflict-marking, predicted-SNR, and
smear-aware PSF helpers so the NavModel ABC can emit one ``STAR``
``NavFeature`` per detectable catalog star.  Heavy lifting lives in
sibling modules; this file is the entry point the orchestrator's
registry sees.
"""

from __future__ import annotations

import math
from typing import TYPE_CHECKING, Any

import numpy as np
from oops import Observation

from spindoctor.annotation import (
    TEXTINFO_BOTTOM,
    TEXTINFO_BOTTOM_LEFT,
    TEXTINFO_BOTTOM_RIGHT,
    TEXTINFO_LEFT,
    TEXTINFO_RIGHT,
    TEXTINFO_TOP,
    TEXTINFO_TOP_LEFT,
    TEXTINFO_TOP_RIGHT,
    Annotation,
    Annotations,
    AnnotationTextInfo,
    TextLocInfo,
)
from spindoctor.config import Config
from spindoctor.feature.constants import MIN_ANISOTROPIC_SMEAR_PX
from spindoctor.feature.feature import NavFeature, NavReliabilityBreakdown
from spindoctor.feature.feature_type import NavFeatureType
from spindoctor.feature.flags import StarFlags
from spindoctor.feature.geometry import StarGeometry
from spindoctor.nav_model.nav_model import NavModel
from spindoctor.nav_model.stars.catalog import reduce_catalogs
from spindoctor.nav_model.stars.conflicts import mark_body_and_ring_conflicts
from spindoctor.nav_model.stars.predicted_snr import (
    SCLASS_TO_B_MINUS_V,
    psf_sigma_px,
)
from spindoctor.nav_model.stars.smeared_psf import compute_smear_vector_px, smear_length_px
from spindoctor.support.flux import clean_sclass
from spindoctor.support.image import draw_rect
from spindoctor.support.time import now_dt

if TYPE_CHECKING:  # pragma: no cover - typing-only import
    from spindoctor.nav_orchestrator.nav_context import NavContext
    from spindoctor.support.types import MutableStar

from spindoctor.support.filters import NavFilterKind, NavFilterSpec

__all__ = [
    'SCLASS_TO_B_MINUS_V',
    'NavModelStars',
]

# Effective-SNR constants for the magnitude-margin reliability shape.
#
# The star gate is purely magnitude based: the catalog reduction and this
# model drop any star fainter than ``obs.star_max_usable_vmag()``.  The
# CRLB-covariance and reliability helpers still want an SNR-like quantity,
# so we synthesise one from how far below the limiting magnitude the star
# sits.  A star exactly at the limit gets ``snr_eff == SNR_REF``; every
# extra magnitude of headroom multiplies the effective SNR by one Pogson
# ratio (``2.512``), matching the flux ratio per magnitude.
#
# ``SNR_REF`` is set to 8.0 — just below the reliability sigmoid centre
# (snr=10) so a star at the very edge of usability lands near the steep
# part of the curve (reliability ~0.34) rather than saturating it, while a
# star a few magnitudes brighter saturates the curve toward 1.0.
# ``SNR_FLOOR`` keeps the synthesised SNR strictly positive so the
# covariance never degenerates to the zero-SNR huge-variance branch.
SNR_REF: float = 8.0
SNR_FLOOR: float = 0.1






def _resolve_mag_offset(obs: Observation) -> float:
    """Return the per-camera-per-filter magnitude offset for a navigation.

    Reads ``obs.inst_config.mag_offset.{fallback_combo, mag_offset_table}``
    populated by ``ObsInst.from_file``.  When no per-instrument
    ``mag_offset`` block is configured (test fixtures, simulated obs
    that skip the ``from_file`` path), the offset is ``0.0`` — the
    legacy hard-coded value.

    The resolved offset is read from
    ``mag_offset_table[fallback_combo]['default']``; the per-color-bin
    sub-keys are reserved for phase-10 calibration.
    """
    inst_config = getattr(obs, 'inst_config', None)
    if inst_config is None:
        return 0.0
    block = inst_config.get('mag_offset')
    if not block:
        return 0.0
    fallback = block.get('fallback_combo')
    table = block.get('mag_offset_table') or {}
    entry = table.get(fallback) if fallback is not None else None
    if entry is None and table:
        entry = next(iter(table.values()), None)
    if entry is None:
        return 0.0
    if isinstance(entry, (int, float)) and not isinstance(entry, bool):
        return float(entry)
    if isinstance(entry, dict):
        default = entry.get('default', 0.0)
        return float(default) if isinstance(default, (int, float)) else 0.0
    return 0.0


def _compute_smear_for_obs(obs: Observation) -> tuple[float, float]:
    """Best-effort smear vector lookup that tolerates obs API quirks.

    Pulls smear from the SPICE bracket via
    ``compute_smear_vector_px`` when the obs supports it.  Snapshots
    that don't have a bracket-capable boresight (simulated obs, partial
    test fixtures) report a zero smear vector — the SNR formula still
    applies and the star covariance falls back to the isotropic case.

    Parameters:
        obs: Observation snapshot.

    Returns:
        ``(my, mx)`` smear vector in pixels.  Returns ``(0.0, 0.0)`` for
        obs stand-ins that don't expose a SPICE-bracketable boresight.
    """
    try:
        return compute_smear_vector_px(obs)
    except (AttributeError, NotImplementedError):
        return (0.0, 0.0)


def _safe_mask_lookup(
    mask: Any,
    v: float,
    u: float,
) -> bool:
    """Return ``mask[round(v), round(u)]`` clamped to bounds, defaulting False."""
    if mask is None:
        return False
    arr = np.asarray(mask)
    if arr.ndim != 2 or arr.size == 0:
        return False
    rows, cols = arr.shape
    vi = int(np.clip(round(v), 0, rows - 1))
    ui = int(np.clip(round(u), 0, cols - 1))
    return bool(arr[vi, ui])


def _crlb_covariance(
    *,
    snr: float,
    sigma_psf: float,
    move_v: float,
    move_u: float,
) -> np.ndarray:
    """Return the 2x2 anisotropic CRLB centroid covariance for a star.

    Implements the formula from the design's STAR section:

    ::

        sigma_along  = sqrt(L^2 / 12 + sigma_PSF^2) / sqrt(SNR)
        sigma_across = sigma_PSF / sqrt(SNR)

    rotated so the major axis aligns with the smear vector ``(my, mx)``.
    Below ``MIN_ANISOTROPIC_SMEAR_PX``, the covariance collapses to an
    isotropic ``(sigma_PSF / sqrt(SNR))^2 * I``.

    Parameters:
        snr: Predicted integrated SNR.
        sigma_psf: Per-pixel PSF Gaussian sigma in pixels.
        move_v: Smear vector V component (pixels).
        move_u: Smear vector U component (pixels).

    Returns:
        2x2 numpy float array with rows / columns in (v, u) order.
    """
    if snr <= 0.0:
        return np.eye(2, dtype=np.float64) * 1e6
    if sigma_psf <= 0.0:
        raise ValueError(f'sigma_psf must be > 0; got {sigma_psf!r}')
    smear = math.hypot(move_v, move_u)
    s_across_sq = (sigma_psf * sigma_psf) / snr
    if smear < MIN_ANISOTROPIC_SMEAR_PX:
        return np.eye(2, dtype=np.float64) * s_across_sq
    s_along_sq = (smear * smear / 12.0 + sigma_psf * sigma_psf) / snr
    cos_t = move_v / smear
    sin_t = move_u / smear
    rot = np.array([[cos_t, -sin_t], [sin_t, cos_t]], dtype=np.float64)
    diag = np.diag([s_along_sq, s_across_sq])
    cov = rot @ diag @ rot.T
    return 0.5 * (cov + cov.T)


def _snr_reason_score(snr: float, min_snr: float) -> float:
    """Map an effective SNR onto the reliability-breakdown ``predicted_snr`` field.

    The ``NavReliabilityBreakdown.predicted_snr`` is a [0, 1] contribution
    score, not the raw SNR.  ``snr`` here is the magnitude-margin-derived
    effective SNR (``SNR_REF * 2.512 ** (mag_limit - vmag)``), not a
    photometric DN SNR.  This helper folds the configured floor in: when
    ``min_snr > 0``, the score is ``snr / min_snr`` capped at 1; when no
    floor is configured, an SNR of 50 saturates the score (the same centre
    the reliability sigmoid uses).

    Parameters:
        snr: Effective magnitude-margin SNR.
        min_snr: Configured floor (``stars.min_predicted_snr``); 0
            when unset.

    Returns:
        Contribution score in ``[0, 1]``.
    """
    if min_snr > 0.0:
        return float(min(1.0, snr / min_snr))
    return float(min(1.0, snr / 50.0))


def _reliability_from_snr(
    *,
    snr: float,
    in_body: bool,
    in_ring: bool,
    in_saturation: bool,
) -> float:
    """Return the [0, 1] reliability for a star.

    Sigmoid-of-SNR core multiplied by hard zero terms for occlusion and
    saturation.  Matches the design's STAR formula:
    ``sigmoid(alpha_0 + alpha_1*(snr - threshold))`` with the
    ``not_in_body_silhouette`` / ``not_in_saturation_or_cosmic`` factors
    applied multiplicatively.  Excessive-smear stars are excluded from
    the feature list before this helper runs.

    Parameters:
        snr: Predicted SNR.
        in_body: True if the star is occluded by a body silhouette.
        in_ring: True if the star is occluded by a ring annulus.
        in_saturation: True if the predicted pixel is saturated or
            tagged as a cosmic-ray hit.

    Returns:
        Reliability scalar in ``[0, 1]``.
    """
    if in_body or in_ring or in_saturation:
        return 0.0
    # Sigmoid centred near snr=10 with a span of ~5 — stars at snr=20 are
    # saturated to ~1, stars at snr=5 are around 0.27.  Coefficients are
    # the calibration starting point; phase-5 fitting will refine them.
    z = 0.2 * (snr - 10.0)
    return float(1.0 / (1.0 + math.exp(-z)))


def _star_feature_id(star: MutableStar) -> str:
    """Return the ``feature_id`` for ``star``: ``star:<catalog>:<id>``."""
    cat = (star.catalog_name or 'unknown').upper()
    uid = star.unique_number if star.unique_number is not None else star.pretty_name
    return f'star:{cat}:{uid}'


def _star_short_info(star: MutableStar) -> str:
    """Return a one-line text summary of a star, suitable for INFO logging.

    Mirrors the shape of the legacy ``NavModelStars._star_short_info`` line so
    the per-image log keeps a familiar listing that operators can grep.
    """
    jb = getattr(star, 'johnson_mag_b', None) or 0.0
    jv = getattr(star, 'johnson_mag_v', None) or 0.0
    temp = getattr(star, 'temperature', None) or 0.0
    return (
        f'Star {star.catalog_name:>6s}/{star.pretty_name:>9s} '
        f'U {star.u:9.3f}+/-{abs(star.move_u):7.3f} '
        f'V {star.v:9.3f}+/-{abs(star.move_v):7.3f} '
        f'VMAG {(-1.0 if star.vmag is None else star.vmag):6.3f} '
        f'JBMAG {jb:6.3f} '
        f'JVMAG {jv:6.3f} '
        f'SCLASS {clean_sclass(star.spectral_class or ""):>3s} '
        f'TEMP {temp:6.0f} '
        f'CONFLICT {star.conflicts}'
    )


def _star_summary(star: MutableStar) -> dict[str, Any]:
    """Return a compact JSON-friendly summary of a star (for metadata)."""
    return {
        'catalog_name': star.catalog_name,
        'unique_number': star.unique_number,
        'pretty_name': star.pretty_name,
        'ra_deg': float(np.degrees(star.ra_pm)),
        'dec_deg': float(np.degrees(star.dec_pm)),
        'vmag': star.vmag,
        'u': star.u,
        'v': star.v,
        'move_u': star.move_u,
        'move_v': star.move_v,
        'spectral_class': star.spectral_class,
        'conflicts': star.conflicts,
    }


def _star_label(
    *,
    star: MutableStar,
    config: Any,
    v_int: int,
    u_int: int,
    label_margin: int,
) -> AnnotationTextInfo:
    """Return the label info for one star (8 candidate placements)."""
    sclass = clean_sclass(star.spectral_class or '')
    line1 = star.pretty_name[:10]
    line2 = f'{star.vmag:.3f} {sclass}' if star.vmag is not None else sclass
    text_loc: list[TextLocInfo] = [
        TextLocInfo(TEXTINFO_BOTTOM, v_int + label_margin, u_int),
        TextLocInfo(TEXTINFO_TOP, v_int - label_margin, u_int),
        TextLocInfo(TEXTINFO_LEFT, v_int, u_int - label_margin),
        TextLocInfo(TEXTINFO_RIGHT, v_int, u_int + label_margin),
        TextLocInfo(TEXTINFO_TOP_LEFT, v_int - label_margin, u_int - label_margin),
        TextLocInfo(TEXTINFO_TOP_RIGHT, v_int - label_margin, u_int + label_margin),
        TextLocInfo(TEXTINFO_BOTTOM_LEFT, v_int + label_margin, u_int - label_margin),
        TextLocInfo(TEXTINFO_BOTTOM_RIGHT, v_int + label_margin, u_int + label_margin),
    ]
    return AnnotationTextInfo(
        f'{line1}\n{line2}',
        ref_vu=(v_int, u_int),
        text_loc=text_loc,
        font=config.label_font,
        font_size=config.label_font_size,
        color=config.label_font_color,
    )