"""Git-history supervision audit: teacher SFT payloads vs student SWE50 runs (read-only).

Inputs are the dumps of extract_train_tool_calls.py and extract_eval_tool_calls.mjs.
One taxonomy (git_history_taxonomy.py) is applied to both sides, so prevalence,
probe outcomes, teacher fingerprints and cost are measured on the same scale.
Optional --labels adds the double-annotated stratified sample (Cohen's kappa,
Wilson intervals, population-weighted rates). Output is diagnostic statistics,
never training data; no gold, hidden tests or judge details are read.

    python -X utf8 analysis/swe/audit_git_history_supervision.py --dump <dir> \
        --out <json> [--labels <labels.json> --strata <strata_key.json>]
"""
from __future__ import annotations

import argparse
from collections import Counter, defaultdict
import hashlib
import json
import math
from pathlib import Path
import random
import re
import sys

sys.path.insert(0, str(Path(__file__).resolve().parent))
from git_history_taxonomy import (REASONING_STYLES, SEGMENT_SPLIT, classify_command, classify_git_segment,  # noqa: E402
                                  git_templates, intent_hits, style_hits)

RUNS = ["base", "run1", "run2"]
FINGERPRINTS = {
    "fsck --no-reflogs --unreachable": r"--no-reflogs\s+--unreachable|--unreachable\s+--no-reflogs",
    "git fsck": r"\bgit\s+fsck\b",
    "git count-objects -v": r"count-objects\s+-v",
    "git branch -a": r"\bgit\s+branch\s+-a\b",
    "git tag": r"\bgit\s+tag\b",
    "git log --all": r"\bgit\s+log\b[^|&;\n]*--all\b",
    "--decorate": r"--decorate",
    "git reflog": r"\bgit\s+reflog\b",
    "git remote -v": r"\bgit\s+remote\s+-v",
    "git log --since/--after/--until": r"\bgit\s+log\b[^|&;\n]*--(?:since|after|until|before)\b",
    "git log --grep": r"\bgit\s+log\b[^|&;\n]*--grep",
    "git log --oneline -N（定位）": r"\bgit\s+log\s+--oneline\s+-\d+\s*($|&&|\||;)",
}
WRITES = re.compile(r"sed -i|\.write_text\(|open\([^)]*['\"][wa]['\"]|cat\s*>|\btee\s|>\s*[\w./-]+\.(py|pyx|c|h|cfg|toml|txt|rst|js|ts)\b|patch\s+-p|git apply")
FILTER = re.compile(r"^\s*(head|tail|wc|sort|uniq|grep|cat|awk|sed -n|echo|true)\b")
NULL_LINE = re.compile(r"^\s*(\*\s*\S+|\(no output\)|\[no output\]|\[stderr\]|\[exit code: -?\d+\]|fatal:.*|error:.*|warning:.*"
                       r"|count: 0|size: 0|in-pack: \d+|packs: \d+|size-pack: \d+|prune-packable: 0|garbage: 0|size-garbage: 0"
                       r"|Checking .*|Dangling .*none.*)\s*$")
BEYOND = re.compile(r"\b(dangling|unreachable) (commit|blob|tree|tag)\b|\(tag: |origin/|remotes/|HEAD@\{\d+\}")


def read(path: Path):
    with path.open(encoding="utf-8") as handle:
        for line in handle:
            if line.strip():
                yield json.loads(line)


def digest(path: Path) -> str:
    with path.open("rb") as handle:
        return hashlib.file_digest(handle, "sha256").hexdigest()


def quantiles(values):
    values = sorted(values)
    if not values:
        return dict(n=0)
    def q(f):
        x = (len(values) - 1) * f
        i = int(x)
        return round(values[i] + (values[min(i + 1, len(values) - 1)] - values[i]) * (x - i), 1)
    return dict(n=len(values), p50=q(.5), p90=q(.9), max=values[-1], mean=round(sum(values) / len(values), 1))


def wilson(k, n, z=1.96):
    if n == 0:
        return [None, None]
    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 [round(max(0, centre - half), 4), round(min(1, centre + half), 4)]


