# linguabase-dictionary-bundle.py -- build per-word dictionary bundles
# from the Linguabase core release (CC0, linguabase.org/download.html).
#
# Multi-word entries are first-class: "ice cream" and "cold shoulder"
# head entries exactly the way "key" does, and phrase forms fold into
# their base ("ice creams" -> "ice cream") the same as "dogs" -> "dog".
#
# Reads six tables (parquet, in the working directory):
#   vocabulary.core.parquet          word, core_rank, ...
#   definitions.core.parquet         word, definition (one paragraph)
#   senses.core.parquet              word, sense_labels (list)
#   sense_associations.core.parquet  word, sense, associates (list)
#   associations.core.parquet        word, associates (ranked list)
#   word_families.core.parquet       word, relatives (list)
#
# Writes:
#   dictionary.json   one bundle per headword:
#                     {w, r, d, s:[[sense,[related...]]...], f:[kin...]}
#   crossrefs.tsv     folded_form <tab> headword  (DOGS -> DOG)
#
# The two numbers you might edit:
CAP = 50_000        # familiarity depth: entries come from the top CAP
MIN_SENSES = 0      # 0 = every head gets a bundle; 2 = sense-split only

import pandas as pd, json

# ---- load ------------------------------------------------------------
# The shipped files are NOT sorted by rank (and the TSVs carry a header
# row) -- always sort on core_rank before taking a prefix.
v    = pd.read_parquet('vocabulary.core.parquet').sort_values('core_rank')
defs = pd.read_parquet('definitions.core.parquet').set_index('word').definition
sen  = pd.read_parquet('senses.core.parquet').set_index('word').sense_labels
fam  = pd.read_parquet('word_families.core.parquet').set_index('word').relatives
sa   = pd.read_parquet('sense_associations.core.parquet')

rank = dict(zip(v.word, v.core_rank))
kin  = lambda w: list(fam.get(w, []))

# ---- 1 · the headword list ------------------------------------------
words = v[v.core_rank <= CAP]
words = words[words.word == words.word.str.lower()]     # no capitals
words = words[words.word.str.fullmatch(                 # letters, spaces,
    r"([^\W\d_]|[ -])+")]                               # hyphens: café and
                                                        # ice cream stay;
                                                        # don't, 24/7 go

# ---- 2 · fold derived forms into their base -------------------------
# DOGS, SMILED and EASIER become cross-references, not entries: a word
# folds when its family holds a better-ranked base it derives from by
# suffix, allowing e-drop (smile+ed), y->i (easy+er) and doubling
# (run+n+ing) spellings. Phrase forms fold the same way whenever the
# family lists their base: ICE CREAMS folds into ICE CREAM.
SFX = 's es d ed ing er ers est ly'.split()
def derived(w, b):                    # is w = b + a suffix?
    if not b: return False
    forms = [b, b + b[-1]]
    if b[-1] == 'e': forms.append(b[:-1])
    if b[-1] == 'y': forms.append(b[:-1] + 'i')
    return any(w == f + s for f in forms for s in SFX)

def fold_base(w):                     # best-ranked base, or None
    cands = [b for b in kin(w)
             if rank.get(b, 9e9) < rank[w] and derived(w, b)]
    return min(cands, key=rank.__getitem__) if cands else None

base = {w: fold_base(w) for w in words.word}
heads = [w for w in words.word if base[w] is None]

# ---- 3 · sense sections, family-filtered ----------------------------
# senses.sense_labels supplies each entry's section headings; the
# same-named row of sense_associations supplies the words under each.
# (sense_associations holds many MORE rows per word -- compound and
# phrase contexts. Select through the labels, don't take them all.)
# The family filter: a related word is dropped when it is the headword
# itself, in the headword's family, or a phrase containing the
# headword -- BARK never lists BARKING, or "dog bark". The containment
# check is a token-sequence match so phrase headwords are caught inside
# longer phrases too.
sa_idx = sa.set_index(['word', 'sense']).associates

def leaks(w, a):
    toks = a.replace('-', ' ').replace('_', ' ').lower().split()
    return a == w or a in kin(w) or f' {w} ' in f' {" ".join(toks)} '

def sections(w):
    out = []
    for label in list(sen.get(w, [])):
        try: cloud = list(sa_idx.at[(w, label)])
        except KeyError: continue     # label without a cloud: skip
        kept = [a.replace('_', ' ') for a in cloud if not leaks(w, a)]
        if kept: out.append([label, kept])
    return out

# ---- 4 · write the bundles ------------------------------------------
bundles = {}
for w in heads:
    s = sections(w)
    if len(s) < MIN_SENSES: continue
    bundles[w] = {'w': w, 'r': int(rank[w]),
                  'd': defs.get(w, ''),      # '' when no paragraph yet
                  's': s, 'f': kin(w)}

with open('dictionary.json', 'w') as f:
    json.dump(bundles, f, ensure_ascii=False, indent=1)

with open('crossrefs.tsv', 'w') as f:
    f.write('folded_form\theadword\n')
    for w in words.word:
        if base[w]: f.write(f'{w}\t{base[w]}\n')

n_split = sum(1 for b in bundles.values() if len(b['s']) >= 2)
print(f'{len(bundles):,} bundles ({n_split:,} sense-split), '
      f'{sum(1 for w in words.word if base[w]):,} cross-references')