"""Keep new training data away from the evaluation sets.

Two rules:
1. Task families used only for zero-shot evaluation (data/heldout.jsonl: trec question types,
   formality, spam) must never be trained on. `blocked_task(name)` checks dataset/task names.
2. No training state may overlap an eval state (heldout, test, bench): exact match after
   normalization, or any shared 12-word shingle. `is_contaminated(state)` checks this.

    from jevlike.data.decontam import Decontaminator
    dc = Decontaminator()                       # loads the eval sets once (~1 s)
    if dc.blocked_task(task_name) or dc.is_contaminated(state): skip

CLI: python -m jevlike.data.decontam IN.jsonl OUT.jsonl   (filters a Record JSONL)
"""
from __future__ import annotations

import re
import sys
from pathlib import Path

from jevlike.types import read_jsonl, write_jsonl

ROOT = Path(__file__).resolve().parents[2]
EVAL_FILES = ["data/heldout.jsonl", "data/test.jsonl", "data/bench/decisions.jsonl",
              "data/eval_newsfin/zeroshot.jsonl", "data/eval_newsfin/zeroshot_long.jsonl"]
BLOCKED_PATTERNS = [r"\btrec\b", r"question[\s_-]*(type|classification)", r"formal", r"spam", r"enron"]
SHINGLE = 12


def _norm(s: str) -> str:
    return re.sub(r"\W+", " ", s.lower()).strip()


def _shingles(s: str) -> set[str]:
    w = _norm(s).split()
    if len(w) < SHINGLE:
        return {" ".join(w)}
    return {" ".join(w[i:i + SHINGLE]) for i in range(len(w) - SHINGLE + 1)}


class Decontaminator:
    def __init__(self, eval_files: list[str] = EVAL_FILES):
        self.exact: set[str] = set()
        self.shingles: set[str] = set()
        for f in eval_files:
            p = ROOT / f
            if not p.exists():
                continue
            for r in read_jsonl(str(p)):
                self.exact.add(_norm(r.state))
                self.shingles |= _shingles(r.state)

    @staticmethod
    def blocked_task(name: str) -> bool:
        """True for dataset/task names in the held-out task families (match on name or description)."""
        return any(re.search(p, name.lower()) for p in BLOCKED_PATTERNS)

    def is_contaminated(self, state: str) -> bool:
        return _norm(state) in self.exact or not self.shingles.isdisjoint(_shingles(state))


if __name__ == "__main__":
    src, dst = sys.argv[1], sys.argv[2]
    dc = Decontaminator()
    rs = read_jsonl(src)
    keep = [r for r in rs if not (dc.blocked_task(r.source) or dc.is_contaminated(r.state))]
    write_jsonl(dst, keep)
    print(f"kept {len(keep)}/{len(rs)} ({len(rs) - len(keep)} dropped)")