def eval_repo(instance_id: str) -> str:
    return instance_id.split("__")[1].rsplit("-", 1)[0]


def pure_probe(command: str) -> bool:
    kinds = []
    for segment in SEGMENT_SPLIT.split(command):
        segment = re.sub(r"^\s*cd\s+\S+\s*$", "", segment).strip()
        if not segment:
            continue
        cls = classify_git_segment(segment)
        if cls:
            kinds.append(cls)
        elif not FILTER.match(segment):
            return False
    return bool(kinds) and "H3" in kinds and all(k in {"H3", "H1"} for k in kinds)


def probe_outcome(obs: str | None) -> str:
    if obs is None:
        return "no_observation"
    lines = [l for l in obs.strip().splitlines() if l.strip()]
    if not lines or all(NULL_LINE.match(l) for l in lines):
        return "null"
    if BEYOND.search(obs):
        return "other_refs_listed"
    if any(re.match(r"^\s*\*?\s*[0-9a-f]{7,12}\b", l) for l in lines):
        decorated = re.search(r"\(HEAD -> [^)]*\)", obs)
        return "head_history_only" if decorated or "pre-agent commit" in obs or "SWE-bench" in obs else "commit_list"
    return "other_output"


def is_mutation(call: dict) -> bool:
    ok = call["exit"] == 0 and not (call["obs"] or "").lstrip().startswith(("Error", "error:"))
    if call.get("is_error"):
        ok = False
    if call["tool"] == "edit":
        return ok and call["obs"] is not None
    return ok and call["tool"] == "bash" and bool(WRITES.search(call["command"])) and not re.search(r"(>|tee\s+|write_text\(|open\()\s*['\"]?/tmp/", call["command"])


def load(dump: Path):
    train_calls = [r for r in read(dump / "train_calls.jsonl")]
    train_msgs = [r for r in read(dump / "train_messages.jsonl")]
    eval_calls = [r for r in read(dump / "eval_calls.jsonl")]
    eval_msgs = [r for r in read(dump / "eval_messages.jsonl")]
    for r in train_calls:
        r["cls"] = classify_command(r["command"], r["repo_dir"]) if r["tool"] == "bash" else None
    for r in eval_calls:
        r["cls"] = classify_command(r["command"], eval_repo(r["instance_id"])) if r["tool"] == "bash" else None
    return train_calls, train_msgs, eval_calls, eval_msgs


def prevalence(calls, units, flag):
    hit, n = set(), 0
    for r in calls:
        if r["cls"] and flag(r["cls"]):
            hit.add(r["instance_id"])
            n += 1
    return dict(units=len(hit), of=len(units), share=round(len(hit) / len(units), 4) if units else None, calls=n)


def comparability(train_calls, eval_calls):
    flags = {"legacy_train": lambda c: c["legacy_train"], "legacy_eval": lambda c: c["legacy_eval"],
             **{k: (lambda k: lambda c: k in c["classes"])(k) for k in ["H3", "U", "H2", "H1"]}}
    a_units = {r["instance_id"] for r in train_calls if r["run"] != "C3"}
    out = {}
    for name, flag in flags.items():
        row = dict(train_supervised=prevalence([r for r in train_calls if r["run"] != "C3" and r["target"]], a_units, flag),
                   train_masked=prevalence([r for r in train_calls if r["run"] != "C3" and not r["target"]], a_units, flag))
        for group in ["baseline", "cap3"]:
            units = {r["instance_id"] for r in train_calls if r["run"] == group}
            row["train_" + group] = prevalence([r for r in train_calls if r["run"] == group and r["target"]], units, flag)
        for run in RUNS:
            rows = [r for r in eval_calls if r["run"] == run]
            row[run] = prevalence(rows, {r["instance_id"] for r in rows}, flag)
        out[name] = row
    return out


