#!/usr/bin/env python3
"""Construit les corpus lingala enrichis IN-DOMAIN pour KenLM.

Exclusions strictes : aucune phrase de devhard_lin (pour que devhard reste
utilisable plus tard) ni de la validation officielle (pour que la comparaison
des LM sur la validation ne soit pas circulaire).
"""
import json, re

LM = "/scratch/lm/"
M = "/scratch/prep/manifests/"


def key(s):
    return " ".join(re.sub(r"[^\w ]", " ", str(s).lower()).split())


def texts(path, field="text"):
    out = []
    for l in open(path, encoding="utf-8"):
        r = json.loads(l)
        t = (r.get(field) or "").strip()
        if t:
            out.append((key(t), t))
    return out


ban = set()
for p in ("/root/devhard/devhard_lin.jsonl", M + "waxal_lin_validation.jsonl"):
    ban |= set(k for k, _ in texts(p))
print("phrases interdites (devhard + validation) : %d" % len(ban))

base = [l.strip() for l in open(LM + "corpus_lin.txt", encoding="utf-8") if l.strip()]
kb = set(key(x) for x in base)
print("corpus de base : %d phrases" % len(base))
n_ban = sum(1 for x in base if key(x) in ban)
print("  dont interdites deja presentes : %d  (doit etre 0)" % n_ban)

for tag, src in (("ps", M + "pseudo_lin_raw.jsonl"), ("psf", M + "pseudo_lin_filtered.jsonl")):
    rows = texts(src)
    seen, add = set(kb), []
    skipped_ban = skipped_dup = 0
    for k, t in rows:
        if k in ban:
            skipped_ban += 1; continue
        if k in seen:
            skipped_dup += 1; continue
        seen.add(k); add.append(t)
    out = base + add
    with open(LM + "corpus_lin_%s.txt" % tag, "w", encoding="utf-8") as f:
        f.write("\n".join(out) + "\n")
    print("corpus_lin_%s.txt : %d phrases (%d de base + %d ajoutees ; %d ecartees interdites, %d doublons)"
          % (tag, len(out), len(base), len(add), skipped_ban, skipped_dup))