# /// script
# requires-python = "==3.12.*"
# dependencies = ["matplotlib", "wordfreq==3.1.1"]
# ///
"""Reproduce the board-opportunity analysis for classic 4x4 Boggle.

The dictionary is pinned to the exact ENABLE file used by the original 2016
notebook. A local games checkout is preferred; otherwise the script downloads
that file from the original commit and verifies its SHA-256 digest.
"""

from __future__ import annotations

import hashlib
import math
import os
import random
import sys
import urllib.request
from collections import Counter
from pathlib import Path

import matplotlib.pyplot as plt
from wordfreq import zipf_frequency

sys.path.insert(0, str(Path(__file__).resolve().parents[1]))
from figure_quality import configure_matplotlib, figure_output_path, validate_figure

configure_matplotlib()

SEED = 20_161_126
BOARD_COUNT = 15_000
DICTIONARY_COMMIT = "8407a721b8b904b861b1dc71233cc8370dbe6673"
DICTIONARY_SHA256 = "a392640f14602fcbb36779eb6671d170260c06d43f41c16383e7f0e35b221b4a"
DICTIONARY_URL = (
    "https://raw.githubusercontent.com/arclabs561/games/"
    f"{DICTIONARY_COMMIT}/boggle/enable1.txt"
)

CUBES = (
    ("a", "a", "e", "e", "g", "n"),
    ("a", "b", "b", "j", "o", "o"),
    ("a", "c", "h", "o", "p", "s"),
    ("a", "f", "f", "k", "p", "s"),
    ("a", "o", "o", "t", "t", "w"),
    ("c", "i", "m", "o", "t", "u"),
    ("d", "e", "i", "l", "r", "x"),
    ("d", "e", "l", "r", "v", "y"),
    ("d", "i", "s", "t", "t", "y"),
    ("e", "e", "g", "h", "n", "w"),
    ("e", "e", "i", "n", "s", "u"),
    ("e", "h", "r", "t", "v", "w"),
    ("e", "i", "o", "s", "s", "t"),
    ("e", "l", "r", "t", "t", "y"),
    ("h", "i", "m", "n", "qu", "u"),
    ("h", "l", "n", "n", "r", "z"),
)
END = "\0"


def hoeffding_sample_size(epsilon: float, delta: float) -> int:
    """Return the two-sided Hoeffding sample size for one fixed Bernoulli mean."""
    if not 0 < epsilon < 1:
        raise ValueError("epsilon must be between zero and one")
    if not 0 < delta < 1:
        raise ValueError("delta must be between zero and one")
    return math.ceil(math.log(2 / delta) / (2 * epsilon**2))


def dictionary_bytes() -> bytes:
    override = os.environ.get("BOGGLE_DICTIONARY")
    local_default = (
        Path(__file__).resolve().parents[5] / "games" / "boggle" / "enable1.txt"
    )
    local_path = Path(override).expanduser() if override else local_default
    if local_path.is_file():
        data = local_path.read_bytes()
        source = str(local_path)
    else:
        with urllib.request.urlopen(DICTIONARY_URL, timeout=30) as response:
            data = response.read()
        source = DICTIONARY_URL

    digest = hashlib.sha256(data).hexdigest()
    if digest != DICTIONARY_SHA256:
        raise RuntimeError(
            f"dictionary digest mismatch: got {digest}, expected {DICTIONARY_SHA256}"
        )
    print(f"dictionary={source}\nsha256={digest}")
    return data


def make_trie(words: list[str]) -> dict:
    root: dict = {}
    for word in words:
        node = root
        for character in word:
            node = node.setdefault(character, {})
        node[END] = None
    return root


def advance(node: dict, tile: str) -> dict | None:
    """Advance through every character on a tile, including the ``qu`` tile."""
    for character in tile:
        child = node.get(character)
        if child is None:
            return None
        node = child
    return node


NEIGHBORS = tuple(
    tuple(
        other_row * 4 + other_column
        for other_row in range(max(0, row - 1), min(4, row + 2))
        for other_column in range(max(0, column - 1), min(4, column + 2))
        if (other_row, other_column) != (row, column)
    )
    for row in range(4)
    for column in range(4)
)


def words_on_board(board: list[str], trie: dict) -> set[str]:
    found: set[str] = set()

    def visit(index: int, node: dict, visited: int, prefix: str) -> None:
        next_node = advance(node, board[index])
        if next_node is None:
            return
        word = prefix + board[index]
        next_visited = visited | (1 << index)
        if END in next_node and len(word) >= 3:
            found.add(word)
        for neighbor in NEIGHBORS[index]:
            if not next_visited & (1 << neighbor):
                visit(neighbor, next_node, next_visited, word)

    for start in range(16):
        visit(start, trie, 0, "")
    return found


def boggle_points(word: str) -> int:
    length = len(word)
    if length < 3:
        return 0
    if length <= 4:
        return 1
    if length == 5:
        return 2
    if length == 6:
        return 3
    if length == 7:
        return 5
    return 11


def wilson_interval(
    successes: int, trials: int, z: float = 1.96
) -> tuple[float, float]:
    proportion = successes / trials
    denominator = 1 + z * z / trials
    center = (proportion + z * z / (2 * trials)) / denominator
    margin = (
        z
        * math.sqrt(
            proportion * (1 - proportion) / trials + z * z / (4 * trials * trials)
        )
        / denominator
    )
    return center - margin, center + margin


