"""Scoring for detection and fill quality.

Detection scoring keeps two things apart that are easy to conflate:

* a **false positive** -- the detector flagged a slot the bench did not corrupt, and
  the window's ground truth says was good;
* a **legacy disagreement** -- the detector flagged a slot the old filter passed.

On a vetted window the second is a candidate *find*, not an error: the old filter is
what we are trying to beat, so scoring its misses against a new detector would
penalise exactly the improvement we want.

Fill scoring is reported both in radiance and after the display mapping the videos
actually use, plus a temporal term.  A frame can score well on per-frame PSNR and
still read as a visible stutter at 60 fps, so flicker is measured explicitly.
"""

from dataclasses import dataclass, field

import numpy as np

#: Display mapping used by merger_FITS.py to turn radiance into pixels, by band.
#: (vmin, vmax, gamma).  Fill error is reported through this because it is what the
#: eye sees -- an error in the dim corona matters far less than one on the disc.
DISPLAY_MAPPING = {
    94: (0.050, 8.0, 0.375),
    131: (0.05, 8.0, 0.40),
    171: (0.100, 20.0, 0.425),
    195: (0.10, 30.0, 0.45),
    284: (0.100, 40.0, 0.475),
    304: (0.1, 90.0, 0.5),
}


def to_display(image, wavelength):
    """Map radiance to the 0-1 display range merger_FITS.py renders with."""
    vmin, vmax, gamma = DISPLAY_MAPPING.get(wavelength, (0.0, 1.0, 1.0))
    clean = np.nan_to_num(np.asarray(image, dtype=np.float64), nan=0.0)
    return np.clip((clean - vmin) / vmax, 0.0, 1.0) ** gamma


# ------------------------------------------------------------------------ detection


@dataclass
class DetectionScore:
    """Confusion counts and derived rates for one detector on one bench case."""

    true_positives: int = 0
    false_positives: int = 0
    false_negatives: int = 0
    true_negatives: int = 0
    unknown: int = 0
    #: Slots the detector flagged that the legacy filter had passed and the bench
    #: did not corrupt.  Reported, never counted as errors.
    legacy_disagreements: int = 0
    #: Recall broken down by corruption mode, so a detector that only catches
    #: blackouts cannot hide behind a good overall number.
    recall_by_mode: dict = field(default_factory=dict)
    elapsed_us: int = 0
    frames: int = 0

    @property
    def precision(self):
        flagged = self.true_positives + self.false_positives
        return self.true_positives / flagged if flagged else float("nan")

    @property
    def recall(self):
        actual = self.true_positives + self.false_negatives
        return self.true_positives / actual if actual else float("nan")

    @property
    def f1(self):
        precision, recall = self.precision, self.recall
        if not np.isfinite(precision) or not np.isfinite(recall) or precision + recall == 0:
            return float("nan")
        return 2 * precision * recall / (precision + recall)

    @property
    def false_positive_rate(self):
        negatives = self.false_positives + self.true_negatives
        return self.false_positives / negatives if negatives else float("nan")

    @property
    def microseconds_per_frame(self):
        return self.elapsed_us / self.frames if self.frames else float("nan")

    def as_dict(self):
        return {
            "true_positives": self.true_positives,
            "false_positives": self.false_positives,
            "false_negatives": self.false_negatives,
            "true_negatives": self.true_negatives,
            "unknown": self.unknown,
            "legacy_disagreements": self.legacy_disagreements,
            "precision": self.precision,
            "recall": self.recall,
            "f1": self.f1,
            "false_positive_rate": self.false_positive_rate,
            "us_per_frame": self.microseconds_per_frame,
            "recall_by_mode": dict(self.recall_by_mode),
        }


