"""Collect the evidence for a distilled student that improved on your set and nothing else.

Purpose: the diagnostic half of Part 15's challenge. Four faults produce that same
    symptom, and each leaves a different trace on disk. This reads the run log, the
    teacher's raw output, the training file, the evaluation set and the two
    tokenisers, then writes one report with a verdict line per fault. It changes
    nothing, needs no accelerator and starts no model, so it is safe to run first
    and cheap to run again after a fix.
Platform: all (standard library, plus transformers only for the tokeniser check,
    which is skipped with a stated reason when transformers is not installed)
Minimum memory: 12 GB nominally; the checks themselves are text in memory
Assumes: Python 3.10 or newer. Whatever exists is read; whatever is missing is
    reported as missing rather than guessed, because "not recorded" is itself
    evidence. distillog.py sits next to this file.

Usage: python3 diagnose-distillation.py --labbook labbook.md --train data/train.jsonl \\
           --tasks my-tasks.json --raw raw/teacher.jsonl --report diagnosis.md
       python3 diagnose-distillation.py --labbook labbook.md --train data/train.jsonl \\
           --tasks my-tasks.json --teacher Qwen/Qwen3-8B --student Qwen/Qwen3-1.7B \\
           --report diagnosis.md
       python3 diagnose-distillation.py --labbook labbook.md --report diagnosis.md
       # the minimum: the run log alone, with everything else reported as absent

The four faults, in the order the challenge page works through them:
  1. teacher errors propagated because nothing verified the teacher
  2. the evaluation set contaminated by teacher output
  3. teacher and student tokenisers differ, in a logit distillation run
  4. sampling temperature and loss weighting chosen without being recorded
"""

from __future__ import annotations

import argparse
import json
import re
from collections import Counter
from pathlib import Path
from typing import Any

import distillog

WORD_RE = re.compile(r"[a-z0-9]+")
UNKNOWN = "not available"

FINAL_PATTERNS = (
    re.compile(r"####\s*(-?[0-9][0-9,]*(?:\.[0-9]+)?)"),
    re.compile(r"\\boxed\{\s*(-?[0-9][0-9,]*(?:\.[0-9]+)?)\s*\}"),
    re.compile(r"(?i)\banswer\s*[:=]\s*\$?(-?[0-9][0-9,]*(?:\.[0-9]+)?)"),
)
INVENTED_SPECIFIC = re.compile(
    r"\b(?:\d{1,3}(?:\.\d{1,3}){3}|port\s+\d{2,5}|(?:host(?:name)?|machine)\s+is\s+\S+)\b",
    re.IGNORECASE,
)
REFUSAL_MARKERS = re.compile(
    r"\b(?:cannot|can't|can not|no access|not able|do not have|don't have|unable)\b", re.IGNORECASE
)

PROBE_STRINGS = (
    "<|im_start|>user\nHello, world!<|im_end|>\n",
    "<think>\n2 + 2 = 4\n</think>\nAnswer: 4",
    "quantisation, tokeniser, générateur, 分词器",
)


def normalise(text: str) -> list[str]:
    return WORD_RE.findall(text.lower())


def ngrams(words: list[str], n: int) -> set[tuple[str, ...]]:
    if len(words) < n:
        return {tuple(words)} if words else set()
    return {tuple(words[i:i + n]) for i in range(len(words) - n + 1)}


def containment(a: set[tuple[str, ...]], b: set[tuple[str, ...]]) -> float:
    return len(a & b) / len(a) if a else 0.0


def read_jsonl(path: Path) -> list[dict]:
    if not path.is_file():
        return []
    rows = []
    for line in path.read_text(encoding="utf-8").splitlines():
        line = line.strip()
        if line:
            try:
                rows.append(json.loads(line))
            except json.JSONDecodeError:
                continue
    return rows


def training_texts(rows: list[dict]) -> list[tuple[str, str]]:
    """Accept every dataset shape this course writes, and return (prompt, completion)."""
    out = []
    for row in rows:
        if isinstance(row.get("prompt"), list) and isinstance(row.get("completion"), list):
            prompt = " ".join(m.get("content", "") for m in row["prompt"])
            completion = " ".join(m.get("content", "") for m in row["completion"])
        elif isinstance(row.get("prompt"), str):
            prompt, completion = row["prompt"], row.get("completion", "")
        elif isinstance(row.get("messages"), list):
            prompt = " ".join(m.get("content", "") for m in row["messages"] if m.get("role") != "assistant")
            completion = " ".join(m.get("content", "") for m in row["messages"] if m.get("role") == "assistant")
        elif isinstance(row.get("text"), str):
            prompt, completion = "", row["text"]
        else:
            continue
        out.append((prompt, completion))
    return out


