"""
Practice: from traces to tests. Lab 14's scripts 06 (feedback linked to traces) and 07 (traces
-> tests): tell a complaint from other feedback, sample a stratified test set from traces, gate
a release on what production showed matters, and read the result against its likely range.

A solution.
"""
import math
import random

COMPLAINTS = {"not_helpful", "retry", "edit", "human"}

KINDS = ("what", "how", "list", "none")
LONG_TURN = 8

GATE = {"right": 0.95, "cites": 0.90, "long_input_tokens": 700}


def is_complaint(kinds):
    return bool(set(kinds) & COMPLAINTS)


def build_evalset(traces, per_stratum=10, seed=7):
    rng = random.Random(seed)
    by_conv = {}
    for t in traces:
        by_conv.setdefault(t["conversation_id"], []).append(t)
    clean = [t for t in traces if t["incident"] is None and not t["error"]]
    cases = []
    for kind in KINDS:
        pool = [t for t in clean if t["kind"] == kind and t["turn"] == 1]
        for t in rng.sample(pool, per_stratum):
            cases.append({"from_trace": t["trace_id"], "question": t["question"], "kind": kind,
                          "history": [], "stratum": f"single turn, {kind}"})
    long = [t for t in clean if t["turn"] >= LONG_TURN]
    for t in rng.sample(long, per_stratum):
        earlier = sorted((x for x in by_conv[t["conversation_id"]] if x["turn"] < t["turn"]), key=lambda x: x["turn"])
        cases.append({"from_trace": t["trace_id"], "question": t["question"], "kind": t["kind"],
                      "history": [(x["question"], x["answer"]) for x in earlier if not x["error"]],
                      "stratum": "long conversation (turn 8+)"})
    return cases


def wilson(k, n, z=1.96):
    if n == 0:
        return 0.0, 1.0
    p = k / n
    centre = (p + z * z / (2 * n)) / (1 + z * z / n)
    half = z * math.sqrt(p * (1 - p) / n + z * z / (4 * n * n)) / (1 + z * z / n)
    return max(0.0, centre - half), min(1.0, centre + half)


def gate(rows, rule="bar"):
    right = sum(r["right"] for r in rows) / len(rows)
    ans = [r for r in rows if r["answerable"] and r["right"]]
    k, n = sum(r["cites"] for r in ans), len(ans)
    long = [r["input_tokens"] for r in rows if r["long"]]
    long_tokens = sum(long) / len(long)
    cites_ok = k / n >= GATE["cites"] if rule == "bar" else wilson(k, n)[1] >= GATE["cites"]
    checks = {"right": right >= GATE["right"], "cites": cites_ok,
              "long_tokens": long_tokens <= GATE["long_input_tokens"]}
    return {"right": right, "cites": (k, n), "long_tokens": long_tokens, "checks": checks, "pass": all(checks.values())}


def chance_at_most(k, n, p):
    return sum(math.comb(n, i) * p ** i * (1 - p) ** (n - i) for i in range(k + 1))