"""Turn raw teacher output into a training set, and say what was thrown away and why.

Purpose: the filtering stage, which is the stage that decides whether the student
    inherits the teacher's ability or the teacher's mistakes. Applies length,
    format and refusal checks, drops degenerate repetition, removes exact and
    near-duplicates, decontaminates against the Part 10 evaluation set, splits
    what survives, and writes both training shapes. Every rejected example is
    counted by reason and a sample of each reason is written out, because a
    filter you cannot inspect is a filter you cannot trust.
Platform: all (standard library only; no model, no accelerator, no network)
Minimum memory: 12 GB nominally, and far less in practice: this is text in memory
Assumes: Python 3.10 or newer. The input is the JSON Lines file written by
    generate-teacher-data.py. --tasks, when given, is the Part 10 evaluation set:
    a JSON object with a "tasks" list whose items carry "prompt" and "reference".
    distillog.py sits next to this file.

Usage: python3 filter-and-dedupe.py --raw raw/teacher.jsonl --out-dir . \\
           --tasks my-tasks.json --labbook labbook.md
       python3 filter-and-dedupe.py --raw raw/teacher.jsonl --out-dir . \\
           --min-words 12 --max-words 400 --near-duplicate 0.8 --keep-thinking

Written under --out-dir:
  data/{train,valid}.jsonl       TRL conversational prompt-completion
  data-mlx/{train,valid}.jsonl   mlx-lm completions, two plain strings
  filter-report.json             counts by reason, plus five examples of each
"""

from __future__ import annotations

import argparse
import json
import random
import re
from collections import Counter, defaultdict
from pathlib import Path
from typing import Any

import distillog

WORD_RE = re.compile(r"[a-z0-9]+")
THINK_BLOCK = re.compile(r"<think>.*?</think>\s*", re.DOTALL)
UNCLOSED_THINK = re.compile(r"<think>(?!.*</think>)", re.DOTALL)

# A refusal answer that invents a hostname, a port or a count is the failure the
# refusal category exists to catch. These are deliberately blunt: the point is to
# notice the teacher inventing specifics, and a false positive costs one example.
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+"
    r"|(?:happened|occurred)\s+\d+\s+times?)\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|no way (?:for me )?to know)\b",
    re.IGNORECASE,
)


def normalise(text: str) -> list[str]:
    """Lowercase, drop punctuation, split on words, so formatting differences do not hide a duplicate."""
    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:
    """Fraction of a's n-grams that also appear in b. Asymmetric on purpose."""
    return len(a & b) / len(a) if a else 0.0


def looks_repetitive(words: list[str], n: int = 8) -> bool:
    """True when the same 8-gram appears three or more times: a decoding loop.

    A model that falls into a loop produces long, fluent, useless text, and a
    student trained on it learns to loop. Length filters miss this because the
    output is often within the length budget.
    """
    if len(words) < n * 3:
        return False
    counts = Counter(tuple(words[i:i + n]) for i in range(len(words) - n + 1))
    return counts.most_common(1)[0][1] >= 3


def json_is_valid(text: str) -> bool:
    stripped = text.strip()
    if stripped.startswith("```"):
        return False  # a fenced block is not the object the prompt asked for
    try:
        json.loads(stripped)
    except json.JSONDecodeError:
        return False
    return True


def check(row: dict, args: argparse.Namespace) -> str | None:
    """Return the reason to reject this row, or None to keep it."""
    completion = (row.get("completion") or "").strip()
    if not completion:
        return "empty"

    if UNCLOSED_THINK.search(completion):
        return "unclosed-thinking-block"

    body = completion if args.keep_thinking else THINK_BLOCK.sub("", completion).strip()
    if not body:
        return "thinking-only"

    # The category checks come before the length checks on purpose. A refusal that
    # invents a hostname is usually short, and reporting it as "too short" would
    # hide the fault that matters: the teacher answered a question it could not know.
    category = row.get("category")
    if category == "format":
        wants_json = '"symptom"' in row.get("prompt", "") or "JSON object" in row.get("prompt", "")
        if wants_json and not json_is_valid(body):
            return "format-not-json"
        if not wants_json and "|" not in body:
            return "format-not-a-table"
    if category == "refusal":
        if INVENTED_SPECIFIC.search(body):
            return "refusal-invented-a-specific"
        if not REFUSAL_MARKERS.search(body):
            return "refusal-did-not-refuse"

    words = normalise(body)
    if len(words) < args.min_words:
        return "too-short"
    if len(words) > args.max_words:
        return "too-long"
    if looks_repetitive(words):
        return "repetition-loop"

    return None