# ---------------------------------------------------------------------------
# Fault 1: teacher errors propagated
# ---------------------------------------------------------------------------

def check_teacher_quality(raw_rows: list[dict], sample: int) -> dict[str, Any]:
    if not raw_rows:
        return {"status": "skipped", "reason": "no --raw teacher file given"}
    checked = raw_rows[:sample] if sample else raw_rows
    verifiable = [r for r in checked if "answer" in r]
    wrong = 0
    unparsed = 0
    for row in verifiable:
        text = row.get("completion") or ""
        predicted = None
        for pattern in FINAL_PATTERNS:
            found = pattern.findall(text)
            if found:
                try:
                    predicted = float(found[-1].replace(",", ""))
                except ValueError:
                    predicted = None
                break
        if predicted is None:
            unparsed += 1
        elif abs(predicted - float(row["answer"])) > 1e-6:
            wrong += 1

    refusal_rows = [r for r in checked if r.get("category") == "refusal"]
    bad_refusals = sum(
        1 for r in refusal_rows
        if INVENTED_SPECIFIC.search(r.get("completion") or "")
        or not REFUSAL_MARKERS.search(r.get("completion") or "")
    )
    empty = sum(1 for r in checked if not (r.get("completion") or "").strip())

    return {
        "status": "checked",
        "sampled": len(checked),
        "verifiable": len(verifiable),
        "verifiably_wrong": wrong,
        "answer_not_found": unparsed,
        "refusal_prompts": len(refusal_rows),
        "refusals_that_invented_or_complied": bad_refusals,
        "empty_completions": empty,
        "error_rate": round((wrong + bad_refusals + empty) / len(checked), 3) if checked else None,
    }


# ---------------------------------------------------------------------------
# Fault 2: contaminated evaluation
# ---------------------------------------------------------------------------

def check_contamination(train_rows: list[dict], tasks_path: Path, n: int, threshold: float) -> dict[str, Any]:
    if not train_rows:
        return {"status": "skipped", "reason": "no --train file given"}
    if not tasks_path or not tasks_path.is_file():
        return {"status": "skipped", "reason": "no --tasks evaluation set given"}
    spec = json.loads(tasks_path.read_text(encoding="utf-8"))
    tasks = spec.get("tasks", [])
    pairs = training_texts(train_rows)
    train_exact = {" ".join(normalise(p)) for p, _ in pairs if p}
    train_grams = [ngrams(normalise(f"{p}\n{c}"), n) for p, c in pairs]

    exact_hits, overlap_hits, worst = [], [], []
    for task in tasks:
        prompt = task.get("prompt", "")
        reference = task.get("reference", "") or ""
        key = " ".join(normalise(prompt))
        if key and key in train_exact:
            exact_hits.append(task.get("id"))
            continue
        target = ngrams(normalise(f"{prompt}\n{reference}"), n)
        best = max((containment(target, g) for g in train_grams if g), default=0.0)
        worst.append((round(best, 3), task.get("id")))
        if best >= threshold:
            overlap_hits.append({"id": task.get("id"), "containment": round(best, 3)})
    worst.sort(reverse=True)
    return {
        "status": "checked",
        "tasks": len(tasks),
        "training_examples": len(pairs),
        "exact_matches": exact_hits,
        "high_overlap": overlap_hits,
        "highest_containment": worst[:5],
        "threshold": threshold,
        "n": n,
    }


# ---------------------------------------------------------------------------
# Fault 3: tokeniser mismatch
# ---------------------------------------------------------------------------

