#!/usr/bin/env python3
# Semantic word golf from the Linguabase core tables.
#
# The game: a hole is two words, an origin and a goal. Each turn shows
# the current word's CLOUD -- its strongest associations after three
# filters -- and the player taps one to move there. Reach the goal in
# as few hops as possible. (Played at full scale, this is the iOS game
# In Other Words: https://inotherwords.app)
#
# This script cuts the course and counts every hole at par 3, then
# writes a small playable pack for a handful of holes.
#
# Needs these four files beside it (any mirror on the Downloads page):
#   vocabulary.core.parquet      word, core_rank (1 = most familiar)
#   associations.core.parquet    word -> ranked list of related words
#   scores_labels.core.parquet   word -> content_rating, sensitivity
#   word_families.core.parquet   word -> list of relatives
#
# The course spec (change these and every count changes with them):
#   RANK_CAP    30,000  every word on the course is inside this rank
#   CLOUD_SIZE  15      a cloud is the first 15 associates that pass
#   MIN_INBOUND 8       a goal must appear in at least 8 clouds
#
# Measured against the July 2026 core release, this spec yields:
#   400,000 -> 30,000 -> 28,332 after the filters -> 28,219 with full clouds
#   -> 13,801 goal candidates -> 41,934,946 origin-goal pairs at par 3
#
# Outputs:
#   holes.tsv   origin, goal, three-hop line count, first-hop spread
#               (one row per hole for the DEMO_ORIGINS below; drop the
#               origin filter to enumerate the full 41.9M -- it runs,
#               but the file is large)
#   pack.json   playable pack: per-hole clouds within two taps

import json
import collections
import pandas as pd

RANK_CAP    = 30_000
CLOUD_SIZE  = 15
MIN_INBOUND = 8
DEMO_ORIGINS = ["candle", "soap", "pencil", "wool", "kite"]

# ---- load the four tables (published files arrive unsorted) ----
v   = pd.read_parquet("vocabulary.core.parquet",
                      columns=["word", "core_rank"])
v   = v.sort_values("core_rank")
a   = pd.read_parquet("associations.core.parquet")
sc  = pd.read_parquet("scores_labels.core.parquet",
                      columns=["word", "content_rating", "sensitivity"])
fam = pd.read_parquet("word_families.core.parquet")

rank   = dict(zip(v.word, v.core_rank))
assoc  = dict(zip(a.word, a.associates))
scores = {w: (cr, s) for w, cr, s
          in zip(sc.word, sc.content_rating, sc.sensitivity)}
family = {w: set(x.lower() for x in rel)
          for w, rel in zip(fam.word, fam.relatives)}

# ---- the three filters, then the cloud cut ----
def passes_filters(w):
    """Inside the rank cap, rated G/PG, sensitivity above 3."""
    if rank.get(w, 10**9) > RANK_CAP:
        return False
    t = scores.get(w)
    return (t is not None and t[0] in ("G", "PG")
            and t[1] is not None and t[1] > 3)

def cloud_of(w):
    """First CLOUD_SIZE associates that pass the filters and are not
    relatives of w (checked both directions, case-blind)."""
    out, seen = [], {w.lower()}
    wfam = family.get(w, set())
    for x in assoc.get(w, []):
        xl = x.lower()
        if xl in seen or not passes_filters(x):
            continue
        if xl in wfam or w.lower() in family.get(x, set()):
            continue
        seen.add(xl)
        out.append(x)
        if len(out) == CLOUD_SIZE:
            break
    return out

print("cutting clouds...")
play  = [w for w in v.word if passes_filters(w)]
cloud = {}
for w in play:
    c = cloud_of(w)
    if len(c) == CLOUD_SIZE:          # a full cloud makes w a node
        cloud[w] = c
nodes = set(cloud)

inbound = collections.Counter()
for w in cloud:
    for x in cloud[w]:
        inbound[x] += 1
goals = set(w for w in play if inbound[w] >= MIN_INBOUND)
print(f"{len(cloud):,} words carry a full cloud; "
      f"{len(goals):,} qualify as goals")

# ---- par-3 holes: goal first reachable at exactly three hops ----
def par3_goals(o):
    d1 = set(cloud[o])
    d2 = set()
    for x in cloud[o]:
        if x in nodes:
            d2.update(cloud[x])
    d2 -= d1
    d2.discard(o)
    d3 = set()
    for x in d2:
        if x in nodes:
            d3.update(cloud[x])
    d3 -= d2
    d3 -= d1
    d3.discard(o)
    return d3 & goals

def hole_stats(o, g):
    """Three-hop line count and which first taps can still win."""
    total, first = 0, []
    for a_ in cloud[o]:
        if a_ not in nodes:
            continue
        c = sum(1 for b in cloud[a_]
                if b in nodes and b not in (o, g) and g in cloud[b])
        if c:
            first.append(a_)
            total += c
    return total, first

print("counting holes for the demo origins...")
with open("holes.tsv", "w") as f:
    f.write("origin\tgoal\tthree_hop_lines\tfirst_hop_spread\n")
    for o in DEMO_ORIGINS:
        for g in sorted(par3_goals(o), key=lambda g: rank[g]):
            p, first = hole_stats(o, g)
            f.write(f"{o}\t{g}\t{p}\t{len(first)}\n")

# ---- a playable pack: every cloud within two taps of each origin ----
pack = {}
for o in DEMO_ORIGINS:
    need = [o] + [x for x in cloud[o] if x in nodes]
    lvl2 = set()
    for a_ in need[1:]:
        lvl2.update(x for x in cloud[a_] if x in nodes)
    need += sorted(lvl2 - set(need), key=lambda w: rank[w])
    pack[o] = {w: cloud[w] for w in need}
json.dump(pack, open("pack.json", "w"))
print("wrote holes.tsv and pack.json")