"""Build the request stream from TAB (public train split, MIT): one request = one sentence of an ECHR judgment.

Documents are shuffled (seed 0) and split by document: HOLDOUT_DOCS documents give the spot-check holdout, the rest give
the stream. Gold spans (TAB annotator 1, the three format types) are kept for reference only; the demo learns from the
teacher, never from gold.
usage: prep.py TAB.parquet OUT_DIR
"""
import json
import random
import re
import sys

import pyarrow.parquet as pq

TYPES = {"DATETIME", "CODE", "QUANTITY"}
N_STREAM, N_HOLDOUT, HOLDOUT_DOCS = 3000, 400, 120
ABBR = r"(?<!\bMr)(?<!\bMrs)(?<!\bMs)(?<!\bDr)(?<!\bno)(?<!\bNo)(?<!\bnos)(?<!\bNos)(?<!\bpp)(?<!\bp)(?<!\bpara)(?<!\bparas)(?<!\bArt)(?<!\bv)(?<!\bSt)(?<!\bJr)(?<!\bi\.e)(?<!\be\.g)(?<!\bseq)(?<!\bop)(?<!\bcit)(?<!\b[A-Z])"
SENT_END = re.compile(ABBR + r"[.;](?=\s+[A-Z(\"'])")


def sentences(text):
    out = []
    for m in re.finditer(r"[^\n]+", text):
        start, para = m.start(), m.group()
        s = 0
        for e in SENT_END.finditer(para):
            out.append((start + s, start + e.end()))
            s = e.end()
            while s < len(para) and para[s] == " ":
                s += 1
        if s < len(para):
            out.append((start + s, start + len(para)))
    return out


def mentions(ann):
    for v in (ann or {}).values():
        if isinstance(v, dict) and v.get("entity_mentions"):
            return v["entity_mentions"]
    return []


def main(pqf, out):
    docs = pq.read_table(pqf).to_pylist()
    random.Random(0).shuffle(docs)
    rows = {"holdout": [], "stream": []}
    for k, d in enumerate(docs):
        part = "holdout" if k < HOLDOUT_DOCS else "stream"
        ann = d["annotations"]
        if isinstance(ann, str):
            ann = json.loads(ann)
        ents = [(m["start_offset"], m["end_offset"], m["entity_type"]) for m in mentions(ann) if m.get("entity_type") in TYPES]
        for a, b in sentences(d["text"]):
            t = d["text"][a:b].strip()
            if not (40 <= len(t) <= 400) or t.isupper():
                continue
            a2 = d["text"].index(t, a)
            gold = sorted({(s - a2, e - a2, ty) for s, e, ty in ents if s >= a2 and e <= a2 + len(t)})
            rows[part].append({"doc_id": d["doc_id"], "text": t,
                               "gold": [{"start": s, "end": e, "type": ty, "text": t[s:e]} for s, e, ty in gold]})
    rng = random.Random(1)
    sel = {"stream": rng.sample(rows["stream"], N_STREAM), "holdout": rng.sample(rows["holdout"], N_HOLDOUT)}
    with open(f"{out}/requests.jsonl", "w") as f:
        for part, rs in sel.items():
            for i, r in enumerate(rs):
                f.write(json.dumps({"id": f"{part[0]}{i:04d}", "part": part, **r}) + "\n")
    for part, rs in sel.items():
        print(part, len(rs), "of", len(rows[part]), "docs", len({r["doc_id"] for r in rs}),
              "with gold", sum(1 for r in rs if r["gold"]), "gold spans", sum(len(r["gold"]) for r in rs))


if __name__ == "__main__":
    main(sys.argv[1], sys.argv[2])