"""The reward. Quoted verbatim on the demo page, so written to be read by
someone who does not write Python.

Three components, weights summing to 1.0. Two pull toward clearing the alert
quickly with the right citations. The third — the heaviest — pulls the other
way: it pays for not letting a suspicious customer through, and a policy that
never investigates cannot collect it.

Hours are integers in MINUTES so that the world, the grader and the browser
agree to the last digit; the page shows them as hours.

# region: pig-demo/reward
"""

from __future__ import annotations

from dataclasses import dataclass

# The generator's config, NOT industry measurements. Nobody timed an analyst.
SCREEN_MINUTES = 30      # reading the free screen; charged to everyone
TURN_MINUTES = 3         # every reply, accepted or rejected
LOOKUP_MINUTES = 24      # history, counterparty, prior alerts
DOC_MINUTES = 6          # documents on file

CAUGHT_WEIGHT = 0.55
HOURS_WEIGHT = 0.25
EVIDENCE_WEIGHT = 0.20

# What `caught` pays for filing on a clean customer. Not zero: a defensive
# escalation is a wasted filing, not a missed crime. Not one: otherwise a
# policy that escalates everything takes the counterweight for free.
FALSE_ESCALATION_CREDIT = 0.5


@dataclass
class Episode:
    """Everything the reward needs from one played alert."""

    label: str                      # "benign" | "suspicious" — what the generator planted
    planted: list[list[str]]        # alternates; the best one is scored
    disposition: str | None         # "close" | "escalate" | None
    cites: list[str]
    minutes_spent: int
    reference_minutes: int | None   # cheapest shipped policy that got this seed right
    turns_spent: int
    rejected: int
    truncated: bool = False         # ended early with no disposition: not scored
    typology: str | None = None       # what the agent wrote; diagnostic only
    true_typology: str | None = None  # what the generator planted; diagnostic only
    lookups: int = 0


def correct(ep: Episode) -> bool:
    """A close on a benign alert, or an escalation on a suspicious one."""
    return (ep.label == "benign" and ep.disposition == "close") or (
        ep.label == "suspicious" and ep.disposition == "escalate"
    )


def caught(ep: Episode) -> float:
    """Nothing slipped through. The counterweight, and over half the reward.

    1.0 on any correct call. 0.0 if a suspicious alert was closed, or nothing
    was decided. A false escalation earns FALSE_ESCALATION_CREDIT: the
    customer was clean, so nothing slipped, but a report was filed on them.

    Read the asymmetry off the weights: a miss zeroes this AND the two
    objectives below (the call was wrong), so the episode scores 0.000. The
    most any policy can gain by being fast is `hours`' 0.25. One miss costs
    more than every hour saved.
    """
    if ep.disposition is None:
        return 0.0
    if correct(ep):
        return 1.0
    if ep.label == "benign":
        return FALSE_ESCALATION_CREDIT
    return 0.0


def hours(ep: Episode) -> float | None:
    """Cleared it in the hours budgeted, as a ratio against a reference.

    The reference is the cheapest SHIPPED policy that reached the correct
    disposition on this same alert — `fast` where the screen was enough,
    `targeted` where a lookup was needed. Never a true optimum: grade against
    an optimum and the shipped policy fails its own probe on some seeds.

    Capped at 1.0. A correct call faster than the reference is not paid
    extra; the ceiling is our heuristic, not an optimum. A fast wrong answer
    saves nothing. If no shipped policy got this seed right there is no
    reference and this is None — not scored — which CI asserts never happens.
    """
    if not correct(ep):
        return 0.0
    if ep.reference_minutes is None:
        return None
    return min(1.0, ep.reference_minutes / max(1, ep.minutes_spent))


def f1_exact(cites: list[str], planted: list[str]) -> tuple[int, int]:
    """F1 as an exact fraction (numerator, denominator), so both runtimes agree bit for bit."""
    common = len(set(cites) & set(planted))
    return 2 * common, len(set(cites)) + len(set(planted))


