#!/usr/bin/env python3
"""Retention of a student against the teacher on an RS suite (spec §7): R = s / t and R_c = (s − c) / (t − c) per
task, with each task's metric and chance level from the suite file.

    python scripts/rs_retention.py --suite configs/eval/rs_fast.yaml \
        --teacher $MT_DATA/evals/teacher/marco-nano-instruct/rs_fast/*/summary.json \
        --student $MT_DATA/runs/<run>/rs_fast/*/summary.json [--markdown]
"""
from __future__ import annotations

import argparse
import json

import yaml


def scores(paths: list[str]) -> dict[str, float]:
    out: dict[str, float] = {}
    for p in paths:
        out.update(json.load(open(p))["scores"])
    return out


def main() -> None:
    p = argparse.ArgumentParser()
    p.add_argument("--suite", required=True)
    p.add_argument("--teacher", nargs="+", required=True)
    p.add_argument("--student", nargs="+", required=True)
    p.add_argument("--markdown", action="store_true")
    a = p.parse_args()
    suite = yaml.safe_load(open(a.suite))
    chance = suite.get("chance", {})
    t, s = scores(a.teacher), scores(a.student)
    rows = []
    for task in sorted(set(t) & set(s)):
        c = chance.get(task, chance.get("default", 0.0))
        rows.append({"task": task, "teacher": t[task], "student": s[task], "R": s[task] / t[task] if t[task] else None,
                     "R_c": (s[task] - c) / (t[task] - c) if t[task] != c else None})
    missing = sorted(set(t) ^ set(s))
    if a.markdown:
        print("| task | teacher | student | R | R_c |\n|------|---------|---------|---|-----|")
        for r in rows:
            print(f"| {r['task']} | {r['teacher']:.3f} | {r['student']:.3f} | {r['R']:.3f} | {r['R_c']:.3f} |")
        if missing:
            print(f"\nNot in both: {', '.join(missing)}")
    else:
        print(json.dumps({"rows": rows, "missing": missing}, indent=1))


if __name__ == "__main__":
    main()