"""Second, stricter linking pass over the corpus.

Adds two relations to the clusters of build_corpus.py and removes from test/val every pair
whose extended cluster also contains a training pair:
  (e) character 5-gram containment: >= 70% of the 5-grams of the shorter space-free
      transliteration (at least 12 grams) occur in the longer one;
  (f) near-identical translations in the same language (Indel similarity >= 95, at least
      40 characters and 7 words) whose transliterations are at least 50% similar.
Writes data/leakfree.json with the kept test/val ids and link statistics.
"""
import json
from collections import Counter, defaultdict

import numpy as np
from rapidfuzz import fuzz, process

ROOT = "/Users/humonen/runic_access"
C = [json.loads(l) for l in open(f"{ROOT}/data/corpus.jsonl")]
n = len(C)
parent = list(range(n))


def find(i):
    while parent[i] != i:
        parent[i] = parent[parent[i]]
        i = parent[i]
    return i


def union(i, j):
    a, b = find(i), find(j)
    if a != b:
        parent[b] = a


first = {}
for i, r in enumerate(C):
    c = r["cluster"]
    if c in first:
        union(first[c], i)
    else:
        first[c] = i

ns = [r["norm"].replace(" ", "") for r in C]

# (e) char 5-gram containment
grams = [set(s[k:k + 5] for k in range(len(s) - 4)) for s in ns]
inv = defaultdict(set)
for i, g in enumerate(grams):
    for t in g:
        inv[t].add(i)
e_links = 0
for i, g in enumerate(grams):
    if len(g) < 12:
        continue
    cand = Counter()
    for t in g:
        if len(inv[t]) > 300:
            continue
        for j in inv[t]:
            if j != i and len(grams[j]) >= len(g):
                cand[j] += 1
    for j, c in cand.items():
        if c / len(g) >= 0.7:
            if find(i) != find(j):
                e_links += 1
            union(i, j)
print("(e) new links:", e_links)

# (f) near-identical translations
f_links = 0
for lang in ("en", "sv"):
    idx = [i for i, r in enumerate(C) if r["lang"] == lang and len(r["translation"]) >= 40
           and len(r["translation"].split()) >= 7]
    texts = [C[i]["translation"] for i in idx]
    M = process.cdist(texts, texts, scorer=fuzz.ratio, score_cutoff=95, workers=-1, dtype=np.uint8)
    ii, jj = np.nonzero(M)
    for a, b in zip(ii, jj):
        if a >= b:
            continue
        i, j = idx[a], idx[b]
        if fuzz.ratio(ns[i], ns[j]) >= 50:
            if find(i) != find(j):
                f_links += 1
            union(i, j)
print("(f) new links:", f_links)

root_splits = defaultdict(set)
for i, r in enumerate(C):
    root_splits[find(i)].add(r["split"])
keep = {"test": [], "val": []}
dropped = Counter()
for i, r in enumerate(C):
    if r["split"] in keep:
        if "train" in root_splits[find(i)]:
            dropped[r["split"], r["lang"]] += 1
        else:
            keep[r["split"]].append(r["id"])
kept = Counter((C[i]["split"], C[i]["lang"]) for i in range(n) if C[i]["id"] in set(keep["test"] + keep["val"]))
print("dropped:", dict(dropped))
print("kept:", dict(kept))
json.dump({"keep_test": keep["test"], "keep_val": keep["val"], "dropped": {f"{k[0]}_{k[1]}": v for k, v in dropped.items()},
           "links_e": e_links, "links_f": f_links}, open(f"{ROOT}/data/leakfree.json", "w"))