def evidence(ep: Episode) -> float:
    """Named the facts that decide it: F1 between the cites and the planted set.

    The planted set is what the generator buried — the explaining KYC or
    document fact on a benign alert, the transactions that constitute the
    typology on a suspicious one. Where a template admits two reasonable
    rationales the best alternate is scored.

    F1's precision term is what stops cite-everything: one planted id among
    eight cited scores 0.22. Its recall term is what stops cite-one-thing on
    a pattern of eight deposits. Zero if the disposition was wrong; a false
    escalation is checked against an empty set. The narrative prose is never
    read here.
    """
    if not correct(ep) or not ep.cites:
        return 0.0
    best_num, best_den = 0, 1
    for alt in ep.planted:
        num, den = f1_exact(ep.cites, alt)
        if num * best_den > best_num * den:
            best_num, best_den = num, den
    return best_num / best_den


def score(ep: Episode) -> dict[str, float | None]:
    """The components, unweighted. None means "not scored", never zero."""
    if ep.truncated:
        return {"caught": None, "hours": None, "evidence": None}
    return {"caught": caught(ep), "hours": hours(ep), "evidence": evidence(ep)}


def total(ep: Episode, weights: dict[str, float] | None = None) -> float | None:
    parts = score(ep)
    if any(v is None for v in parts.values()):
        return None
    w = weights or WEIGHTS
    return parts["caught"] * w["caught"] + parts["hours"] * w["hours"] + parts["evidence"] * w["evidence"]  # type: ignore[operator]


WEIGHTS = {"caught": CAUGHT_WEIGHT, "hours": HOURS_WEIGHT, "evidence": EVIDENCE_WEIGHT}
ROLES = {"caught": "counterweight", "hours": "objective", "evidence": "objective"}

# endregion: pig-demo/reward


def escalate_break_even(evidence_if_close: float = 1.0) -> float:
    """P(benign) below which a blind escalation out-scores a well-cited close.

    Escalating blind pays CAUGHT_WEIGHT × (credit × p_b + (1 − p_b)) plus, on
    the suspicious share, hours (paid 1.0: the blind filer is fastest) and
    nothing for evidence. Closing pays p_b × (CAUGHT + HOURS + EVIDENCE × e).
    Printed on the page so a BSA officer can read the threshold the weights
    encode rather than infer it.
    """
    close_per_benign = CAUGHT_WEIGHT + HOURS_WEIGHT + EVIDENCE_WEIGHT * evidence_if_close
    esc_per_benign = CAUGHT_WEIGHT * FALSE_ESCALATION_CREDIT
    esc_per_suspicious = CAUGHT_WEIGHT + HOURS_WEIGHT
    # p_b · close = p_b · esc_b + (1 − p_b) · esc_s  →  p_b = esc_s / (close − esc_b + esc_s)
    return esc_per_suspicious / (close_per_benign - esc_per_benign + esc_per_suspicious)


def score_exact(ep: Episode) -> dict[str, list[int] | None]:
    """The components as exact fractions, for the cross-language digest."""
    if ep.truncated:
        return {"caught": None, "hours": None, "evidence": None}
    if ep.disposition is None:
        c: list[int] = [0, 1]
    elif correct(ep):
        c = [1, 1]
    elif ep.label == "benign":
        c = [1, 2]
    else:
        c = [0, 1]
    if not correct(ep):
        h: list[int] | None = [0, 1]
    elif ep.reference_minutes is None:
        h = None
    else:
        spent = max(1, ep.minutes_spent)
        h = [1, 1] if ep.reference_minutes >= spent else [ep.reference_minutes, spent]
    if not correct(ep) or not ep.cites:
        e = [0, 1]
    else:
        best = (0, 1)
        for alt in ep.planted:
            num, den = f1_exact(ep.cites, alt)
            if num * best[1] > best[0] * den:
                best = (num, den)
        e = [best[0], best[1]]
    return {"caught": c, "hours": h, "evidence": e}


def metrics(ep: Episode) -> dict[str, float | None]:
    """Diagnostics. Weight zero — reported, never summed into the reward."""
    return {
        "hours_spent": ep.minutes_spent / 60,
        "reference_hours": None if ep.reference_minutes is None else ep.reference_minutes / 60,
        "lookups": float(ep.lookups),
        "turns_used": float(ep.turns_spent),
        "rejected_replies": float(ep.rejected),
        "false_escalation": 1.0 if ep.label == "benign" and ep.disposition == "escalate" else 0.0,
        "missed": 1.0 if ep.label == "suspicious" and ep.disposition == "close" else 0.0,
        "typology_match": None if ep.disposition != "escalate" else (1.0 if ep.typology == ep.true_typology else 0.0),
    }
