"""WP-13 D3/D4: ranking reweight + flat-candidate guard (independent-reference asserts)."""

from __future__ import annotations

import numpy as np
import torch
import pytest

from densitometry import VALID
from film_physics import get_film_curve
from hybrid_loss import HybridFilmLoss
from app.api_client import SeparationResult
from app.scoring import rank_candidates, score_separation, RankingResult, _is_flat_pair


def _load_fixture0():
    d = np.load("synth/fixtures/case_000.npz")
    return d


def _hard_p50_split(rgb: np.ndarray) -> tuple[np.ndarray, np.ndarray]:
    """Structured hard-threshold split (demo-style)."""
    lum = 0.2126 * rgb[..., 0] + 0.7152 * rgb[..., 1] + 0.0722 * rgb[..., 2]
    thr = float(np.percentile(lum, 50.0))
    ma = (lum <= thr).astype(np.float32)[..., None]
    mb = 1.0 - ma
    a = (rgb * ma).astype(np.float32)
    b = (rgb * mb).astype(np.float32)
    return a, b


def _flat_pair(h: int, w: int) -> tuple[np.ndarray, np.ndarray]:
    """Solid-red + solid-white degenerate pair (the lab-data failure mode)."""
    red = np.zeros((h, w, 3), dtype=np.float32)
    red[..., 0] = 0.9
    white = np.full((h, w, 3), 0.95, dtype=np.float32)
    return red, white


class TestRankingIntegration:
    """Honest-fix assert: lab-warped fixture ranking (spec test 3)."""

    def test_lab_warped_ranking_affine_beats_flat_and_guard(self):
        """Build lab-warped D_obs; flat must rank below structured under full new scoring;
        with flat_guard the flat pair is rejected (count==1).
        """
        data = _load_fixture0()
        scan = data["scan"].astype(np.float32)
        # Working size: fixture is already small (64×64)
        h, w = scan.shape[:2]
        curve = get_film_curve("Generic")

        # Density from true fixture path when present; else synthesize via curve of scan lum
        if "density" in data.files:
            d_true = data["density"].astype(np.float32)
        else:
            # Fall back: preprocess densitometry
            from app.preprocessing import preprocess_negative
            from PIL import Image
            pil = Image.fromarray((scan * 255).clip(0, 255).astype(np.uint8))
            pre = preprocess_negative(pil, stock="Generic")
            d_true = pre.density.astype(np.float32)
            scan = pre.rgb.astype(np.float32)
            h, w = scan.shape[:2]

        # Lab warp: D_obs = 0.7·D + 0.3 (systematic tone curve)
        d_warped = (0.7 * d_true + 0.3).astype(np.float32)
        mask = np.full(d_warped.shape, VALID, dtype=np.uint8)
        # Confidence mask from fixture if available
        if "confidence_mask" in data.files:
            mask = data["confidence_mask"].astype(np.uint8)

        struct_a, struct_b = _hard_p50_split(scan)
        flat_a, flat_b = _flat_pair(h, w)

        log_exp = torch.full((1, 1, h, w), -0.3)

        # (i) discrimination: affine-mode physics ratio flat/structured ≥ 3× pointwise ratio
        loss_pt = HybridFilmLoss(
            film_curve=curve,
            physics_weight=1.0,
            perceptual_weight=0.0,
            exclusivity_weight=0.0,
            balance_weight=0.0,
            naturalness_weight=0.0,
            calibration="none",
        )
        loss_af = HybridFilmLoss(
            film_curve=curve,
            physics_weight=1.0,
            perceptual_weight=0.0,
            exclusivity_weight=0.0,
            balance_weight=0.0,
            naturalness_weight=0.0,
            calibration="affine",
        )
        bd_flat_pt = loss_pt.evaluate(log_exp, scan, flat_a, flat_b, density=d_warped, confidence_mask=mask)
        bd_str_pt = loss_pt.evaluate(log_exp, scan, struct_a, struct_b, density=d_warped, confidence_mask=mask)
        bd_flat_af = loss_af.evaluate(log_exp, scan, flat_a, flat_b, density=d_warped, confidence_mask=mask)
        bd_str_af = loss_af.evaluate(log_exp, scan, struct_a, struct_b, density=d_warped, confidence_mask=mask)

        ratio_pt = bd_flat_pt.physics_loss / max(bd_str_pt.physics_loss, 1e-12)
        ratio_af = bd_flat_af.physics_loss / max(bd_str_af.physics_loss, 1e-12)
        assert ratio_af >= 3.0 * ratio_pt, (
            f"affine discrimination {ratio_af:.3f} not ≥ 3× pointwise {ratio_pt:.3f}"
        )

        # (ii) FULL new ranking (affine + grad 0.5 + reweight, guard OFF): flat ranks BELOW structured
        flat_c = SeparationResult(
            image_a=flat_a, image_b=flat_b, method="flat", message="", candidate_id="flat_red_white"
        )
        struct_c = SeparationResult(
            image_a=struct_a, image_b=struct_b, method="hard_p50", message="", candidate_id="hard_p50"
        )
        ranked = rank_candidates(
            candidates=[flat_c, struct_c],
            observed_log_exposure=log_exp,
            observed_rgb=scan,
            film_curve=curve,
            physics_weight=1.0,
            perceptual_weight=0.5,
            density=d_warped,
            confidence_mask=mask,
            calibration="affine",
            physics_grad_weight=0.5,
            flat_guard=False,
        )
        ids = [r.candidate_id for r in ranked]
        assert ids.index("hard_p50") < ids.index("flat_red_white"), (
            f"structured must rank above flat without guard; order={ids}; "
            f"scores={[ (r.candidate_id, r.score.total_loss) for r in ranked ]}"
        )

        # (iii) flat_guard=True → flat rejected, rejected_count == 1
        ranked_g = rank_candidates(
            candidates=[flat_c, struct_c],
            observed_log_exposure=log_exp,
            observed_rgb=scan,
            film_curve=curve,
            density=d_warped,
            confidence_mask=mask,
            calibration="affine",
            physics_grad_weight=0.5,
            flat_guard=True,
        )
        assert isinstance(ranked_g, RankingResult)
        assert ranked_g.rejected_count == 1, f"expected 1 rejected, got {ranked_g.rejected_count}"
        assert all(r.candidate_id != "flat_red_white" for r in ranked_g.ranked)
        assert len(ranked_g.ranked) == 1
        assert ranked_g.ranked[0].candidate_id == "hard_p50"