def fingerprints(train_calls, eval_calls):
    teacher = [r for r in train_calls if r["run"] != "C3" and r["target"] and r["tool"] == "bash"]
    t_units = {r["instance_id"] for r in train_calls if r["run"] != "C3"}
    out = []
    for name, pattern in FINGERPRINTS.items():
        rx = re.compile(pattern)
        row = dict(name=name, teacher=len({r["instance_id"] for r in teacher if rx.search(r["command"])}), teacher_of=len(t_units))
        for run in RUNS:
            rows = [r for r in eval_calls if r["run"] == run and r["tool"] == "bash"]
            row[run] = len({r["instance_id"] for r in rows if rx.search(r["command"])})
            row[run + "_calls"] = sum(bool(rx.search(r["command"])) for r in rows)
        out.append(row)
    return out


def template_overlap(train_calls, eval_calls, seed=20261030, boots=1000):
    def h3_templates(command):
        return [template for segment in SEGMENT_SPLIT.split(command) if classify_git_segment(segment) == "H3"
                for template in git_templates(segment)]
    teacher = Counter(t for r in train_calls if r["run"] != "C3" and r["target"] and r["cls"] and r["cls"]["top"] == "H3"
                      for t in h3_templates(r["command"]))
    known = set(teacher)
    def jsd(p: Counter, q: Counter):
        keys = set(p) | set(q)
        sp, sq = sum(p.values()), sum(q.values())
        if not sp or not sq:
            return None
        total = 0.0
        for k in keys:
            a, b = p[k] / sp, q[k] / sq
            m = (a + b) / 2
            total += (a * math.log2(a / m) if a else 0) + (b * math.log2(b / m) if b else 0)
        return total / 2
    rng = random.Random(seed)
    out = {}
    for run in RUNS:
        per_task = defaultdict(Counter)
        for r in eval_calls:
            if r["run"] == run and r["cls"] and r["cls"]["top"] == "H3":
                per_task[r["instance_id"]].update(h3_templates(r["command"]))
        pooled = sum(per_task.values(), Counter())
        matched = sum(c for t, c in pooled.items() if t in known)
        tasks = sorted(per_task)
        samples = []
        for _ in range(boots):
            pick = [tasks[rng.randrange(len(tasks))] for _ in tasks] if tasks else []
            boot = sum((per_task[t] for t in pick), Counter())
            value = jsd(teacher, boot)
            if value is not None:
                samples.append(value)
        samples.sort()
        out[run] = dict(h3_templates=sum(pooled.values()), distinct=len(pooled),
                        exact_teacher_template_share=round(matched / sum(pooled.values()), 4) if pooled else None,
                        jsd_vs_teacher=round(jsd(teacher, pooled), 4) if pooled else None,
                        jsd_ci95=[round(samples[int(.025 * len(samples))], 4), round(samples[int(.975 * len(samples)) - 1], 4)] if samples else None,
                        top=[[t, c, t in known] for t, c in pooled.most_common(12)])
    out["teacher_top"] = [[t, c] for t, c in teacher.most_common(20)]
    return out


def probe_outcomes(train_calls, eval_calls):
    out = {}
    for key, rows in [("train_supervised", [r for r in train_calls if r["target"]]),
                      *[(run, [r for r in eval_calls if r["run"] == run]) for run in RUNS]]:
        probes = [r for r in rows if r["cls"] and r["cls"]["top"] == "H3"]
        pure = [r for r in probes if pure_probe(r["command"])]
        out[key] = dict(h3_calls=len(probes), pure_probe_calls=len(pure), outcomes=dict(Counter(probe_outcome(r["obs"]) for r in pure)))
    # Date-window searches are the student's generalization; record how often they return nothing.
    for run in RUNS:
        dated = [r for r in eval_calls if r["run"] == run and r["tool"] == "bash" and re.search(r"\bgit\s+log\b[^|&;\n]*--(since|after)\b", r["command"])]
        out[run]["date_window_calls"] = len(dated)
        out[run]["date_window_empty"] = sum(probe_outcome(r["obs"]) == "null" for r in dated)
    return out


