"""Check a supervised fine-tuning dataset for duplicates and for overlap with an evaluation set.

Purpose: the honesty check that has to run before any fine-tune is measured. Reads a
         training file and the Part 10 task file, reports exact duplicates inside the
         training data, exact matches against the evaluation set, and near-duplicates
         found by shared word n-grams, and can write a cleaned training file with the
         contaminated examples removed.
Platform: all (standard library only; no model, no accelerator, no network)
Minimum memory: 8 GB, and far less in practice: the whole check is text in memory
Assumes: Python 3.10 or newer. The training file is JSON Lines in any of the shapes
         Part 11's dataset lesson describes: {"prompt", "completion"}, {"messages": [...]}
         or {"text"}. The task file is the Part 10 evaluation set, a JSON object with a
         "tasks" list whose items carry "prompt" and usually "reference".

Usage: python3 decontaminate.py --train data/train.jsonl --tasks my-tasks.json
       python3 decontaminate.py --train data/train.jsonl --tasks my-tasks.json \
           --write-clean data/train-clean.jsonl --labbook labbook.md
       python3 decontaminate.py --train data/train.jsonl --valid data/valid.jsonl \
           --tasks my-tasks.json --n 13 --threshold 0.4 --strict

Two kinds of contamination are reported separately because they have different fixes.
An exact match means the same text is in both files, and the fix is to delete it from
the training data. A high n-gram containment means one text is largely contained in the
other, which is what happens when an evaluation prompt was paraphrased into a training
example, and the fix is a judgement call you have to make by reading the pair.
"""
from __future__ import annotations

import argparse
import hashlib
import json
import re
import sys
from datetime import datetime, timezone
from pathlib import Path
from typing import Any

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


def normalise(text: str) -> list[str]:
    """Lowercase, drop punctuation, split on words. Two texts that differ only in
    formatting have to compare equal, or the check misses the commonest duplicate."""
    return WORD_RE.findall(text.lower())


def ngrams(words: list[str], n: int) -> set[tuple[str, ...]]:
    """The set of word n-grams. Short texts fall back to one gram of the whole text so
    that a two-word answer is still comparable rather than silently empty."""
    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:
    """How much of a is also in b, between 0 and 1. Containment rather than Jaccard,
    because a short evaluation prompt buried inside a long training example is exactly
    the case that matters and Jaccard would score it low."""
    if not a:
        return 0.0
    return len(a & b) / len(a)


def digest(words: list[str]) -> str:
    return hashlib.sha256(" ".join(words).encode("utf-8")).hexdigest()


def record_text(record: dict[str, Any]) -> str:
    """Flatten one training record to the text a comparison should see, whichever of the
    three dataset shapes it is in."""
    if "messages" in record:
        parts = []
        for message in record["messages"]:
            content = message.get("content")
            if isinstance(content, str):
                parts.append(content)
        return "\n".join(parts)
    if "prompt" in record or "completion" in record:
        parts = []
        for key in ("prompt", "completion"):
            value = record.get(key)
            if isinstance(value, str):
                parts.append(value)
            elif isinstance(value, list):
                for message in value:
                    content = message.get("content") if isinstance(message, dict) else None
                    if isinstance(content, str):
                        parts.append(content)
        return "\n".join(parts)
    if isinstance(record.get("text"), str):
        return record["text"]
    return ""


def load_jsonl(path: Path) -> list[dict[str, Any]]:
    records = []
    with path.open("r", encoding="utf-8") as handle:
        for number, line in enumerate(handle, start=1):
            line = line.strip()
            if not line:
                continue
            try:
                records.append(json.loads(line))
            except json.JSONDecodeError as exc:
                raise SystemExit(f"{path}:{number}: not valid JSON ({exc.msg})") from exc
    return records


def load_tasks(path: Path) -> list[dict[str, str]]:
    data = json.loads(path.read_text(encoding="utf-8"))
    tasks = data.get("tasks") if isinstance(data, dict) else data
    if not isinstance(tasks, list):
        raise SystemExit(f"{path} does not contain a list of tasks; expected the Part 10 task file")
    out = []
    for index, task in enumerate(tasks):
        text = "\n".join(
            str(task[key]) for key in ("prompt", "reference") if isinstance(task.get(key), str)
        )
        out.append({"id": str(task.get("id", f"task-{index:03d}")), "text": text})
    return out


def summarise(text: str, limit: int = 90) -> str:
    flat = " ".join(text.split())
    return flat if len(flat) <= limit else flat[: limit - 1] + "…"


def check_split(
    name: str,
    records: list[dict[str, Any]],
    task_grams: list[tuple[str, set[tuple[str, ...]]]],
    task_hashes: dict[str, str],
    n: int,
    threshold: float,
) -> dict[str, Any]:
    """One split against the evaluation set, plus its own internal duplicates."""
    seen: dict[str, int] = {}
    internal: list[dict[str, Any]] = []
    exact: list[dict[str, Any]] = []
    near: list[dict[str, Any]] = []
    contaminated: set[int] = set()
    empty = 0

    for index, record in enumerate(records):
        text = record_text(record)
        words = normalise(text)
        if not words:
            empty += 1
            continue
        key = digest(words)
        if key in seen:
            internal.append({"index": index, "duplicate_of": seen[key], "text": summarise(text)})
            contaminated.add(index)
            continue
        seen[key] = index

        if key in task_hashes:
            exact.append({"index": index, "task": task_hashes[key], "text": summarise(text)})
            contaminated.add(index)
            continue

        grams = ngrams(words, n)
        worst_task, worst_score = None, 0.0
        for task_id, gram_set in task_grams:
            score = max(containment(grams, gram_set), containment(gram_set, grams))
            if score > worst_score:
                worst_task, worst_score = task_id, score
        if worst_score >= threshold:
            near.append({
                "index": index, "task": worst_task,
                "containment": round(worst_score, 3), "text": summarise(text),
            })
            contaminated.add(index)

    return {
        "split": name,
        "examples": len(records),
        "empty_examples": empty,
        "internal_duplicates": internal,
        "exact_matches": exact,
        "near_duplicates": near,
        "contaminated_indices": sorted(contaminated),
    }


