app/scoring.py

"""Hybrid physics + perceptual scoring and candidate ranking."""

from __future__ import annotations

from dataclasses import dataclass
from typing import List, Optional, Sequence

import numpy as np

from film_physics import PiecewiseFilmCurve
from hybrid_loss import HybridFilmLoss, HybridLossBreakdown, _luminance_from_rgb
from app.api_client import SeparationResult
from scoring_policy import ScoringPolicy, DEFAULT_POLICY, loss_kwargs

# WP-13 D4: flat-candidate guard (κ · std(observed)); value from APP_POLICY era
_FLAT_GUARD_KAPPA = 0.05


@dataclass
class RankedCandidate:
    """A separation candidate with hybrid loss breakdown."""

    rank: int
    separation: SeparationResult
    score: HybridLossBreakdown
    candidate_id: str


@dataclass
class RankingResult:
    """Scored, ordered candidates plus flat-guard metadata (WP-13.1 C)."""

    ranked: List[RankedCandidate]
    rejected_count: int = 0

    def __len__(self) -> int:
        return len(self.ranked)

    def __iter__(self):
        return iter(self.ranked)

    def __getitem__(self, index):
        return self.ranked[index]


def _is_flat_pair(
    image_a: np.ndarray,
    image_b: np.ndarray,
    observed_rgb: np.ndarray,
    kappa: float = _FLAT_GUARD_KAPPA,
) -> bool:
    """True if min(std(lum A), std(lum B)) < kappa · std(lum observed)."""
    # Shared Rec.709 helper (hybrid_loss._luminance_from_rgb) — WP-13.1 D
    la = _luminance_from_rgb(image_a)
    lb = _luminance_from_rgb(image_b)
    lo = _luminance_from_rgb(observed_rgb)
    std_o = float(np.std(lo))
    if std_o < 1e-12:
        # Degenerate observation — do not reject (avoid empty ranking)
        return False
    thr = kappa * std_o
    return min(float(np.std(la)), float(np.std(lb))) < thr


def _resolve_policy(
    policy: Optional[ScoringPolicy],
    calibration: str,
    physics_grad_weight: float,
    flat_guard: bool,
) -> ScoringPolicy:
    """Prefer explicit policy; else build from legacy kwargs (compat)."""
    if policy is not None:
        return policy
    # Reconstruct reweights when calibration matches APP (behavior-preserving)
    from scoring_policy import APP_POLICY

    if calibration == APP_POLICY.calibration:
        return ScoringPolicy(
            calibration=calibration,
            physics_grad_weight=physics_grad_weight,
            flat_guard=flat_guard,
            exclusivity_weight=APP_POLICY.exclusivity_weight,
            balance_weight=APP_POLICY.balance_weight,
        )
    return ScoringPolicy(
        calibration=calibration,
        physics_grad_weight=physics_grad_weight,
        flat_guard=flat_guard,
    )


def score_separation(
    observed_log_exposure,
    observed_rgb: np.ndarray,
    image_a_rgb: np.ndarray,
    image_b_rgb: np.ndarray,
    film_curve: PiecewiseFilmCurve,
    physics_weight: float = 1.0,
    perceptual_weight: float = 0.5,
    density: Optional[np.ndarray] = None,
    confidence_mask: Optional[np.ndarray] = None,
    excl_obs: Optional[float] = None,
    calibration: str = "none",
    physics_grad_weight: float = 0.0,
    policy: Optional[ScoringPolicy] = None,
) -> HybridLossBreakdown:
    """Score a candidate pair using ``HybridFilmLoss`` (WP-3 density path when provided).

    Prefer ``policy=APP_POLICY`` from the app path. Legacy kwargs remain for
    benches/tests (None policy → build from kwargs).
    """
    pol = _resolve_policy(policy, calibration, physics_grad_weight, flat_guard=False)
    loss_fn = HybridFilmLoss(
        film_curve=film_curve,
        physics_weight=physics_weight,
        perceptual_weight=perceptual_weight,
        **loss_kwargs(pol),
    )
    return loss_fn.evaluate(
        observed_log_exposure=observed_log_exposure,
        observed_rgb=observed_rgb,
        image_a_rgb=image_a_rgb,
        image_b_rgb=image_b_rgb,
        density=density,
        confidence_mask=confidence_mask,
        excl_obs=excl_obs,
    )