def intent(train_calls, train_msgs, eval_msgs):
    probe_turns = {(r["instance_id"], r["step"]) for r in train_calls if r["cls"] and r["cls"]["top"] in {"H3", "U"}}
    sup = [m for m in train_msgs if m["target"]]
    hits = [m for m in sup if intent_hits(m["reasoning"])]
    near = [m for m in hits if any((m["instance_id"], m["step"] + d) in probe_turns for d in (-1, 0, 1))]
    phrase = Counter(p for m in hits for p in intent_hits(m["reasoning"]))
    out = dict(train_supervised_turns=len(sup), train_intent_turns=len(hits), train_intent_instances=len({m["instance_id"] for m in hits}),
               train_intent_near_probe=len(near), top_phrases=phrase.most_common(15))
    for run in RUNS:
        msgs = [m for m in eval_msgs if m["run"] == run]
        h = [m for m in msgs if intent_hits(m["reasoning"])]
        out[run] = dict(turns=len(h), tasks=len({m["instance_id"] for m in h}), of_turns=len(msgs))
    return out


def reasoning_styles(train_msgs, eval_msgs):
    """Answer-seeking reasoning in supervised teacher turns vs student turns, and next to runaways."""
    sides = [("teacher", [m for m in train_msgs if m["target"] and m["run"] != "C3"])]
    sides += [(run, [m for m in eval_msgs if m["run"] == run]) for run in RUNS]
    out = {}
    for key, msgs in sides:
        units = {m["instance_id"] for m in msgs}
        row = dict(units=len(units), turns=len(msgs))
        for name in REASONING_STYLES:
            hit = [m for m in msgs if name in style_hits(m["reasoning"])]
            row[name] = dict(units=len({m["instance_id"] for m in hit}), turns=len(hit))
        any_hit = [m for m in msgs if style_hits(m["reasoning"])]
        row["any"] = dict(units=len({m["instance_id"] for m in any_hit}), turns=len(any_hit),
                          target_tokens=sum(m.get("target_tokens", 0) for m in any_hit))
        if key != "teacher":
            runaway = [m for m in msgs if m["stop_reason"] == "length" and m["reasoning_chars"] >= 20000]
            other = [m for m in msgs if m["stop_reason"] != "length" and m["reasoning_chars"] >= 2000]
            row["runaway_vs_long"] = {name: dict(runaway=sum(name in style_hits(m["reasoning"]) for m in runaway), runaway_n=len(runaway),
                                                 other=sum(name in style_hits(m["reasoning"]) for m in other), other_n=len(other))
                                      for name in REASONING_STYLES}
        out[key] = row
    return out


