app/stats_math.py

"""Pure probability and ranking math for log and daily statistics.

No I/O: callers pass entry/daily dicts; this module only aggregates.
"""

from __future__ import annotations

from collections import defaultdict
from typing import Any


def _normalize_remedy(value: str) -> str:
    return " ".join(value.strip().lower().split())


def intensity_bucket(intensity: int) -> str:
    """Map intensity 1-10 into coarse buckets."""

    if intensity <= 3:
        return "1-3"
    if intensity <= 6:
        return "4-6"
    return "7-10"


def scored_entries(entries: list[dict[str, Any]]) -> list[dict[str, Any]]:
    """Exclude pending outcomes from probability denominators."""

    return [e for e in entries if e.get("result") != "pending"]


def outcome_histogram(entries: list[dict[str, Any]]) -> dict[str, int]:
    """Count outcomes including pending."""

    counts = {"worked": 0, "partial": 0, "failed": 0, "pending": 0}
    for entry in entries:
        result = entry.get("result")
        if result in counts:
            counts[result] += 1
    return counts


def by_remedy(
    entries: list[dict[str, Any]],
    *,
    min_n: int,
    shrink_k: float,
) -> list[dict[str, Any]]:
    """Compute per-remedy n, p_worked, p_helped, and shrinkage rank."""

    scored = scored_entries(entries)
    totals: dict[str, int] = defaultdict(int)
    worked: dict[str, int] = defaultdict(int)
    helped: dict[str, int] = defaultdict(int)
    for entry in scored:
        key = _normalize_remedy(str(entry.get("remedy") or ""))
        if not key:
            continue
        totals[key] += 1
        result = entry.get("result")
        if result == "worked":
            worked[key] += 1
            helped[key] += 1
        elif result == "partial":
            helped[key] += 1

    rows: list[dict[str, Any]] = []
    for key, n in totals.items():
        if n < min_n:
            continue
        p_worked = worked[key] / n
        p_helped = helped[key] / n
        rank = p_helped * (n / (n + shrink_k))
        rows.append(
            {
                "key": key,
                "n": n,
                "p_worked": p_worked,
                "p_helped": p_helped,
                "rank": rank,
            }
        )
    rows.sort(key=lambda row: (-row["rank"], -row["n"], row["key"]))
    return rows


def by_emotion(entries: list[dict[str, Any]]) -> list[dict[str, Any]]:
    """Emotion multi-label buckets with p_helped among scored uses."""

    totals: dict[str, int] = defaultdict(int)
    helped: dict[str, int] = defaultdict(int)
    for entry in scored_entries(entries):
        result = entry.get("result")
        is_helped = result in ("worked", "partial")
        for emotion in entry.get("emotions") or []:
            key = str(emotion).strip().lower()
            if not key:
                continue
            totals[key] += 1
            if is_helped:
                helped[key] += 1
    rows = [
        {
            "key": key,
            "n": n,
            "p_helped": (helped[key] / n) if n else 0.0,
        }
        for key, n in totals.items()
    ]
    rows.sort(key=lambda row: (-row["n"], row["key"]))
    return rows


def by_tag(entries: list[dict[str, Any]]) -> list[dict[str, Any]]:
    """Tag buckets with failure rate among scored uses."""

    totals: dict[str, int] = defaultdict(int)
    failed: dict[str, int] = defaultdict(int)
    for entry in scored_entries(entries):
        is_failed = entry.get("result") == "failed"
        for tag in entry.get("tags") or []:
            key = str(tag).strip().lower()
            if not key:
                continue
            totals[key] += 1
            if is_failed:
                failed[key] += 1
    rows = [
        {
            "key": key,
            "n": n,
            "p_failed": (failed[key] / n) if n else 0.0,
        }
        for key, n in totals.items()
    ]
    rows.sort(key=lambda row: (-row["p_failed"], -row["n"], row["key"]))
    return rows


def by_intensity_bucket(entries: list[dict[str, Any]]) -> list[dict[str, Any]]:
    """Intensity bucket counts and helped rates."""

    totals = {"1-3": 0, "4-6": 0, "7-10": 0}
    helped = {"1-3": 0, "4-6": 0, "7-10": 0}
    for entry in scored_entries(entries):
        try:
            intensity = int(entry.get("intensity", 0))
        except (TypeError, ValueError):
            continue
        if intensity < 1 or intensity > 10:
            continue
        bucket = intensity_bucket(intensity)
        totals[bucket] += 1
        if entry.get("result") in ("worked", "partial"):
            helped[bucket] += 1
    return [
        {
            "key": key,
            "n": totals[key],
            "p_helped": (helped[key] / totals[key]) if totals[key] else 0.0,
        }
        for key in ("1-3", "4-6", "7-10")
    ]


