#!/usr/bin/env python3
"""Build HANTAI boards from the Linguabase antonyms + vocabulary tables.

HANTAI (the game on demo-7-antonyms.html): word tiles on a grid; the
player swaps tiles until no two opposites are orthogonally adjacent.
This script is the whole build in one file:

  1. HARVEST mutual pairs — a pair {a, b} qualifies only if the table
     lists b as a bare-word opposite of a AND a as a bare-word opposite
     of b (any axis), with a strength rating of 6+ in at least one
     direction, both words single, lowercase, and inside the 20,000
     most familiar. One-way listings are the data's stretches
     (hot -> old exists; old -> hot does not) — the mutual join drops
     them: 47,140 one-way-eligible pairs -> 21,105 mutual.
  2. GROW boards — dense connected patches of the pair network (every
     tile has at least one opposite on the board; most have 2+), then
     prove each solvable by backtracking search.
  3. GATE by simulated difficulty — a board ships only if the dumbest
     policy (swap a random clashing tile with a random calm one) runs
     past a 300-move cap, and greedy best-swap stalls in local minima
     often while still winning often enough to stay humanly fair.
  4. PACK — no two shipped boards share more than 2 pairs.

Needs beside it:  vocabulary.core.parquet, antonyms.core.parquet
Writes:           hantai_boards.json  (50 boards: grid, words, pairs,
                  shuffled start, one verified solution, difficulty stats)

Run:  python3 linguabase-hantai-boards.py        (~4 minutes)
"""
import json, random, time, statistics
import pandas as pd

RANK_CAP     = 20_000   # familiarity cap, both sides of every pair
MIN_STRENGTH = 6        # required in at least one listing direction
MAX_LEN      = 9        # tile-friendly words
MOVE_CAP     = 300      # simulated-policy move budget
N_BOARDS     = 50
rng = random.Random(20260813)

STOP = set("""the and to of in a is that it for on with as at by an be this have
from or had not are but were which you all she he his her they them we our your
its was will would could should can may might must did does do been being who
whom what when where why how than then there here also just only very much some
any each other more most such no nor so too own same""".split())

# ---------------------------------------------------------------- 1. harvest
vocab = pd.read_parquet("vocabulary.core.parquet", columns=["word", "core_rank"])
fam = {w: r for w, r in zip(vocab.word, vocab.core_rank)
       if r <= RANK_CAP and " " not in w}
ant = pd.read_parquet("antonyms.core.parquet")

def clean(w):
    return (w.isalpha() and w.islower() and 3 <= len(w) <= MAX_LEN
            and w in fam and w not in STOP)

# directed[(a, b)] = best strength with which b appears as a bare opposite
# in any row headed (case-insensitively) by a
directed = {}
for word, opps in zip(ant.word, ant.opposites):
    if opps is None:
        continue
    wl = word.lower()
    for o in opps:
        name, _, st = o.rpartition(":")
        try:
            st = int(st)
        except ValueError:
            continue
        nl = name.lower()
        if nl != wl:
            key = (wl, nl)
            if st > directed.get(key, 0):
                directed[key] = st

pairs = set()
for (a, b), st in directed.items():
    if not (clean(a) and clean(b)):
        continue
    # the mutual screen: (b, a) must exist too — this is the line that
    # removes the stretches
    if max(st, directed.get((b, a), 0)) >= MIN_STRENGTH and (b, a) in directed:
        pairs.add(frozenset((a, b)))
assert len(pairs) == 21_105, f"mutual pool drifted: {len(pairs)}"
print(f"mutual pool: {len(pairs)} pairs")

adj = {}
for p in pairs:
    a, b = tuple(p)
    adj.setdefault(a, set()).add(b)
    adj.setdefault(b, set()).add(a)

def same_rootish(a, b):
    """Block near-duplicates (rainy/raining, billion/billions)."""
    if a.startswith(b) or b.startswith(a):
        return True
    n = 0
    for x, y in zip(a, b):
        if x != y:
            break
        n += 1
    return n >= 5

# ---------------------------------------------------------------- helpers
def board_pairs(words):
    ws = set(words)
    return sorted(tuple(sorted(p)) for p in pairs if p <= ws)

def conflicts(order, R, C, pset):
    out = []
    for i, w in enumerate(order):
        r, c = divmod(i, C)
        if c + 1 < C and frozenset((w, order[i + 1])) in pset:
            out.append((i, i + 1))
        if r + 1 < R and frozenset((w, order[i + C])) in pset:
            out.append((i, i + C))
    return out