def cost(train_calls, train_msgs, eval_calls, eval_msgs):
    out = {}
    probe_turns = defaultdict(set)
    for r in train_calls:
        if r["target"] and r["cls"] and r["cls"]["top"] in {"H3", "U"}:
            probe_turns[r["cls"]["top"]].add((r["instance_id"], r["step"]))
    sup = [m for m in train_msgs if m["target"]]
    total = sum(m["target_tokens"] for m in sup)
    for k in ["H3", "U"]:
        turns = [m for m in sup if (m["instance_id"], m["step"]) in probe_turns[k]]
        out["train_" + k] = dict(turns=len(turns), turn_share=round(len(turns) / len(sup), 4),
                                 target_tokens=sum(m["target_tokens"] for m in turns), token_share=round(sum(m["target_tokens"] for m in turns) / total, 4))
    out["train_supervised_tokens"] = total

    def first_mutation(calls):
        steps = [c["step"] for c in calls if is_mutation(c)]
        return min(steps) if steps else None

    for key, calls in [("train", [r for r in train_calls if r["run"] != "C3"]), *[(run, [r for r in eval_calls if r["run"] == run]) for run in RUNS]]:
        by = defaultdict(list)
        for c in calls:
            by[c["instance_id"]].append(c)
        before, after, steps, obs_chars, per_task = 0, 0, [], 0, []
        for iid, cs in by.items():
            fm = first_mutation(cs)
            probes = [c for c in cs if c["cls"] and c["cls"]["top"] == "H3" and (key != "train" or c["target"])]
            per_task.append(len(probes))
            for c in probes:
                steps.append(c["step"])
                obs_chars += c["obs_len"] or 0
                if fm is None or c["step"] < fm:
                    before += 1
                else:
                    after += 1
        out[key] = dict(h3_calls=before + after, before_first_mutation=before,
                        before_share=round(before / (before + after), 4) if before + after else None,
                        step=quantiles(steps), obs_chars=obs_chars, per_unit=quantiles(per_task))
    for run in RUNS:
        msgs = {(m["instance_id"], m["step"]): m for m in eval_msgs if m["run"] == run}
        probe_steps = {(c["instance_id"], c["step"]) for c in eval_calls if c["run"] == run and c["cls"] and c["cls"]["top"] == "H3"}
        runaway = {k for k, m in msgs.items() if m["stop_reason"] == "length" and m["reasoning_chars"] >= 20000}
        after_probe = sum(any((i, s + d) in runaway for d in (1, 2, 3)) for i, s in probe_steps)
        base_rate = sum(any((i, s + d) in runaway for d in (1, 2, 3)) for i, s in msgs) / len(msgs)
        out[run]["runaway_within3_after_probe"] = dict(k=after_probe, n=len(probe_steps), rate=round(after_probe / len(probe_steps), 4) if probe_steps else None,
                                                     any_step_rate=round(base_rate, 4))
        outcome = {}
        for c in eval_calls:
            if c["run"] == run:
                outcome[c["instance_id"]] = c["outcome"]
        counts = Counter(c["instance_id"] for c in eval_calls if c["run"] == run and c["cls"] and c["cls"]["top"] == "H3")
        for label, test in [("h3_ge5", lambda n: n >= 5), ("h3_lt5", lambda n: n < 5)]:
            tasks = [i for i in outcome if test(counts.get(i, 0))]
            k = sum(outcome[i] == "resolved" for i in tasks)
            out[run][label] = dict(tasks=len(tasks), resolved=k, rate_ci=wilson(k, len(tasks)))
    return out


def edit_error_kind(obs: str | None) -> str:
    obs = (obs or "").strip()
    if not obs.startswith("Error"):
        return "ok"
    for kind, pattern in [("path_outside_workspace", r"outside the session workspace"), ("file_not_found", r"File not found|No such file"),
                          ("must_read_first", r"requires reading|must read|not been read"), ("no_change", r"must differ"),
                          ("ambiguous_match", r"matched \d+ times"), ("old_string_not_found", r"not found|could not find|did not match")]:
        if re.search(pattern, obs, re.I):
            return kind
    return "other_error"


def tool_protocol(train_calls, eval_calls):
    """Native read/edit demonstrations: how many supervised calls actually succeeded."""
    out = {}
    for tool in ["read", "edit"]:
        rows = [r for r in train_calls if r["run"] != "C3" and r["tool"] == tool]
        out["teacher_" + tool] = dict(supervised=dict(Counter(edit_error_kind(r["obs"]) for r in rows if r["target"])),
                                      masked=dict(Counter(edit_error_kind(r["obs"]) for r in rows if not r["target"])))
    teacher = Counter(r["tool"] for r in train_calls if r["run"] != "C3" and r["target"])
    out["teacher_tool_mix"] = dict(teacher)
    for run in RUNS:
        rows = [r for r in eval_calls if r["run"] == run]
        out[run] = dict(tool_mix=dict(Counter(r["tool"] for r in rows)),
                        edit=dict(Counter(edit_error_kind(r["obs"]) for r in rows if r["tool"] == "edit")))
    return out