def check_tokenisers(teacher_id: str | None, student_id: str | None) -> dict[str, Any]:
    if not teacher_id or not student_id:
        return {"status": "skipped",
                "reason": "pass --teacher and --student to compare vocabularies"}
    try:
        from transformers import AutoTokenizer  # noqa: PLC0415 - optional here on purpose
    except ImportError:
        return {"status": "skipped", "reason": "transformers is not installed in this environment"}
    try:
        teacher_tok = AutoTokenizer.from_pretrained(teacher_id)
        student_tok = AutoTokenizer.from_pretrained(student_id)
    except (OSError, ValueError) as exc:
        return {"status": "skipped", "reason": f"could not load a tokeniser: {exc}"}

    mismatches = []
    for probe in PROBE_STRINGS:
        a = teacher_tok(probe, add_special_tokens=False)["input_ids"]
        b = student_tok(probe, add_special_tokens=False)["input_ids"]
        if a != b:
            mismatches.append({"probe": probe[:50], "teacher": a[:16], "student": b[:16]})
    teacher_specials = teacher_tok.get_added_vocab()
    student_specials = student_tok.get_added_vocab()
    special_diffs = [
        {"token": t, "teacher_id": i, "student_id": student_specials.get(t)}
        for t, i in sorted(teacher_specials.items()) if student_specials.get(t) != i
    ]
    return {
        "status": "checked",
        "teacher": teacher_id,
        "student": student_id,
        "teacher_class": type(teacher_tok).__name__,
        "student_class": type(student_tok).__name__,
        "teacher_vocab_size": len(teacher_tok),
        "student_vocab_size": len(student_tok),
        "probe_mismatches": mismatches,
        "special_token_differences": special_diffs[:10],
        "compatible": (len(teacher_tok) == len(student_tok) and not mismatches and not special_diffs),
    }


# ---------------------------------------------------------------------------
# Fault 4: settings nobody recorded
# ---------------------------------------------------------------------------

def check_settings(labbook: str) -> dict[str, Any]:
    stages = distillog.read_stages(labbook)
    if not stages:
        return {"status": "skipped", "reason": f"no Part 15 records found in {labbook}"}
    by_stage: dict[str, dict] = {}
    for record in stages:
        by_stage[record.get("stage", "?")] = record

    generate = by_stage.get("generate", {})
    train_logit = by_stage.get("train-logit", {})
    gen_hyper = generate.get("hyperparameters") or {}
    logit_hyper = train_logit.get("hyperparameters") or {}

    missing = []
    for stage in ("generate", "filter", "train"):
        if stage not in by_stage:
            missing.append(stage)
    return {
        "status": "checked",
        "stages_present": sorted(by_stage),
        "stages_missing": missing,
        "generation_temperature": gen_hyper.get("temperature", UNKNOWN),
        "generation_top_p": gen_hyper.get("top_p", UNKNOWN),
        "generation_mode": gen_hyper.get("mode", UNKNOWN),
        "generation_samples_per_prompt": gen_hyper.get("samples", UNKNOWN),
        "distillation_beta": logit_hyper.get("beta", UNKNOWN),
        "distillation_temperature": logit_hyper.get("temperature", UNKNOWN),
        "distillation_trainer": logit_hyper.get("trainer", UNKNOWN),
        "filter_thresholds": (by_stage.get("filter", {}).get("hyperparameters") or {}) or UNKNOWN,
    }


# ---------------------------------------------------------------------------

def verdict(name: str, failed: bool, skipped: bool, detail: str) -> str:
    if skipped:
        return f"- **{name}: not checked.** {detail}"
    return f"- **{name}: {'EVIDENCE FOUND' if failed else 'no evidence'}.** {detail}"