def solve(words, R, C, pset, budget=150_000):
    """Backtracking placement; None = no solution inside the node budget."""
    n = R * C
    order, used = [None] * n, [False] * len(words)
    ws = sorted(words, key=lambda w: -len(adj.get(w, set()) & set(words)))
    nodes = [0]
    def bt(i):
        nodes[0] += 1
        if nodes[0] > budget:
            return False
        if i == n:
            return True
        r, c = divmod(i, C)
        for k, w in enumerate(ws):
            if used[k]:
                continue
            if c > 0 and frozenset((w, order[i - 1])) in pset:
                continue
            if r > 0 and frozenset((w, order[i - C])) in pset:
                continue
            used[k] = True; order[i] = w
            if bt(i + 1):
                return True
            used[k] = False; order[i] = None
        return False
    return list(order) if bt(0) else None

# ---------------------------------------------------------------- 3. the gates
def policy_A(board, rng):
    """The dumbest policy: random clashing tile <-> random calm tile."""
    R, C = board["rows"], board["cols"]
    pset = {frozenset(p) for p in board["pairs"]}
    order = board["start"][:]
    for move in range(1, MOVE_CAP + 1):
        conf = conflicts(order, R, C, pset)
        if not conf:
            return move - 1, True
        bad = sorted({i for p in conf for i in p})
        good = [i for i in range(len(order)) if i not in set(bad)]
        i = rng.choice(bad)
        j = rng.choice(good) if good else rng.choice([x for x in bad if x != i])
        order[i], order[j] = order[j], order[i]
    return MOVE_CAP, False

def policy_B(board, rng):
    """Greedy best-swap; stalls when no swap strictly reduces clashes."""
    R, C = board["rows"], board["cols"]
    pset = {frozenset(p) for p in board["pairs"]}
    order = board["start"][:]
    n = len(order)
    for move in range(1, MOVE_CAP + 1):
        cur = len(conflicts(order, R, C, pset))
        if cur == 0:
            return move - 1, True, False
        best, choices = cur, []
        for i in range(n):
            for j in range(i + 1, n):
                order[i], order[j] = order[j], order[i]
                v = len(conflicts(order, R, C, pset))
                order[i], order[j] = order[j], order[i]
                if v < best:
                    best, choices = v, [(i, j)]
                elif v == best and v < cur:
                    choices.append((i, j))
        if not choices:
            return move - 1, False, True
        i, j = rng.choice(choices)
        order[i], order[j] = order[j], order[i]
    return MOVE_CAP, False, False

def measure(board, trials=50, seed=7):
    r = random.Random(seed)
    a = [policy_A(board, r) for _ in range(trials)]
    b = [policy_B(board, r) for _ in range(trials)]
    return dict(
        A_median=statistics.median(m for m, _ in a),
        A_pct20=100 * sum(1 for m, s in a if s and m <= 20) / trials,
        B_stall_pct=100 * sum(1 for *_ , st in b if st) / trials,
        B_solve_pct=100 * sum(1 for _, s, _ in b if s) / trials,
    )

def opt_ub(board):
    """Swap-distance from start to the known solution (upper bound)."""
    a, b = board["start"], board["solution"]
    pos = {w: i for i, w in enumerate(b)}
    target = [pos[w] for w in a]
    seen, cycles = [False] * len(a), 0
    for i in range(len(a)):
        if seen[i]:
            continue
        cycles += 1
        j = i
        while not seen[j]:
            seen[j] = True
            j = target[j]
    return len(a) - cycles

# ---------------------------------------------------------------- 2. grow
def build_sparse(seed, size, hub_t=2, extra_s=3):
    """The teaching board: a small hub + disjoint pairs, no dense weave."""
    partners = [p for p in sorted(adj.get(seed, set()), key=lambda w: fam[w])
                if clean(p) and not same_rootish(p, seed)]
    if len(partners) < hub_t:
        return None
    words = {seed}
    for p in partners:
        if len(words) >= 1 + hub_t:
            break
        if not any(same_rootish(p, w) for w in words):
            words.add(p)
    pool = list(pairs)
    rng.shuffle(pool)
    for pr in pool:
        if len(words) + 2 > size:
            break
        a, b = tuple(pr)
        if a in words or b in words:
            continue
        if any(same_rootish(a, w) or same_rootish(b, w) for w in words):
            continue
        if (adj[a] & words) or (adj[b] & words):
            continue
        words |= {a, b}
    return sorted(words) if len(words) == size else None