def mask_scenarios(train_calls, train_msgs):
    """What turn-level masking of probe turns would and would not remove."""
    probe = {(r["instance_id"], r["step"]) for r in train_calls if r["run"] != "C3" and r["cls"] and r["cls"]["top"] in {"H3", "U"}}
    failed_tool = {(r["instance_id"], r["step"]) for r in train_calls if r["run"] != "C3" and r["tool"] in {"read", "edit"} and edit_error_kind(r["obs"]) != "ok"}
    sup = [m for m in train_msgs if m["target"] and m["run"] != "C3"]
    total = sum(m["target_tokens"] for m in sup)
    key = lambda m: (m["instance_id"], m["step"])
    def share(rows):
        return dict(turns=len(rows), turn_share=round(len(rows) / len(sup), 4), token_share=round(sum(m["target_tokens"] for m in rows) / total, 4),
                    instances=len({m["instance_id"] for m in rows}))
    rest = [m for m in sup if key(m) not in probe]
    seeking = lambda m: bool(set(style_hits(m["reasoning"])) & {"R2_recall_upstream", "R3_target_commit"})
    residual = [m for m in rest if seeking(m)]
    dirty = {m["instance_id"] for m in sup if key(m) in probe or seeking(m)}
    return dict(supervised_turns=len(sup), supervised_tokens=total, probe_turns=share([m for m in sup if key(m) in probe]),
                failed_tool_turns=share([m for m in sup if key(m) in failed_tool]),
                residual_seeking_after_probe_mask=dict(share(residual), of_remaining_turns=round(len(residual) / len(rest), 4),
                                                       of_remaining_tokens=round(sum(m["target_tokens"] for m in residual) / sum(m["target_tokens"] for m in rest), 4)),
                clean_instances=len({m["instance_id"] for m in sup} - dirty), instances=len({m["instance_id"] for m in sup}),
                style_token_share={name: round(sum(m["target_tokens"] for m in sup if name in style_hits(m["reasoning"])) / total, 4) for name in REASONING_STYLES})


def state_coverage(train_calls, train_msgs, eval_calls, eval_msgs):
    """State exposure at the start of each turn; file writes are heuristic, not semantic fixes."""
    out = {}
    for group, calls, msgs in [("teacher", train_calls, train_msgs)] + [
        (run, [r for r in eval_calls if r["run"] == run], [m for m in eval_msgs if m["run"] == run]) for run in RUNS
    ]:
        call_by, msg_by = defaultdict(list), defaultdict(list)
        for c in calls:
            call_by[c["instance_id"]].append(c)
        for m in msgs:
            msg_by[m["instance_id"]].append(m)
        totals, units, eligible_units = Counter(), defaultdict(set), set()
        turns = 0
        for iid, messages in msg_by.items():
            stream = sorted(call_by[iid], key=lambda c: (c["step"], c.get("call_index", 0)))
            prior, index, prev = [], 0, None
            for m in sorted(messages, key=lambda m: m["step"]):
                while index < len(stream) and stream[index]["step"] < m["step"]:
                    prior.append(stream[index])
                    index += 1
                if group == "teacher" and not m["target"]:
                    prev = m
                    continue
                eligible_units.add(iid)
                turns += 1
                last_write = next((j for j in range(len(prior) - 1, -1, -1) if is_mutation(prior[j])), -1)
                signatures = [json.dumps([c["tool"], c["command"], c.get("args", {})], sort_keys=True) for c in prior[-5:]]
                context = m.get("context_start") if group == "teacher" else m.get("input_tokens")
                flags = {
                    "step_gt50": m["step"] > 50, "step_gt100": m["step"] > 100,
                    "no_write15": len(prior) - last_write - 1 >= 15,
                    "repeat3of5": bool(signatures) and Counter(signatures).most_common(1)[0][1] >= 3,
                    "after_length": prev is not None and prev.get("stop_reason") == "length",
                    "after_long8000": prev is not None and prev["reasoning_chars"] >= 8000,
                    "context_ge64000": context is not None and context >= 64000,
                    "reasoning_ge20000": m["reasoning_chars"] >= 20000,
                }
                for name, hit in flags.items():
                    if hit:
                        totals[name] += 1
                        units[name].add(iid)
                prev = m
        out[group] = dict(turns=turns, instances=len(eligible_units),
                          features={name: dict(turns=totals[name], share=round(totals[name] / turns, 6),
                                               instances=len(units[name])) for name in flags})
    return out


