"""Build the inscription-level corpus and leakage-free splits.

Sources (all scraped by the thesis authors):
  SRD 2014 export (signum, transliteration, English)            -> target en
  RunesDB / RuneS (transliteration, translation, image URLs)    -> target en (langid-filtered)
  christerhamp.se Gamla runor (Anton's clean scrape, Swedish)   -> target sv

Rows are grouped into inscription clusters by (a) identical normalised transliteration,
(b) identical signum, (c) near-duplicate normalised transliteration (Indel ratio >= 90 for
strings of >= 20 chars) and (d) word-trigram containment >= 0.8 of the shorter string.
Splits are assigned per cluster, so no inscription appears in two splits.
"""
import json
import random
import re
import sys
from collections import Counter, defaultdict

import langid
import numpy as np
import pandas as pd
from rapidfuzz import fuzz, process

sys.path.insert(0, "/Users/humonen/runic_access/code")
from runic_norm import normalize  # noqa: E402

P = "/Users/humonen/runic_access/sources/repos/runic_ai/parsers/"
CH = "/Users/humonen/runic_access/sources/gdrive/Christerhamp/christerhamp_images/christerhamp_gamla_runor_with_filenames.csv"
OUT = "/Users/humonen/runic_access/data/"
SEED = 13


def base_signum(s):
    if not isinstance(s, str):
        return None
    s = re.split(r"[$§]", s)[0]
    s = re.sub(r"\s+", " ", s).strip()
    return s or None


def clean_translation(t):
    if not isinstance(t, str):
        return ""
    t = re.sub(r"\s+", " ", t).strip()
    return t


rows = []
srd = pd.read_csv(P + "srd_parallel.csv")
for r in srd.itertuples():
    rows.append(dict(src="srd", signum=base_signum(r.signum), translit_raw=r.transliteration,
                     translation=clean_translation(r.english), lang="en", image_urls="", page=""))
rdb = pd.read_csv(P + "runes_parsed_parallel_1_10k.csv")
for r in rdb.itertuples():
    rows.append(dict(src="rdb", signum=None, translit_raw=r.transliteration,
                     translation=clean_translation(r.translation), lang="en",
                     image_urls=r.images_urls if isinstance(r.images_urls, str) else "", page=r.url))
ch = pd.read_csv(CH)
for r in ch.itertuples():
    sig = r.stone_id if isinstance(r.stone_id, str) and re.match(r"^[A-ZÅÄÖÖÜa-zåäö]{1,4}\s", r.stone_id) else None
    rows.append(dict(src="ch", signum=base_signum(sig), translit_raw=r.transliteration,
                     translation=clean_translation(r.translation), lang="sv",
                     image_urls=r.image_url if isinstance(r.image_url, str) else "", page=r.page_url))
df = pd.DataFrame(rows)
df["norm"] = df.translit_raw.map(normalize)
n0 = len(df)

# translation must carry content
alpha = df.translation.map(lambda t: len(re.findall(r"[A-Za-zÀ-ÿþðÞÐǫ]", t)))
df = df[(df.norm.str.len() >= 2) & (alpha >= 3)].copy()

# RunesDB: keep English (or short name-like glosses, which langid cannot judge)
def rdb_keep(t):
    words = t.split()
    if len(words) <= 3:
        return True
    return langid.classify(t)[0] == "en"

mask_rdb = df.src.eq("rdb")
keep = df[mask_rdb].translation.map(rdb_keep)
df = pd.concat([df[~mask_rdb], df[mask_rdb][keep]])
df = df.reset_index(drop=True)
print(f"rows: raw {n0} -> kept {len(df)}", df.groupby(["src", "lang"]).size().to_dict())

# ---------- clustering ----------
parent = list(range(len(df)))


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


def union_groups(keys):
    first = {}
    for i, k in enumerate(keys):
        if k is None or k == "":
            continue
        if k in first:
            union(first[k], i)
        else:
            first[k] = i


df["nospace"] = df.norm.str.replace(" ", "", regex=False)
union_groups(df.nospace.tolist())
union_groups(df.signum.tolist())

uniq_ns = sorted(set(df.nospace))
idx_ns = defaultdict(list)
for i, s in enumerate(df.nospace):
    idx_ns[s].append(i)