def grow_dense(seed, size, jitter=4):
    """Hard boards: grow the patch that maximizes internal opposition."""
    words = {seed}
    while len(words) < size:
        cands = set()
        for w in words:
            cands |= adj.get(w, set())
        cands -= words
        cands = [c for c in cands if clean(c)
                 and not any(same_rootish(c, w) for w in words)]
        if not cands:
            return None
        cands.sort(key=lambda w: (-len(adj[w] & words), fam[w]))
        words.add(rng.choice(cands[:jitter]))
    return sorted(words)

def gate_teach(s, np_):
    return 4 <= np_ <= 7 and 3 <= s["A_median"] <= 14 and s["B_stall_pct"] <= 15
def gate_ramp(s, np_):
    return (np_ >= 8 and 5 <= s["B_stall_pct"] <= 40 and s["B_solve_pct"] >= 55
            and s["A_median"] >= 30)
def gate_hard(min_pairs, stall_lo=30):
    def g(s, np_):
        return (np_ >= min_pairs and s["A_median"] >= MOVE_CAP * 0.8
                and s["B_stall_pct"] >= stall_lo and s["B_solve_pct"] >= 35)
    return g

def spec_for(k):
    if k == 1:  return 3, 3, gate_teach, "sparse"
    if k == 2:  return 3, 3, gate_ramp, "dense"
    if k <= 8:  return 3, 3, gate_hard(10, 25), "dense"
    if k <= 18: return 3, 4, gate_hard(14, 25), "dense"
    if k <= 34: return 4, 4, gate_hard(18), "dense"
    return 4, 5, gate_hard(22), "dense"

hubs = [w for w in sorted(adj, key=lambda w: (-len(adj[w]), fam[w]))
        if fam[w] <= 8000][:400]
rng.shuffle(hubs)
seed_iter = iter(hubs * 50)

# ---------------------------------------------------------------- 4. build + pack
t0 = time.time()
boards, pairsets = [], []
k, misses = 1, 0
while len(boards) < N_BOARDS:
    R, C, gate, mode = spec_for(k)
    seed = next(seed_iter)
    words = build_sparse(seed, R * C) if mode == "sparse" else grow_dense(seed, R * C)
    if not words:
        misses += 1; continue
    bp = board_pairs(words)
    ps_t = set(map(tuple, bp))
    if any(len(ps_t & prev) > 2 for prev in pairsets):   # the packing rule
        misses += 1; continue
    pset = {frozenset(p) for p in bp}
    sol = solve(words, R, C, pset)
    if sol is None:
        misses += 1; continue
    start = None
    for _ in range(400):
        cand = words[:]
        rng.shuffle(cand)
        if len(conflicts(cand, R, C, pset)) >= 3:
            start = cand; break
    if start is None:
        misses += 1; continue
    b = dict(rows=R, cols=C, words=words, pairs=bp, start=start, solution=sol, seed=seed)
    stats = measure(b)
    if not gate(stats, len(bp)):
        misses += 1; continue
    b["difficulty"] = stats
    boards.append(b); pairsets.append(ps_t); misses = 0
    print(f"board {k}: {R}x{C} seed={seed} pairs={len(bp)} "
          f"Amed={stats['A_median']} stall={stats['B_stall_pct']:.0f}% "
          f"[{time.time()-t0:.0f}s]")
    k += 1
    assert misses < 3000, "generator exhausted — relax a gate or widen seeds"

# final asserts: solvable, honest starts, no free parking, packing holds
for b in boards:
    ps = {frozenset(p) for p in b["pairs"]}
    assert not conflicts(b["solution"], b["rows"], b["cols"], ps)
    assert len(conflicts(b["start"], b["rows"], b["cols"], ps)) >= 3
    on = set(b["words"])
    assert all(len(adj.get(w, set()) & on) >= 1 for w in b["words"])
for i in range(len(boards)):
    for j in range(i + 1, len(boards)):
        assert len(pairsets[i] & pairsets[j]) <= 2

with open("hantai_boards.json", "w") as f:
    json.dump(boards, f, ensure_ascii=False)
print(f"wrote hantai_boards.json — {len(boards)} boards "
      f"[{(time.time()-t0)/60:.1f} min]")