def upstream(train_calls):
    rows = [r for r in train_calls if r["cls"] and "U" in r["cls"]["classes"]]
    by = defaultdict(list)
    for r in rows:
        by[r["instance_id"]].append(r)
    later = defaultdict(list)
    for r in train_calls:
        payload = r["command"] if r["tool"] == "bash" else json.dumps(r["args"], ensure_ascii=False)
        later[r["instance_id"]].append((r["step"], r["call_index"], payload, r["tool"] == "edit" or bool(WRITES.search(payload))))
    flagged = []
    for iid, us in by.items():
        copied = 0
        for u in us:
            lines = {l.strip() for l in (u["obs"] or "").splitlines() if len(l.strip()) >= 25 and re.search(r"[A-Za-z_]\w*\s*[\(=:]", l)}
            writes = " ".join(p for s, ci, p, w in later[iid] if w and (s, ci) > (u["step"], u["call_index"]))
            copied = max(copied, sum(l in writes for l in lines))
        flagged.append(dict(instance_id=iid, group=us[0]["run"], calls=len(us), supervised_calls=sum(u["target"] for u in us),
                            copied_lines=copied, commands=[u["command"][:200] for u in us[:3]]))
    flagged.sort(key=lambda x: -x["copied_lines"])
    return dict(instances=len(by), supervised_instances=len({r["instance_id"] for r in rows if r["target"]}),
                copy_ge2=[f for f in flagged if f["copied_lines"] >= 2], all=flagged)


def kappa(a, b):
    cats = sorted(set(a) | set(b))
    n = len(a)
    if not n:
        return None
    po = sum(x == y for x, y in zip(a, b)) / n
    ca, cb = Counter(a), Counter(b)
    pe = sum(ca[c] * cb[c] for c in cats) / (n * n)
    return round((po - pe) / (1 - pe), 4) if pe < 1 else 1.0


def annotation(labels_path: Path, strata_path: Path):
    labels = json.loads(labels_path.read_text(encoding="utf-8"))
    strata = json.loads(strata_path.read_text(encoding="utf-8"))
    key, population = strata["key"], strata["population"]
    pairs = {}
    first = {l["item_id"]: l for batch in ["b1", "b2"] for l in labels[batch + "_opus"]["labels"]}
    second = {l["item_id"]: l for batch in ["b1", "b2"] for l in labels[batch + "_sonnet"]["labels"]}
    for item, pair in labels.get("additional", {}).items():
        first[item], second[item] = pair["a"], pair["b"]
    for item in key:
        assert item in first and item in second, f"Missing independent label: {item}"
        pairs[item] = (first[item], second[item])
    out = dict(items=len(pairs), design=strata.get("design"), agreement={}, strata={}, categories={})
    for field in ["imitate", "primary_category", "info_gain", "intent_to_obtain_answer"]:
        a = [str(p[0][field]) for p in pairs.values()]
        b = [str(p[1][field]) for p in pairs.values()]
        out["agreement"][field] = dict(raw=round(sum(x == y for x, y in zip(a, b)) / len(a), 4), kappa=kappa(a, b))
    bad = lambda l: l["imitate"] == "no"
    for stratum in population:
        items = [i for i in pairs if key[i]["stratum"] == stratum]
        both = sum(bad(pairs[i][0]) and bad(pairs[i][1]) for i in items)
        either = sum(bad(pairs[i][0]) or bad(pairs[i][1]) for i in items)
        partial_or_no = sum(pairs[i][0]["imitate"] != "yes" and pairs[i][1]["imitate"] != "yes" for i in items)
        privileged = sum(pairs[i][0]["primary_category"] == "privileged_info_seeking" and pairs[i][1]["primary_category"] == "privileged_info_seeking" for i in items)
        out["strata"][stratum] = dict(n=len(items), population=population[stratum],
                                      no_both=both, no_both_ci=wilson(both, len(items)), no_either=either, no_either_ci=wilson(either, len(items)),
                                      not_yes_both=partial_or_no, not_yes_both_ci=wilson(partial_or_no, len(items)),
                                      privileged_both=privileged, privileged_both_ci=wilson(privileged, len(items)))
    total = sum(population.values())
    for metric in ["no_both", "no_either", "not_yes_both"]:
        est, var = 0.0, 0.0
        for s in population:
            n = out["strata"][s]["n"]
            if not n:
                continue
            p = out["strata"][s][metric] / n
            weight = population[s] / total
            est += weight * p
            var += weight ** 2 * (1 - n / population[s]) * p * (1 - p) / max(n - 1, 1)
        out["weighted_" + metric] = dict(estimate=round(est, 4), ci95=[round(max(0, est - 1.96 * math.sqrt(var)), 4), round(min(1, est + 1.96 * math.sqrt(var)), 4)])
    for who in [0, 1]:
        out["categories"]["annotator_" + "AB"[who]] = dict(Counter(p[who]["primary_category"] for p in pairs.values()))
    out["disagreements"] = [dict(item=i, stratum=key[i]["stratum"], a=p[0]["imitate"] + "/" + p[0]["primary_category"],
                                 b=p[1]["imitate"] + "/" + p[1]["primary_category"]) for i, p in pairs.items() if p[0]["imitate"] != p[1]["imitate"]]
    out["pairs"] = {i: dict(stratum=key[i]["stratum"], a=p[0], b=p[1]) for i, p in pairs.items()}
    return out