def main() -> None:
    parser = argparse.ArgumentParser(description=__doc__.splitlines()[0])
    parser.add_argument("--raw", required=True, help="JSON Lines from generate-teacher-data.py")
    parser.add_argument("--out-dir", default=".")
    parser.add_argument("--tasks", default=None,
                        help="the Part 10 evaluation set, for decontamination")
    parser.add_argument("--min-words", type=int, default=8)
    parser.add_argument("--max-words", type=int, default=400)
    parser.add_argument("--near-duplicate", type=float, default=0.8,
                        help="drop a completion whose 13-gram containment in an earlier one is above this")
    parser.add_argument("--contamination", type=float, default=0.5,
                        help="drop an example whose prompt overlaps an evaluation prompt above this")
    parser.add_argument("--n", type=int, default=13, help="n-gram size for the overlap checks")
    parser.add_argument("--keep-thinking", action="store_true",
                        help="keep <think> blocks in the completion instead of stripping them")
    parser.add_argument("--valid-fraction", type=float, default=0.1)
    parser.add_argument("--seed", type=int, default=0)
    parser.add_argument("--labbook", default=None)
    args = parser.parse_args()

    raw_path = Path(args.raw)
    rows = [json.loads(line) for line in raw_path.read_text(encoding="utf-8").splitlines() if line.strip()]
    print(f"read {len(rows)} raw answer(s) from {raw_path}")

    rejected: Counter[str] = Counter()
    samples: dict[str, list[dict]] = defaultdict(list)

    def reject(row: dict, reason: str) -> None:
        rejected[reason] += 1
        if len(samples[reason]) < 5:
            samples[reason].append({
                "id": row.get("id"),
                "category": row.get("category"),
                "prompt": (row.get("prompt") or "")[:200],
                "completion": (row.get("completion") or "")[:300],
            })

    # Stage 1: per-example checks.
    stage1 = []
    for row in rows:
        reason = check(row, args)
        if reason:
            reject(row, reason)
            continue
        completion = (row["completion"] or "").strip()
        if not args.keep_thinking:
            completion = THINK_BLOCK.sub("", completion).strip()
        stage1.append({**row, "completion": completion})

    # Stage 2: exact and near-duplicate completions, and repeated prompts.
    kept: list[dict] = []
    seen_exact: set[str] = set()
    seen_prompt: set[str] = set()
    seen_grams: list[set[tuple[str, ...]]] = []
    for row in stage1:
        key = " ".join(normalise(row["completion"]))
        if key in seen_exact:
            reject(row, "duplicate-completion")
            continue
        if row["prompt"] in seen_prompt:
            # More than one sample per prompt survived the filters; keep the first
            # so a single prompt cannot dominate the training set.
            reject(row, "duplicate-prompt")
            continue
        grams = ngrams(normalise(row["completion"]), args.n)
        if any(containment(grams, earlier) >= args.near_duplicate for earlier in seen_grams):
            reject(row, "near-duplicate")
            continue
        seen_exact.add(key)
        seen_prompt.add(row["prompt"])
        seen_grams.append(grams)
        kept.append(row)

    # Stage 3: decontamination against the evaluation set.
    contaminated = 0
    if args.tasks:
        spec = json.loads(Path(args.tasks).read_text(encoding="utf-8"))
        eval_texts = []
        for task in spec.get("tasks", []):
            eval_texts.append(task.get("prompt", ""))
            if task.get("reference"):
                eval_texts.append(task["reference"])
        eval_grams = [ngrams(normalise(t), args.n) for t in eval_texts if t]
        eval_exact = {" ".join(normalise(t)) for t in eval_texts if t}
        clean = []
        for row in kept:
            text = f"{row['prompt']}\n{row['completion']}"
            if " ".join(normalise(row["prompt"])) in eval_exact:
                reject(row, "contaminated-exact")
                contaminated += 1
                continue
            grams = ngrams(normalise(text), args.n)
            if any(containment(e, grams) >= args.contamination for e in eval_grams if e):
                reject(row, "contaminated-overlap")
                contaminated += 1
                continue
            clean.append(row)
        kept = clean
        print(f"decontamination: checked against {len(eval_texts)} evaluation text(s), "
              f"dropped {contaminated}")
    else:
        print("decontamination: SKIPPED, no --tasks given. A gain measured on an "
              "evaluation set you did not check is not a gain you can report.")

    if not kept:
        raise SystemExit("every example was rejected; look at filter-report.json before changing thresholds")

    rng = random.Random(args.seed)
    rng.shuffle(kept)
    cut = max(1, int(len(kept) * args.valid_fraction))
    valid, train = kept[:cut], kept[cut:]

    out = Path(args.out_dir)
    written = {}
    for name, part in (("train", train), ("valid", valid)):
        trl_path = out / "data" / f"{name}.jsonl"
        trl_path.parent.mkdir(parents=True, exist_ok=True)
        with trl_path.open("w", encoding="utf-8") as handle:
            for row in part:
                handle.write(json.dumps({
                    "prompt": [{"role": "user", "content": row["prompt"]}],
                    "completion": [{"role": "assistant", "content": row["completion"]}],
                }, ensure_ascii=False) + "\n")
        mlx_path = out / "data-mlx" / f"{name}.jsonl"
        mlx_path.parent.mkdir(parents=True, exist_ok=True)
        with mlx_path.open("w", encoding="utf-8") as handle:
            for row in part:
                handle.write(json.dumps({"prompt": row["prompt"], "completion": row["completion"]},
                                        ensure_ascii=False) + "\n")
        written[name] = {"count": len(part), "trl": str(trl_path), "mlx": str(mlx_path),
                         "sha256": distillog.file_sha256(trl_path)}

    report = {
        "raw": str(raw_path),
        "raw_sha256": distillog.file_sha256(raw_path),
        "raw_examples": len(rows),
        "kept": len(kept),
        "rejected_total": sum(rejected.values()),
        "rejected_by_reason": dict(rejected.most_common()),
        "examples_by_reason": {k: v for k, v in samples.items()},
        "thresholds": {
            "min_words": args.min_words, "max_words": args.max_words,
            "near_duplicate": args.near_duplicate, "contamination": args.contamination,
            "n": args.n, "keep_thinking": args.keep_thinking,
        },
        "splits": written,
    }
    report_path = out / "filter-report.json"
    report_path.write_text(json.dumps(report, indent=2, ensure_ascii=False), encoding="utf-8")

    print()
    print(f"kept {len(kept)} of {len(rows)} ({100 * len(kept) / len(rows):.1f}%)")
    for reason, count in rejected.most_common():
        print(f"  rejected {count:>5}  {reason}")
    print(f"train {len(train)}  valid {len(valid)}")
    print(f"report written to {report_path}: read the examples before you accept the counts")

    if args.labbook:
        rec = distillog.record(
            labbook=args.labbook,
            lab="part-15/filter-and-dedupe",
            stage="filter",
            teacher=(rows[0].get("teacher") if rows else None),
            student=None,
            dataset={
                "raw": str(raw_path),
                "raw_sha256": report["raw_sha256"],
                "raw_examples": len(rows),
                "kept": len(kept),
                "train": len(train),
                "valid": len(valid),
                "train_sha256": written["train"]["sha256"],
                "tasks": args.tasks,
                "contaminated_dropped": contaminated,
            },
            hyperparameters=report["thresholds"],
            seed=args.seed,
            scores={"rejected_by_reason": dict(rejected.most_common())},
            config_path=__file__,
            notes=None if args.tasks else "decontamination skipped: no evaluation set supplied",
        )
        print(f"recorded {rec['stage']} as {rec['run_id']} in {args.labbook}")


if __name__ == "__main__":
    main()