def main() -> None:
    parser = argparse.ArgumentParser(description=__doc__.splitlines()[0])
    parser.add_argument("--train", required=True, help="training file, JSON Lines")
    parser.add_argument("--valid", default=None, help="validation file, checked the same way")
    parser.add_argument("--tasks", required=True, help="the Part 10 evaluation task file")
    parser.add_argument("--n", type=int, default=13,
                        help="n-gram size for the near-duplicate check; 13 is the size the "
                             "decontamination sections of several pretraining papers use")
    parser.add_argument("--threshold", type=float, default=0.5,
                        help="report a pair when either text's n-grams are this fraction "
                             "contained in the other's")
    parser.add_argument("--report", default=None, help="write the full findings here as JSON")
    parser.add_argument("--write-clean", default=None,
                        help="write the training file with every flagged example removed")
    parser.add_argument("--labbook", default=None, help="append one JSON line recording this check")
    parser.add_argument("--strict", action="store_true",
                        help="exit 1 when anything was flagged, for use in a script")
    args = parser.parse_args()

    train_path = Path(args.train)
    tasks_path = Path(args.tasks)
    for path in (train_path, tasks_path):
        if not path.is_file():
            raise SystemExit(f"{path} does not exist")

    tasks = load_tasks(tasks_path)
    task_grams = [(t["id"], ngrams(normalise(t["text"]), args.n)) for t in tasks]
    task_hashes = {digest(normalise(t["text"])): t["id"] for t in tasks}

    train_records = load_jsonl(train_path)
    results = [check_split("train", train_records, task_grams, task_hashes, args.n, args.threshold)]
    if args.valid:
        valid_path = Path(args.valid)
        if not valid_path.is_file():
            raise SystemExit(f"{valid_path} does not exist")
        results.append(check_split("validation", load_jsonl(valid_path), task_grams,
                                   task_hashes, args.n, args.threshold))

    flagged = 0
    print(f"tasks in evaluation set: {len(tasks)}   n-gram size: {args.n}   "
          f"threshold: {args.threshold}")
    for result in results:
        counts = (len(result["internal_duplicates"]), len(result["exact_matches"]),
                  len(result["near_duplicates"]))
        flagged += sum(counts)
        print(f"\n{result['split']}: {result['examples']} example(s), "
              f"{counts[0]} internal duplicate(s), {counts[1]} exact match(es) against the "
              f"evaluation set, {counts[2]} near-duplicate(s)")
        for row in result["exact_matches"][:5]:
            print(f"  EXACT  line {row['index'] + 1} == task {row['task']}: {row['text']}")
        for row in result["near_duplicates"][:5]:
            print(f"  NEAR   line {row['index'] + 1} ~ task {row['task']} "
                  f"(containment {row['containment']}): {row['text']}")
        for row in result["internal_duplicates"][:5]:
            print(f"  DUP    line {row['index'] + 1} repeats line {row['duplicate_of'] + 1}: "
                  f"{row['text']}")

    if args.report:
        Path(args.report).write_text(
            json.dumps({"tasks_file": str(tasks_path), "n": args.n,
                        "threshold": args.threshold, "results": results}, indent=2),
            encoding="utf-8")
        print(f"\nfull findings written to {args.report}")

    if args.write_clean:
        drop = set(results[0]["contaminated_indices"])
        kept = [r for i, r in enumerate(train_records) if i not in drop]
        with Path(args.write_clean).open("w", encoding="utf-8") as handle:
            for record in kept:
                handle.write(json.dumps(record, ensure_ascii=False) + "\n")
        print(f"clean training file: {args.write_clean} "
              f"({len(kept)} kept, {len(train_records) - len(kept)} removed)")

    if args.labbook:
        line = {
            "check": "part-13/decontaminate",
            "date": datetime.now(timezone.utc).strftime("%Y-%m-%dT%H:%M:%SZ"),
            "train": str(train_path),
            "tasks": str(tasks_path),
            "n": args.n,
            "threshold": args.threshold,
            "flagged": flagged,
            "per_split": [
                {"split": r["split"], "examples": r["examples"],
                 "internal_duplicates": len(r["internal_duplicates"]),
                 "exact_matches": len(r["exact_matches"]),
                 "near_duplicates": len(r["near_duplicates"])}
                for r in results
            ],
        }
        book = Path(args.labbook)
        if not book.exists():
            book.write_text("# Lab notebook\n\n## Results\n\n", encoding="utf-8")
        with book.open("a", encoding="utf-8") as handle:
            handle.write(json.dumps(line, sort_keys=True) + "\n")
        print(f"recorded the check in {args.labbook}")

    if flagged == 0:
        print("\nNothing flagged. The training data and the evaluation set are disjoint at "
              "this n-gram size and threshold, which is what a comparable measurement needs.")
    else:
        print(f"\n{flagged} item(s) flagged. Read them before you train: a number measured on a "
              "contaminated evaluation set is not a measurement.")
    if args.strict and flagged:
        sys.exit(1)


if __name__ == "__main__":
    main()