def score_detection(verdicts, injected, legacy_good=frozenset(), elapsed_us=0):
    """Score detector output against what the bench actually injected.

    `verdicts` maps slot -> Verdict.  `injected` maps slot -> corruption mode name
    for every slot the bench damaged (a deleted slot has no frame to judge, so it is
    not scored here).  `legacy_good` is the set of slots the old filter passed.

    'unknown' verdicts are counted separately and excluded from precision and recall
    rather than folded into 'good' -- a detector that abstains has not made an error,
    but it has not made a call either, and hiding that would flatter it.
    """
    score = DetectionScore(elapsed_us=elapsed_us, frames=len(verdicts))
    by_mode = {}

    for slot, verdict in verdicts.items():
        mode = injected.get(slot)
        corrupted = mode is not None

        if verdict.verdict == "unknown":
            score.unknown += 1
            if corrupted:
                by_mode.setdefault(mode, [0, 0])[1] += 1
            continue

        flagged = verdict.verdict == "bad"
        if corrupted:
            stats = by_mode.setdefault(mode, [0, 0])
            stats[1] += 1
            if flagged:
                score.true_positives += 1
                stats[0] += 1
            else:
                score.false_negatives += 1
        elif flagged:
            score.false_positives += 1
            if slot in legacy_good:
                score.legacy_disagreements += 1
        else:
            score.true_negatives += 1

    score.recall_by_mode = {
        mode: (caught / total if total else float("nan"))
        for mode, (caught, total) in sorted(by_mode.items())
    }
    return score


#: How several detectors' verdicts are folded into one.
COMBINATION_POLICIES = ("any", "all", "majority")


def combine_verdicts(verdict_maps, policy="any"):
    """Fold several detectors' verdicts into one per slot.

    Which detectors to run in production is a trade-off between recall and the
    false-positive rate, and the only honest way to choose is to score the
    combinations the same way as the individuals.  Because `detection` rows already
    store a verdict per frame, this is pure post-processing.

    * ``any``      -- bad if any detector says bad.  Maximum recall; false positives
                      accumulate across detectors.
    * ``all``      -- bad only if every detector that voted says bad.  Minimum false
                      positives, and the right shape for a near-zero FP budget.
    * ``majority`` -- bad if more than half the votes say bad.

    ``unknown`` abstains rather than voting.  A detector that declines to judge must
    not be silently counted as saying "good": under ``all`` that would let one
    abstention veto a real detection, and under ``any`` it would inflate recall.  A
    slot where every detector abstains is itself ``unknown``.

    `verdict_maps` is a sequence of {slot: Verdict}.  Returns {slot: Verdict}.
    """
    if policy not in COMBINATION_POLICIES:
        raise ValueError(f"Unknown policy {policy!r}; expected one of {COMBINATION_POLICIES}")
    if not verdict_maps:
        return {}

    from .detectors import BAD, GOOD, UNKNOWN, Verdict

    combined = {}
    slots = set()
    for verdicts in verdict_maps:
        slots.update(verdicts)

    for slot in slots:
        votes = []
        reasons = []
        for verdicts in verdict_maps:
            verdict = verdicts.get(slot)
            if verdict is None or verdict.verdict == UNKNOWN:
                continue
            votes.append(verdict.verdict == BAD)
            if verdict.verdict == BAD and verdict.reason:
                reasons.append(verdict.reason)

        if not votes:
            combined[slot] = Verdict(UNKNOWN, "no detector judged this frame", {})
            continue

        if policy == "any":
            is_bad = any(votes)
        elif policy == "all":
            is_bad = all(votes)
        else:
            is_bad = sum(votes) * 2 > len(votes)

        combined[slot] = Verdict(
            BAD if is_bad else GOOD,
            "; ".join(reasons[:3]) if is_bad else None,
            {"votes_bad": float(sum(votes)), "votes_total": float(len(votes))},
        )
    return combined


def precision_recall_curve(scores, labels):
    """Precision/recall across every threshold of a continuous detector score.

    `scores` are higher-is-more-suspicious. Returns (thresholds, precision, recall).
    """
    scores = np.asarray(scores, dtype=np.float64)
    labels = np.asarray(labels, dtype=bool)
    finite = np.isfinite(scores)
    scores, labels = scores[finite], labels[finite]
    if scores.size == 0 or not labels.any():
        return np.array([]), np.array([]), np.array([])

    order = np.argsort(-scores)
    scores, labels = scores[order], labels[order]
    true_positives = np.cumsum(labels)
    flagged = np.arange(1, scores.size + 1)
    precision = true_positives / flagged
    recall = true_positives / labels.sum()
    return scores, precision, recall


def average_precision(scores, labels):
    """Area under the precision-recall curve; the threshold-free summary."""
    _, precision, recall = precision_recall_curve(scores, labels)
    if recall.size == 0:
        return float("nan")
    return float(np.sum(np.diff(np.concatenate([[0.0], recall])) * precision))