def main() -> None:
    parser = argparse.ArgumentParser(description=__doc__.splitlines()[0])
    parser.add_argument("--labbook", default="labbook.md")
    parser.add_argument("--raw", default=None, help="the teacher's raw output, for the verifier check")
    parser.add_argument("--train", default=None, help="the training file the student learned from")
    parser.add_argument("--tasks", default=None, help="the evaluation set the gain was measured on")
    parser.add_argument("--teacher", default=None, help="teacher model id or path, for the tokeniser check")
    parser.add_argument("--student", default=None, help="student model id or path, for the tokeniser check")
    parser.add_argument("--sample", type=int, default=200, help="teacher answers to verify")
    parser.add_argument("--n", type=int, default=13)
    parser.add_argument("--threshold", type=float, default=0.5)
    parser.add_argument("--report", default="diagnosis.md")
    args = parser.parse_args()

    raw_rows = read_jsonl(Path(args.raw)) if args.raw else []
    train_rows = read_jsonl(Path(args.train)) if args.train else []

    teacher_quality = check_teacher_quality(raw_rows, args.sample)
    contamination = check_contamination(train_rows, Path(args.tasks) if args.tasks else None,
                                        args.n, args.threshold)
    tokenisers = check_tokenisers(args.teacher, args.student)
    settings = check_settings(args.labbook)

    fault1 = (teacher_quality.get("status") == "checked"
              and (teacher_quality.get("error_rate") or 0) > 0.02)
    fault2 = (contamination.get("status") == "checked"
              and (contamination.get("exact_matches") or contamination.get("high_overlap")))
    fault3 = tokenisers.get("status") == "checked" and not tokenisers.get("compatible")
    fault4 = (settings.get("status") == "checked"
              and (settings.get("stages_missing")
                   or settings.get("generation_temperature") == UNKNOWN))

    lines = [
        "# Distillation diagnosis",
        "",
        "Evidence collected without changing anything. Read the verdicts, then read the",
        "sections: a verdict is a pointer to the numbers under it, not a conclusion.",
        "",
        "## Verdicts",
        "",
        verdict("1. Teacher errors propagated",
                bool(fault1), teacher_quality.get("status") != "checked",
                teacher_quality.get("reason", "")
                or f"{teacher_quality.get('verifiably_wrong', 0)} verifiably wrong, "
                   f"{teacher_quality.get('refusals_that_invented_or_complied', 0)} bad refusals, "
                   f"{teacher_quality.get('empty_completions', 0)} empty, out of "
                   f"{teacher_quality.get('sampled', 0)} sampled."),
        verdict("2. Evaluation contaminated by teacher output",
                bool(fault2), contamination.get("status") != "checked",
                contamination.get("reason", "")
                or f"{len(contamination.get('exact_matches', []))} exact match(es) and "
                   f"{len(contamination.get('high_overlap', []))} task(s) above containment "
                   f"{args.threshold} against {contamination.get('training_examples', 0)} "
                   "training examples."),
        verdict("3. Tokeniser mismatch",
                bool(fault3), tokenisers.get("status") != "checked",
                tokenisers.get("reason", "")
                or (f"teacher {tokenisers.get('teacher_vocab_size')} entries, student "
                    f"{tokenisers.get('student_vocab_size')} entries, "
                    f"{len(tokenisers.get('probe_mismatches', []))} probe mismatch(es).")),
        verdict("4. Settings not recorded, or a stage missing",
                bool(fault4), settings.get("status") != "checked",
                settings.get("reason", "")
                or f"stages present: {', '.join(settings.get('stages_present', []))}; "
                   f"missing: {', '.join(settings.get('stages_missing', [])) or 'none'}; "
                   f"generation temperature: {settings.get('generation_temperature')}."),
        "",
        "## 1. Did anything verify the teacher?",
        "",
        "```json",
        json.dumps(teacher_quality, indent=2),
        "```",
        "",
        "A teacher error rate above a couple of per cent, propagated into training data with",
        "no verifier, is enough to teach a student a wrong habit that no evaluation on the",
        "same distribution will catch.",
        "",
        "## 2. Does the training data overlap the evaluation set?",
        "",
        "```json",
        json.dumps(contamination, indent=2),
        "```",
        "",
        "An exact match is contamination and the fix is deletion. A high containment score is",
        "a pair to read yourself: paraphrase or coincidence is a judgement, not a threshold.",
        "",
        "## 3. Do the teacher and the student share a vocabulary?",
        "",
        "```json",
        json.dumps(tokenisers, indent=2),
        "```",
        "",
        "This matters only for logit distillation, where a teacher probability is applied to",
        "the student's vocabulary at the same index. Sequence-level distillation moves text,",
        "so a mismatch here is not a fault for that route.",
        "",
        "## 4. What settings produced this?",
        "",
        "```json",
        json.dumps(settings, indent=2, default=str),
        "```",
        "",
        f"Records read from `{args.labbook}`. A stage with no record did not necessarily not",
        "happen, but it cannot be reproduced, which for this purpose is the same thing.",
        "",
    ]
    Path(args.report).write_text("\n".join(lines), encoding="utf-8")

    verdict_start = lines.index("## Verdicts")
    print("\n".join(lines[verdict_start:verdict_start + 6]))
    print(f"\nfull report written to {args.report}")
    if not any([fault1, fault2, fault3, fault4]):
        print("No fault found from the evidence supplied. Check what you did not pass in: "
              "the checks that were skipped are listed as 'not checked' above.")


if __name__ == "__main__":
    main()