def rank_candidates(
    candidates: Sequence[SeparationResult],
    observed_log_exposure,
    observed_rgb: np.ndarray,
    film_curve: PiecewiseFilmCurve,
    physics_weight: float = 1.0,
    perceptual_weight: float = 0.5,
    density: Optional[np.ndarray] = None,
    confidence_mask: Optional[np.ndarray] = None,
    calibration: str = "none",
    physics_grad_weight: float = 0.0,
    flat_guard: bool = False,
    policy: Optional[ScoringPolicy] = None,
) -> RankingResult:
    """
    Score and rank separation candidates by hybrid loss (lower is better).

    Uses density-space physics + regularizers + K-selection when density/mask provided.
    Returns RankingResult (ranked list + rejected_count).
    """
    pol = _resolve_policy(policy, calibration, physics_grad_weight, flat_guard)
    use_flat_guard = pol.flat_guard

    # Flat-guard filter (opt-in)
    rejected_count = 0
    pool: List[SeparationResult] = list(candidates)
    if use_flat_guard and pool:
        kept: List[SeparationResult] = []
        for c in pool:
            if _is_flat_pair(c.image_a, c.image_b, observed_rgb):
                rejected_count += 1
            else:
                kept.append(c)
        if kept:
            pool = kept
        else:
            # All rejected → unfiltered fallback; keep rejected_count for status
            pool = list(candidates)

    scored: List[RankedCandidate] = []
    # Hoist frame-level K evidence once (Fix 4)
    from hybrid_loss import _lum_split_gradient_overlap
    excl_obs = _lum_split_gradient_overlap(observed_rgb) if observed_rgb is not None else None
    for candidate in pool:
        breakdown = score_separation(
            observed_log_exposure=observed_log_exposure,
            observed_rgb=observed_rgb,
            image_a_rgb=candidate.image_a,
            image_b_rgb=candidate.image_b,
            film_curve=film_curve,
            physics_weight=physics_weight,
            perceptual_weight=perceptual_weight,
            density=density,
            confidence_mask=confidence_mask,
            excl_obs=excl_obs,
            policy=pol,
        )
        scored.append(
            RankedCandidate(
                rank=0,
                separation=candidate,
                score=breakdown,
                candidate_id=candidate.candidate_id,
            )
        )

    scored.sort(key=lambda c: c.score.total_loss)
    for i, item in enumerate(scored, start=1):
        item.rank = i
    return RankingResult(ranked=scored, rejected_count=rejected_count)


def merge_new_candidates_into_ranking(
    base: RankingResult,
    new_candidates: Sequence[SeparationResult],
    observed_log_exposure,
    observed_rgb: np.ndarray,
    film_curve: PiecewiseFilmCurve,
    physics_weight: float = 1.0,
    perceptual_weight: float = 0.5,
    density: Optional[np.ndarray] = None,
    confidence_mask: Optional[np.ndarray] = None,
    policy: Optional[ScoringPolicy] = None,
    calibration: str = "none",
    physics_grad_weight: float = 0.0,
) -> RankingResult:
    """Score ONLY new candidates and merge by total_loss (WP-14.1 P0-rank-once).

    Flat guard applies to new entries using the same policy as rank_candidates.
    Does not re-score base.ranked.
    """
    if policy is not None:
        pol = policy
        use_flat = bool(policy.flat_guard)
    else:
        pol = _resolve_policy(None, calibration, physics_grad_weight, flat_guard=False)
        use_flat = False

    rejected = int(base.rejected_count)
    from hybrid_loss import _lum_split_gradient_overlap

    excl_obs = (
        _lum_split_gradient_overlap(observed_rgb) if observed_rgb is not None else None
    )

    pool_new: List[SeparationResult] = list(new_candidates)
    if use_flat and pool_new:
        kept: List[SeparationResult] = []
        for c in pool_new:
            if _is_flat_pair(c.image_a, c.image_b, observed_rgb):
                rejected += 1
            else:
                kept.append(c)
        # Same fallback as rank_candidates: if all new flat and no base, keep unfiltered
        if kept:
            pool_new = kept
        elif base.ranked:
            pool_new = []  # base covers the ranking
        else:
            pool_new = list(new_candidates)

    scored_new: List[RankedCandidate] = []
    for candidate in pool_new:
        breakdown = score_separation(
            observed_log_exposure=observed_log_exposure,
            observed_rgb=observed_rgb,
            image_a_rgb=candidate.image_a,
            image_b_rgb=candidate.image_b,
            film_curve=film_curve,
            physics_weight=physics_weight,
            perceptual_weight=perceptual_weight,
            density=density,
            confidence_mask=confidence_mask,
            excl_obs=excl_obs,
            policy=pol,
        )
        scored_new.append(
            RankedCandidate(
                rank=0,
                separation=candidate,
                score=breakdown,
                candidate_id=candidate.candidate_id,
            )
        )

    combined = list(base.ranked) + scored_new
    combined.sort(key=lambda c: c.score.total_loss)
    for i, item in enumerate(combined, start=1):
        item.rank = i
    return RankingResult(ranked=combined, rejected_count=rejected)