def corn_ok(row: dict[str, Any]) -> bool:
    """True when corn sessions are within the delay policy."""

    sessions = int(row.get("corn_sessions") or 0)
    delay_ok = bool(row.get("delay_ok", True))
    return sessions == 0 or (sessions <= 1 and delay_ok)


def daily_rates(daily_rows: list[dict[str, Any]]) -> dict[str, float]:
    """Aggregate daily scoreboard rates for a set of days."""

    n = len(daily_rows)
    if n == 0:
        return {
            "p_brick_done": 0.0,
            "p_corn_ok": 0.0,
            "p_no_fc": 0.0,
            "p_rerun_clean": 0.0,
            "p_court_closed": 0.0,
            "avg_points": 0.0,
        }
    brick = sum(1 for row in daily_rows if row.get("brick_done"))
    corn = sum(1 for row in daily_rows if corn_ok(row))
    no_fc = sum(1 for row in daily_rows if row.get("daydream") != "fc")
    rerun = sum(1 for row in daily_rows if row.get("rerun") == "clean")
    court = sum(1 for row in daily_rows if row.get("court") == "closed")
    avg_points = sum(float(row.get("points") or 0) for row in daily_rows) / n
    return {
        "p_brick_done": brick / n,
        "p_corn_ok": corn / n,
        "p_no_fc": no_fc / n,
        "p_rerun_clean": rerun / n,
        "p_court_closed": court / n,
        "avg_points": avg_points,
    }


def helped_tags_by_remedy(entries: list[dict[str, Any]]) -> dict[str, set[str]]:
    """Tags that appear on helped uses of each remedy."""

    tags_by_remedy: dict[str, set[str]] = defaultdict(set)
    for entry in scored_entries(entries):
        if entry.get("result") not in ("worked", "partial"):
            continue
        key = _normalize_remedy(str(entry.get("remedy") or ""))
        if not key:
            continue
        for tag in entry.get("tags") or []:
            token = str(tag).strip().lower()
            if token:
                tags_by_remedy[key].add(token)
    return tags_by_remedy


def match_score(current_tags: set[str], remedy_tags: set[str]) -> float:
    """Fraction of current tags that match a remedy's helped-tag set."""

    if not current_tags:
        return 0.0
    return len(current_tags & remedy_tags) / max(len(current_tags), 1)


def server_picks(
    entries: list[dict[str, Any]],
    current_tags: list[str] | set[str],
    *,
    min_n: int,
    shrink_k: float,
    match_alpha: float,
    limit: int = 5,
) -> list[dict[str, Any]]:
    """Rank remedies by shrinkage + optional tag-match boost."""

    tag_set = {str(tag).strip().lower() for tag in current_tags if str(tag).strip()}
    remedy_tags = helped_tags_by_remedy(entries)
    picks: list[dict[str, Any]] = []
    for row in by_remedy(entries, min_n=min_n, shrink_k=shrink_k):
        match = match_score(tag_set, remedy_tags.get(row["key"], set()))
        pick = row["rank"] * (1.0 + match_alpha * match)
        picks.append(
            {
                "remedy_key": row["key"],
                "pick": pick,
                "n": row["n"],
                "p_helped": row["p_helped"],
                "p_worked": row["p_worked"],
                "rank": row["rank"],
                "match": match,
            }
        )
    picks.sort(key=lambda item: (-item["pick"], -item["n"], item["remedy_key"]))
    return picks[:limit]


def data_thin(n_scored: int) -> bool:
    """True when scored history is too thin for strong coaching."""

    return n_scored < 10


FORMULAS = {
    "p_worked": "N(worked,r) / N(r); pending excluded",
    "p_helped": "N(worked|partial,r) / N(r); pending excluded",
    "rank": "p_helped * n/(n+k)",
    "pick": "rank * (1 + alpha * match)",
    "match": "|T intersect T_r| / max(|T|,1)",
    "DATA_THIN": "n_scored < 10",
}


def build_stats(
    entries: list[dict[str, Any]],
    daily_rows: list[dict[str, Any]],
    *,
    min_n: int,
    shrink_k: float,
    generated_at: str,
) -> dict[str, Any]:
    """Assemble the /api/stats response payload."""

    scored = scored_entries(entries)
    return {
        "n_entries_total": len(entries),
        "n_entries_scored": len(scored),
        "outcomes": outcome_histogram(entries),
        "by_remedy": by_remedy(entries, min_n=min_n, shrink_k=shrink_k),
        "by_emotion": by_emotion(entries),
        "by_tag": by_tag(entries),
        "by_intensity_bucket": by_intensity_bucket(entries),
        "daily": daily_rates(daily_rows),
        "formulas": FORMULAS,
        "generated_at": generated_at,
        "min_n": min_n,
        "shrink_k": shrink_k,
        "DATA_THIN": data_thin(len(scored)),
    }