# ----------------------------------------------------------------------------- fill


@dataclass
class FillScore:
    """Reconstruction error for one filled frame."""

    rmse: float
    mae: float
    log_rmse: float
    psnr: float
    ssim: float
    gap_frames: int = 0
    wavelength: int = 0

    def as_dict(self):
        return {
            "rmse": self.rmse,
            "mae": self.mae,
            "log_rmse": self.log_rmse,
            "psnr": self.psnr,
            "ssim": self.ssim,
            "gap_frames": self.gap_frames,
            "wavelength": self.wavelength,
        }


def score_fill(filled, truth, wavelength, gap_frames=0):
    """Compare a reconstructed frame against the frame that was withheld.

    Radiance errors (rmse/mae) are dominated by the bright disc; the log term keeps
    the faint corona visible in the score; psnr and ssim are computed after the
    display mapping, because that is the image a viewer actually sees.
    """
    filled = np.nan_to_num(np.asarray(filled, dtype=np.float64), nan=0.0)
    truth = np.nan_to_num(np.asarray(truth, dtype=np.float64), nan=0.0)
    if filled.shape != truth.shape:
        raise ValueError(f"shape mismatch: {filled.shape} vs {truth.shape}")

    residual = filled - truth
    rmse = float(np.sqrt(np.mean(residual**2)))
    mae = float(np.mean(np.abs(residual)))

    log_residual = np.log1p(np.clip(filled, 0, None)) - np.log1p(np.clip(truth, 0, None))
    log_rmse = float(np.sqrt(np.mean(log_residual**2)))

    # Imported here, not at module scope: DISPLAY_MAPPING and to_display are the only
    # parts of this module the training host needs, and scikit-image on that host is
    # built against a different numpy ABI than the ROCm image's torch.  A constants
    # table should not drag an image-processing library in behind it.
    from skimage.metrics import structural_similarity

    shown_fill, shown_truth = to_display(filled, wavelength), to_display(truth, wavelength)
    display_mse = float(np.mean((shown_fill - shown_truth) ** 2))
    psnr = float("inf") if display_mse == 0 else float(10.0 * np.log10(1.0 / display_mse))
    ssim = float(structural_similarity(shown_truth, shown_fill, data_range=1.0))

    return FillScore(rmse, mae, log_rmse, psnr, ssim, gap_frames, wavelength)


def temporal_flicker(sequence, truth_sequence, wavelength):
    """How much the reconstruction's frame-to-frame motion departs from reality.

    Per-frame PSNR is blind to the artifact that matters most in a 60 fps video: a
    filled run can match each frame tolerably and still freeze then jump.  This
    compares the *rate of change* rather than the frames, in display space.

    Returns mean absolute difference of successive-frame deltas; 0 is perfect.
    """
    if len(sequence) != len(truth_sequence):
        raise ValueError("sequences must be the same length")
    if len(sequence) < 2:
        return float("nan")

    shown = [to_display(frame, wavelength) for frame in sequence]
    shown_truth = [to_display(frame, wavelength) for frame in truth_sequence]
    deltas = [np.mean(np.abs(shown[i] - shown[i - 1])) for i in range(1, len(shown))]
    truth_deltas = [
        np.mean(np.abs(shown_truth[i] - shown_truth[i - 1])) for i in range(1, len(shown_truth))
    ]
    return float(np.mean(np.abs(np.array(deltas) - np.array(truth_deltas))))


def summarise_fills(scores):
    """Aggregate per-frame fill scores, overall and by gap length."""
    if not scores:
        return {}
    by_gap = {}
    for score in scores:
        by_gap.setdefault(score.gap_frames, []).append(score)

    def mean(values):
        finite = [v for v in values if np.isfinite(v)]
        return float(np.mean(finite)) if finite else float("nan")

    return {
        "n": len(scores),
        "rmse": mean([s.rmse for s in scores]),
        "log_rmse": mean([s.log_rmse for s in scores]),
        "psnr": mean([s.psnr for s in scores]),
        "ssim": mean([s.ssim for s in scores]),
        "by_gap": {
            gap: {
                "n": len(group),
                "rmse": mean([s.rmse for s in group]),
                "psnr": mean([s.psnr for s in group]),
                "ssim": mean([s.ssim for s in group]),
            }
            for gap, group in sorted(by_gap.items())
        },
    }