"""Script check of paraphrase questions: shared salient tokens with the gold chunk text.

Threshold (frozen in params.json): at most MAX_SHARED shared tokens with any single gold-page chunk
(lowercase, split on non-alphanumerics, stopwords and 1-2 letter tokens removed). Also verifies that every
keyword_exact query appears verbatim in at least one chunk of a gold URL.
"""
import json, re, sys
MAX_SHARED = 2
STOP = set("""a an the and or of to in on for with from by at as is are was be been it its this that these those i my me we our you your
do does did how what which when where why who whom not no can could should would will may might if then than so but about into onto over
up down out off again once only own same too very s t just don have has had having there here all any both each few more most other some
such get got set use used using via per""".split())

def toks(s):
    return {t for t in re.findall(r"[a-z0-9]+", s.lower()) if t not in STOP and len(t) > 2}

def main(path, chunks="../data/chunks.jsonl"):
    rows = [json.loads(l) for l in open(chunks)]
    items = json.load(open(path))
    out, bad = [], 0
    for it in items:
        gold = [r for r in rows if r["url"] in it["gold_urls"]]
        if not gold:
            print("NO GOLD CHUNKS", it); bad += 1; continue
        if it["type"] == "keyword_exact":
            ok = any(it["query"] in r["text"] for r in gold)
            out.append({"id": it["id"], "type": it["type"], "token_in_gold_chunk": ok})
            bad += (not ok)
        elif it["type"] == "paraphrase":
            q = toks(it["query"])
            best = max(gold, key=lambda r: len(q & toks(r["heading"] + " " + r["text"])))
            sh = sorted(q & toks(best["heading"] + " " + best["text"]))
            out.append({"id": it["id"], "type": it["type"], "query": it["query"], "shared": sh, "n_shared": len(sh), "ok": len(sh) <= MAX_SHARED})
            bad += (len(sh) > MAX_SHARED)
        else:
            q = toks(it["query"].split()[-1])
            out.append({"id": it["id"], "type": it["type"], "token_in_gold_chunk": any(re.search(r"[a-z]+_[a-z_]+", it["query"]) and m in r["text"] for r in gold for m in re.findall(r"[A-Za-z]+_[A-Za-z_]+", it["query"]))})
    json.dump({"max_shared": MAX_SHARED, "results": out}, open("overlap.json", "w"), indent=1)
    for o in out:
        if o["type"] == "paraphrase": print(o["id"], o["n_shared"], o["shared"], "OK" if o["ok"] else "REWRITE")
        elif not o["token_in_gold_chunk"]: print(o["id"], o["type"], "TOKEN NOT IN GOLD CHUNK")
    print("problems:", bad)

if __name__ == "__main__":
    main(sys.argv[1] if len(sys.argv) > 1 else "eval_set_draft.json")