"""Summaries of graded generations, and paired comparisons between two runs."""
from __future__ import annotations

import json
import math
from pathlib import Path


def read_jsonl(path):
    with open(path) as f:
        return [json.loads(line) for line in f if line.strip()]


def _mean(xs):
    xs = list(xs)
    return sum(xs) / len(xs) if xs else None


def _se(xs):
    xs = list(xs)
    if len(xs) < 2:
        return None
    m = sum(xs) / len(xs)
    return math.sqrt(sum((x - m) ** 2 for x in xs) / (len(xs) - 1) / len(xs))


def summarize(recs: list[dict]) -> dict:
    S = [s for r in recs for s in r["samples"]]
    fin = [s for s in S if not s["cut"]]
    cut = [s for s in S if s["cut"]]
    per_problem = [_mean(s["correct"] for s in r["samples"]) for r in recs]
    out = {
        "problems": len(recs),
        "samples": len(S),
        "acc": _mean(s["correct"] for s in S),
        "acc_se": _se(per_problem),
        "cut_share": _mean(s["cut"] for s in S),
        "mean_tokens": _mean(s["tokens"] for s in S),
        "closed_share": _mean(s["closed"] for s in S),
        "p_finish": len(fin) / len(S) if S else None,
        "p_correct_given_finished": _mean(s["correct"] for s in fin),
        "p_correct_given_cut": _mean(s["correct"] for s in cut),
        "no_answer_finished_share": _mean((s.get("pred") is None and not s["cut"]) for s in S),
        "loop_end_share_of_cut": _mean((s["verbatim_loop_end"] or s["repetitive_end"]) for s in cut),
        "tail_dup8_cut_mean": _mean(s["tail_dup8"] for s in cut),
        "median_tokens_correct": _median([s["tokens"] for s in S if s["correct"]]),
    }
    srcs = sorted({r.get("source") for r in recs if r.get("source")})
    if len(srcs) > 1:
        out["by_source"] = {src: {k: v for k, v in summarize([r for r in recs if r.get("source") == src]).items()
                                  if k in ("problems", "samples", "acc", "acc_se", "cut_share", "mean_tokens")}
                            for src in srcs}
    return out


def _median(xs):
    xs = sorted(xs)
    if not xs:
        return None
    m = len(xs) // 2
    return xs[m] if len(xs) % 2 else (xs[m - 1] + xs[m]) / 2


def compare(a: list[dict], b: list[dict]) -> dict:
    """Paired comparison of run b against run a on the problems both answered."""
    A = {r["id"]: r for r in a}
    B = {r["id"]: r for r in b}
    ids = [i for i in A if i in B]
    d = [_mean(s["correct"] for s in B[i]["samples"]) - _mean(s["correct"] for s in A[i]["samples"]) for i in ids]
    a_any = {i for i in ids if any(s["correct"] for s in A[i]["samples"])}
    b_any = {i for i in ids if any(s["correct"] for s in B[i]["samples"])}
    return {
        "problems": len(ids),
        "gain": _mean(d),
        "gain_se": _se(d),
        "problems_better": sum(x > 0 for x in d),
        "problems_worse": sum(x < 0 for x in d),
        # a problem is "lost" if a solved it in some sample and b never did
        "lost": len(a_any - b_any),
        "rescued": len(b_any - a_any),
    }


def load_run(run_dir) -> list[dict]:
    p = Path(run_dir) / "merged.jsonl"
    return read_jsonl(p)