class TestFlatGuardFallback:
    """Spec test 4: all-flat list → unfiltered fallback, non-empty."""

    def test_all_flat_fallback(self):
        h, w = 32, 32
        obs = np.random.default_rng(0).uniform(0.2, 0.8, (h, w, 3)).astype(np.float32)
        fa, fb = _flat_pair(h, w)
        # Second flat pair (different solid colors, still flat)
        fa2 = np.zeros((h, w, 3), dtype=np.float32)
        fa2[..., 2] = 0.8
        fb2 = np.full((h, w, 3), 0.5, dtype=np.float32)

        cands = [
            SeparationResult(fa, fb, "flat", "", "flat1"),
            SeparationResult(fa2, fb2, "flat", "", "flat2"),
        ]
        curve = get_film_curve("Generic")
        log_exp = torch.full((1, 1, h, w), -0.3)
        ranked = rank_candidates(
            candidates=cands,
            observed_log_exposure=log_exp,
            observed_rgb=obs,
            film_curve=curve,
            flat_guard=True,
        )
        assert len(ranked) >= 1, "all-flat must fall back to non-empty unfiltered list"
        assert ranked.rejected_count == 2
        # Independent check: both pairs really are flat under the guard definition
        assert _is_flat_pair(fa, fb, obs)
        assert _is_flat_pair(fa2, fb2, obs)


class TestDefaultsByteIdentical:
    """Defaults must not change scoring behavior for existing call sites."""

    def test_score_separation_defaults_match_class_weights(self):
        data = _load_fixture0()
        scan = data["scan"].astype(np.float32)
        h, w = scan.shape[:2]
        a, b = _hard_p50_split(scan)
        curve = get_film_curve("Generic")
        log_exp = torch.full((1, 1, h, w), -0.3)
        # No density → legacy path; just ensure no crash and finite
        bd = score_separation(log_exp, scan, a, b, curve)
        assert np.isfinite(bd.total_loss)
        assert bd.affine_a == 1.0  # default unset path