words = [
    word
    for raw_word in dictionary_bytes().decode("utf-8").splitlines()
    if (word := raw_word.strip().lower()).isalpha() and len(word) >= 3
]
trie = make_trie(words)
rng = random.Random(SEED)
appearances: Counter[str] = Counter()

for _ in range(BOARD_COUNT):
    cube_order = rng.sample(CUBES, k=len(CUBES))
    board = [rng.choice(cube) for cube in cube_order]
    appearances.update(sorted(words_on_board(board, trie)))

rows = []
for word, count in appearances.items():
    probability = count / BOARD_COUNT
    rows.append(
        {
            "word": word,
            "count": count,
            "probability": probability,
            "points": boggle_points(word),
            "opportunity": probability * boggle_points(word),
            "zipf": zipf_frequency(word, "en"),
        }
    )

most_common = sorted(rows, key=lambda row: (-row["probability"], row["word"]))[:15]
most_common_four_plus = sorted(
    (row for row in rows if len(row["word"]) >= 4),
    key=lambda row: (-row["probability"], row["word"]),
)[:15]
print(f"seed={SEED} boards={BOARD_COUNT} dictionary_words={len(words)}")
sample_size = hoeffding_sample_size(epsilon=0.01, delta=0.1)
assert sample_size == 14_979
print(f"Hoeffding samples for epsilon=0.01, delta=0.1: {sample_size:,}")
print("\nMost common board words")
print("| word | boards | probability | 95% Wilson interval |")
print("|:---|---:|---:|:---|")
for row in most_common:
    low, high = wilson_interval(row["count"], BOARD_COUNT)
    print(
        f"| {row['word']} | {row['count']} | {row['probability']:.4%} "
        f"| [{low:.4%}, {high:.4%}] |"
    )

print("\nMost common words with at least four letters")
print("| word | boards | probability |")
print("|:---|---:|---:|")
for row in most_common_four_plus:
    print(f"| {row['word']} | {row['count']} | {row['probability']:.4%} |")

by_word = {row["word"]: row for row in rows}
print("\nSelected low-corpus-frequency four-letter words")
print("| word | board probability | solo opportunity | wordfreq Zipf |")
print("|:---|---:|---:|---:|")
for word in ("toea", "seta", "nett", "teat", "stet", "tret", "rete", "sett"):
    row = by_word[word]
    zipf = f"{row['zipf']:.2f}" if row["zipf"] > 0 else "unlisted"
    print(f"| {word} | {row['probability']:.4%} | {row['opportunity']:.4f} | {zipf} |")

eligible = [
    row
    for row in rows
    if row["count"] >= 15 and row["zipf"] > 0 and row["opportunity"] >= 0.003
]
frontier = []
best_frequency = math.inf
for row in sorted(eligible, key=lambda item: (-item["opportunity"], item["zipf"])):
    if row["zipf"] < best_frequency:
        frontier.append(row)
        best_frequency = row["zipf"]

print("\nNon-dominated measured opportunity / corpus-frequency frontier")
print("| word | expected solo points/board | wordfreq Zipf frequency |")
print("|:---|---:|---:|")
for row in frontier[:20]:
    print(f"| {row['word']} | {row['opportunity']:.4f} | {row['zipf']:.2f} |")

fig, (common_ax, scatter_ax) = plt.subplots(1, 2, figsize=(10.4, 4.5))
fig.patch.set_facecolor("#ffffff")
for axis in (common_ax, scatter_ax):
    axis.set_facecolor("#ffffff")
    axis.spines["top"].set_visible(False)
    axis.spines["right"].set_visible(False)
    axis.tick_params(labelsize=10, colors="#30343B")
    axis.spines["bottom"].set_color("#3F4650")
    axis.spines["left"].set_color("#3F4650")

top = most_common[:10][::-1]
common_ax.barh(
    [row["word"] for row in top],
    [100 * row["probability"] for row in top],
    color="#1E5A8A",
)
common_ax.set_xlabel("Boards containing word (%)", fontsize=12)
common_ax.set_title("Board opportunity", fontsize=14)

plotted = [row for row in rows if row["count"] >= 15 and row["opportunity"] > 0]
listed = [row for row in plotted if row["zipf"] > 0]
unlisted = [row for row in plotted if row["zipf"] == 0]
scatter_ax.scatter(
    [row["opportunity"] for row in listed],
    [row["zipf"] for row in listed],
    s=12,
    alpha=0.28,
    color="#1E5A8A",
    linewidths=0,
    label="listed by wordfreq",
)
scatter_ax.scatter(
    [row["opportunity"] for row in unlisted],
    [0.65 for _ in unlisted],
    s=16,
    alpha=0.35,
    facecolors="none",
    edgecolors="#8B3A2B",
    label="unlisted, not zero usage",
)
scatter_ax.scatter(
    [row["opportunity"] for row in frontier],
    [row["zipf"] for row in frontier],
    s=34,
    color="#8B3A2B",
    linewidths=0,
    label="measured non-dominated frontier",
)

scatter_ax.set_xscale("log")
scatter_ax.set_xlabel("Expected solo points per board", fontsize=12)
scatter_ax.set_ylabel("English Zipf frequency", fontsize=12)
scatter_ax.set_title("Corpus frequency is not findability", fontsize=14)
scatter_ax.legend(
    fontsize=10,
    frameon=True,
    fancybox=False,
    edgecolor="#8A8F98",
    loc="upper left",
)

fig.tight_layout()
validate_figure(fig)
output = figure_output_path(__file__, "boggle_opportunity.png")
fig.savefig(output, dpi=300, bbox_inches="tight", facecolor="white")
print(f"\nsaved={output}")