# finetune/purity.py
"""Bengali script-purity quality gate for teacher labels.

The teacher (Gemma) writes mostly clean Bengali, but occasionally code-switches
(e.g. a stray Latin or Cyrillic word — we saw `зеленая` leak once). Distillation
caps the student at the label quality, so we filter/repair bad labels before
training. This is the gate referenced in the Bengali-quality investigation.

Heuristics only — no model needed. A human Bengali speaker should still spot-check
a sample of what passes.
"""

import re
import unicodedata

# Bengali Unicode block: U+0980–U+09FF.
_BENGALI = re.compile(r"[ঀ-৿]")
# Letters from other scripts that must NOT appear in a Bengali story body.
_LATIN = re.compile(r"[A-Za-z]")
_CYRILLIC = re.compile(r"[Ѐ-ӿ]")
# Allowed non-letter noise: digits, whitespace, common punctuation, emoji, danda.
_ALLOWED_NONLETTER = re.compile(r"[\s\d\.,!?…\"'“”‘’—\-–—:;()।॥☀-➿\U0001F300-\U0001FAFF]")


def script_stats(text: str) -> dict:
    """Counts of Bengali vs foreign letters and the Bengali-letter ratio."""
    text = unicodedata.normalize("NFC", text or "")
    bengali = len(_BENGALI.findall(text))
    latin = len(_LATIN.findall(text))
    cyrillic = len(_CYRILLIC.findall(text))
    letters = bengali + latin + cyrillic
    ratio = (bengali / letters) if letters else 0.0
    return {
        "bengali": bengali,
        "latin": latin,
        "cyrillic": cyrillic,
        "foreign": latin + cyrillic,
        "bengali_ratio": round(ratio, 4),
    }


def is_clean(
    text: str,
    min_bengali_ratio: float = 0.98,
    max_foreign_letters: int = 0,
    min_length: int = 40,
) -> tuple[bool, dict]:
    """Decide whether a teacher label is clean enough to train on.

    Defaults are strict: essentially zero foreign-script letters. Loosen
    max_foreign_letters to 1–2 if you'd rather repair than drop.

    Returns (ok, stats).
    """
    stats = script_stats(text)
    ok = (
        len((text or "").strip()) >= min_length
        and stats["foreign"] <= max_foreign_letters
        and stats["bengali_ratio"] >= min_bengali_ratio
    )
    return ok, stats


def foreign_words(text: str) -> list[str]:
    """Return whitespace tokens that contain any Latin/Cyrillic letter — useful
    for eyeballing exactly what leaked (e.g. ['зеленая'])."""
    out = []
    for tok in (text or "").split():
        if _LATIN.search(tok) or _CYRILLIC.search(tok):
            out.append(tok)
    return out


if __name__ == "__main__":
    samples = {
        "good": "আচ্ছা রূপা, চোখ বুজে নাও। চাঁদমামা হাসছে, পুকুরের ধারে ঘাস দুলছে। শুভরাত্রি।",
        "leak": "দেখছো зеленая গোল বলটা? ওটা রূপার প্রিয় খেলনা সবুজ আপেল!",
        "english": "Once upon a time there was a small red house under the sun.",
    }
    for name, s in samples.items():
        ok, stats = is_clean(s)
        print(f"{name:8} ok={ok!s:5} {stats} leaks={foreign_words(s)}")