def main():
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--dump", type=Path, required=True)
    parser.add_argument("--out", type=Path, required=True)
    parser.add_argument("--labels", type=Path)
    parser.add_argument("--strata", type=Path)
    args = parser.parse_args()
    train_calls, train_msgs, eval_calls, eval_msgs = load(args.dump)
    report = dict(schema_version=1, training_eligible=False, benchmark_training_eligible=False,
                  scope="Teacher SFT payloads (merged A/C mask + C3) vs student SWE50 base/run1/run2; regex taxonomy, descriptive.",
                  inputs={p.name: digest(p) for p in sorted(args.dump.glob("*.jsonl"))},
                  code_sha256={p.name: digest(p) for p in [Path(__file__), Path(__file__).with_name("git_history_taxonomy.py")]},
                  comparability=comparability(train_calls, eval_calls), fingerprints=fingerprints(train_calls, eval_calls),
                  templates=template_overlap(train_calls, eval_calls), probe_outcomes=probe_outcomes(train_calls, eval_calls),
                  intent=intent(train_calls, train_msgs, eval_msgs), styles=reasoning_styles(train_msgs, eval_msgs),
                  cost=cost(train_calls, train_msgs, eval_calls, eval_msgs), tool_protocol=tool_protocol(train_calls, eval_calls),
                  mask_scenarios=mask_scenarios(train_calls, train_msgs),
                  state_coverage=state_coverage(train_calls, train_msgs, eval_calls, eval_msgs), upstream=upstream(train_calls))
    if args.labels and args.strata:
        report["annotation"] = annotation(args.labels, args.strata)
    args.out.parent.mkdir(parents=True, exist_ok=True)
    args.out.write_text(json.dumps(report, ensure_ascii=False, indent=1) + "\n", encoding="utf-8")
    brief = {k: report[k] for k in ["probe_outcomes", "intent"]}
    brief["cost"] = {k: v for k, v in report["cost"].items() if k != "train_supervised_tokens"}
    brief["templates"] = {r: {k: v for k, v in report["templates"][r].items() if k != "top"} for r in RUNS}
    brief["upstream"] = dict(instances=report["upstream"]["instances"], supervised=report["upstream"]["supervised_instances"],
                             copy_ge2=[(f["instance_id"], f["copied_lines"], f["supervised_calls"]) for f in report["upstream"]["copy_ge2"]])
    print(json.dumps(brief, ensure_ascii=False, indent=1))


if __name__ == "__main__":
    main()