long_ = [s for s in uniq_ns if len(s) >= 8]
M = process.cdist(long_, long_, scorer=fuzz.ratio, score_cutoff=90, workers=-1, dtype=np.uint8)
ii, jj = np.nonzero(M)
near_pairs = [(long_[a], long_[b], int(M[a, b])) for a, b in zip(ii, jj) if a < b]
for a, b, _ in near_pairs:
    union(idx_ns[a][0], idx_ns[b][0])
uniq = sorted(set(df.norm))
idx_of = defaultdict(list)
for i, s in enumerate(df.norm):
    idx_of[s].append(i)
print("near-dup pairs (ratio>=90 on space-free strings >= 8 chars):", len(near_pairs))


def trigrams(s):
    w = s.split()
    return {" ".join(w[k:k + 3]) for k in range(len(w) - 2)}

tg = {s: trigrams(s) for s in uniq}
inv = defaultdict(set)
for s, g in tg.items():
    for t in g:
        inv[t].add(s)
contain_pairs = 0
for s, g in tg.items():
    if len(g) < 3:
        continue
    cand = Counter()
    for t in g:
        if len(inv[t]) > 200:  # formula trigrams shared by hundreds of stones carry no identity
            continue
        for o in inv[t]:
            if o != s:
                cand[o] += 1
    for o, c in cand.items():
        if len(tg[o]) >= len(g) and c / len(g) >= 0.8:
            union(idx_of[s][0], idx_of[o][0])
            contain_pairs += 1
print("containment pairs:", contain_pairs)

df["cluster"] = [find(i) for i in range(len(df))]
sizes = df.cluster.value_counts()
print("clusters:", len(sizes), "largest:", sizes.head(5).tolist(), "singletons:", int((sizes == 1).sum()))

# save boundary examples for manual inspection
random.seed(SEED)
border = [p for p in near_pairs if p[2] < 93]
with open(OUT + "inspect_near_dup_pairs.txt", "w") as f:
    for a, b, sc in random.sample(border, min(40, len(border))):
        f.write(f"{sc}\n  {a}\n  {b}\n")
big = sizes.index[0]
with open(OUT + "inspect_largest_cluster.txt", "w") as f:
    for s in df[df.cluster == big].norm.head(60):
        f.write(s + "\n")

# ---------- split ----------
has_img = df.image_urls.str.len() > 0
cl_img = df[has_img].groupby("cluster").src.apply(set)
clusters = sorted(sizes.index)
rng = random.Random(SEED)
rng.shuffle(clusters)
# cap the test/val clusters at size <= 8 so a single formulaic mega-cluster cannot dominate
eligible = [c for c in clusters if sizes[c] <= 8]
img_ch = [c for c in eligible if c in cl_img.index and "ch" in cl_img[c]]
img_rdb = [c for c in eligible if c in cl_img.index and "rdb" in cl_img[c] and "ch" not in cl_img[c]]
test = set(img_ch[:220]) | set(img_rdb[:220])
rest = [c for c in eligible if c not in test]
target_test_rows = 1300
for c in rest:
    if df.cluster.isin(test).sum() >= target_test_rows:
        break
    test.add(c)
val = set()
for c in [c for c in eligible if c not in test]:
    if df.cluster.isin(val).sum() >= 600:
        break
    val.add(c)
df["split"] = np.where(df.cluster.isin(test), "test", np.where(df.cluster.isin(val), "val", "train"))
print(df.groupby(["split", "lang"]).size().unstack().to_string())
print("test rows with images:", df[(df.split == "test") & has_img].groupby("src").size().to_dict())

# leakage audit: max similarity of every test norm to any train norm
tr = sorted(set(df[df.split == "train"].nospace))
te = df[df.split == "test"].nospace.tolist()
best = process.cdist(te, tr, scorer=fuzz.ratio, workers=-1, dtype=np.uint8).max(axis=1)
print("test->train max ratio: >=90:", int((best >= 90).sum()), ">=80:", int((best >= 80).sum()), "of", len(te))
df.loc[df.split == "test", "max_ratio_to_train"] = best

df["id"] = [f"{s}{i:05d}" for i, s in zip(range(len(df)), df.src)]
df.to_json(OUT + "corpus.jsonl", orient="records", lines=True, force_ascii=False)
print("wrote", OUT + "corpus.jsonl")