def rank_then_merge_asymmetric(
    base_candidates: Sequence[SeparationResult],
    preprocessed,
    *,
    film_curve: PiecewiseFilmCurve,
    physics_weight: float = 1.0,
    perceptual_weight: float = 0.5,
    policy: Optional[ScoringPolicy] = None,
    enable_asymmetric: bool = False,
    anchor: str = "auto",
    complete: bool = False,
    status_notes: Optional[List[str]] = None,
    inpaint_fn=None,
    debug: bool = False,
) -> tuple:
    """Rank base pool ONCE; optionally score-only-merge asym candidates (WP-14.1).

    Shared by app/main.py and synth/real_protocol.py. Never re-scores the base pool.

    Returns:
        (RankingResult, api_contacted: bool) — api_contacted drives consent UI.
    """
    density = getattr(preprocessed, "density", None)
    confidence_mask = getattr(preprocessed, "confidence_mask", None)
    observed_rgb = getattr(preprocessed, "rgb", None)
    observed_log_exposure = getattr(preprocessed, "log_exposure", None)

    base_ranked = rank_candidates(
        candidates=base_candidates,
        observed_log_exposure=observed_log_exposure,
        observed_rgb=observed_rgb,
        film_curve=film_curve,
        physics_weight=physics_weight,
        perceptual_weight=perceptual_weight,
        density=density,
        confidence_mask=confidence_mask,
        policy=policy,
    )
    if not enable_asymmetric or len(base_ranked) == 0:
        return base_ranked, False

    notes = status_notes if status_notes is not None else []
    api_contacted = False
    try:
        from app.asymmetric import append_asymmetric

        scratch: list = []
        run = append_asymmetric(
            scratch,
            preprocessed,
            base_ranked[0].separation,
            anchor=anchor,
            complete=complete,
            status_notes=notes,
            inpaint_fn=inpaint_fn,
            debug=debug,
        )
        api_contacted = bool(run.api_contacted)
        asym = run.candidates
    except Exception as exc:
        notes.append(f"Asymmetric recovery failed: {exc}")
        return base_ranked, False

    if not asym:
        return base_ranked, api_contacted

    merged = merge_new_candidates_into_ranking(
        base_ranked,
        asym,
        observed_log_exposure=observed_log_exposure,
        observed_rgb=observed_rgb,
        film_curve=film_curve,
        physics_weight=physics_weight,
        perceptual_weight=perceptual_weight,
        density=density,
        confidence_mask=confidence_mask,
        policy=policy,
    )
    return merged, api_contacted


def format_score_summary(
    ranked: RankedCandidate,
    num_candidates: int,
    physics_weight: float,
    perceptual_weight: float,
) -> str:
    """Human-readable score summary for the UI."""
    s = ranked.score
    return (
        f"**Best candidate:** #{ranked.rank} of {num_candidates} "
        f"(`{ranked.candidate_id}`)\n\n"
        f"**Hybrid loss:** {s.total_loss:.6f} _(lower is better)_\n\n"
        f"- Physics MSE: {s.physics_loss:.6f} (weight {physics_weight:.1f})\n"
        f"- LPIPS: {s.perceptual_loss:.6f} (weight {perceptual_weight:.1f})"
    )


def format_ranking_table(ranked_list) -> str:
    """Compact markdown table of all candidate scores (includes K-selection)."""
    items = ranked_list.ranked if isinstance(ranked_list, RankingResult) else ranked_list
    lines = ["| Rank | ID | Hybrid | Physics | LPIPS | K-sel | Method |", "|---|---|---|---|---|---|---|"]
    for item in items:
        s = item.score
        ksel = s.k_selection_score
        lines.append(
            f"| {item.rank} | {item.candidate_id} | {s.total_loss:.4f} | "
            f"{s.physics_loss:.4f} | {s.perceptual_loss:.4f} | "
            f"{ksel:.2f} | {item.separation.method} |"
        )
    return "\n".join(lines)