PIG-Demo/envs/wordle_five/wordle_five/solver.py

"""The reference player: greedy maximum-entropy with full constraint filtering.

Two jobs, and it is important they are the same code:

  1. It is the `oracle` rung of the probe ladder.
  2. Its solve depth on a seed is the denominator of the `economy` reward.

Deriving the reward's reference from the shipped policy rather than from a
separately-computed optimum is deliberate. If `economy` were measured against a
depth-optimal search the oracle could not reach 1.0 — entropy-greedy is not
depth-optimal — and the probe's own assertion would fail on some seeds. The
optimum is still published, as a diagnostic, in `optimal_depth_report()`.

Note what this solver does NOT get: it scores 1.0 on solving and on economy,
and strictly below 1.0 on `consistency`, because maximising information means
deliberately guessing words that cannot win. That tension is the point — see
reward.py.
"""

from __future__ import annotations

import json
from functools import lru_cache
from pathlib import Path

import numpy as np

from .engine import GREEN, WORD_LENGTH, answers, score_guess

_CACHE = Path(__file__).parent.parent / ".pattern-cache" / "answers.npy"

# Base-3 code for a 5-tile pattern: 0..242, so it fits in a uint8. The all-green
# pattern is the largest code, which is the only value the solver compares
# against by name.
_TILE_VALUE = {"X": 0, "Y": 1, "G": 2}
ALL_GREEN_CODE = sum(2 * 3**i for i in range(WORD_LENGTH))


def pattern_code(pattern: str) -> int:
    return sum(_TILE_VALUE[t] * 3**i for i, t in enumerate(pattern))


def _compute_matrix(words: tuple[str, ...]) -> np.ndarray:
    """M[g, a] = the pattern code for guessing words[g] against answer words[a].

    Vectorised one letter at a time and chunked over guesses. The naive
    (n_guess, n_answer, 5) boolean intermediate is ~106 MB per letter at this
    list size, so chunking is what keeps peak memory in the tens of megabytes
    instead of gigabytes.
    """
    n = len(words)
    grid = np.frombuffer("".join(words).encode(), dtype=np.uint8).reshape(n, WORD_LENGTH)
    grid = grid - ord("a")

    out = np.zeros((n, n), dtype=np.uint8)
    place = np.array([3**i for i in range(WORD_LENGTH)], dtype=np.uint16)

    chunk = 256
    for start in range(0, n, chunk):
        stop = min(start + chunk, n)
        g = grid[start:stop]                       # (c, 5)
        # Greens first, exactly as engine.py does: every green in the word is
        # resolved before any yellow is considered.
        green = g[:, None, :] == grid[None, :, :]  # (c, n, 5)
        tiles = green.astype(np.uint16) * 2

        for letter in range(26):
            g_is = g == letter                      # (c, 5)
            a_is = grid == letter                   # (n, 5)
            if not g_is.any() or not a_is.any():
                continue
            # How many of this letter the answer has left after greens took theirs.
            green_here = green & g_is[:, None, :]    # (c, n, 5)
            available = a_is.sum(1)[None, :] - green_here.sum(2)   # (c, n)

            # Candidate yellow positions, ranked left to right. A position wins a
            # yellow only if its rank is inside the remaining allocation, which
            # is what makes SASSY/BASIS come out YGGXX rather than YGGYX.
            candidate = g_is[:, None, :] & ~green_here              # (c, n, 5)
            rank = np.cumsum(candidate, axis=2) - 1
            yellow = candidate & (rank < available[:, :, None])
            tiles += yellow.astype(np.uint16)

        out[start:stop] = (tiles * place[None, None, :]).sum(2).astype(np.uint8)

    return out


@lru_cache(maxsize=1)
def pattern_matrix() -> np.ndarray:
    """The answer-vs-answer pattern matrix, memoised on disk.

    Roughly 21 MB at the shipped list size. Gitignored and regenerated on
    demand — committing it would put a build artefact in a repo whose whole
    argument is that you can rebuild everything it claims.
    """
    words = answers()
    if _CACHE.exists():
        cached = np.load(_CACHE)
        if cached.shape == (len(words), len(words)):
            return cached
    matrix = _compute_matrix(words)
    _CACHE.parent.mkdir(parents=True, exist_ok=True)
    np.save(_CACHE, matrix)
    return matrix


def _verify_matrix_row(index: int) -> None:
    """Cross-check one matrix row against engine.score_guess. Used by tests."""
    words = answers()
    row = pattern_matrix()[index]
    for j, answer in enumerate(words):
        expected = pattern_code(score_guess(words[index], answer))
        if row[j] != expected:
            raise AssertionError(f"matrix disagrees with engine at ({index}, {j})")


def best_guess(candidates: np.ndarray, first: bool = False) -> int:
    """The index of the guess with the highest expected information gain.

    Ties break on being a candidate — among equally informative guesses, prefer
    one that could actually win, which is free expected value and also costs
    nothing in `consistency`.
    """
    matrix = pattern_matrix()
    if first:
        # The opening guess never depends on the seed, so it is computed once
        # per process rather than once per game.
        return _best_opener()
    if len(candidates) == 1:
        return int(candidates[0])

    sub = matrix[:, candidates]                        # (n_guess, n_candidates)
    best_index, best_score = -1, (-1.0, -1)
    total = len(candidates)
    candidate_set = set(candidates.tolist())
    for g in range(matrix.shape[0]):
        counts = np.bincount(sub[g], minlength=243)
        counts = counts[counts > 0]
        probability = counts / total
        entropy = float(-(probability * np.log2(probability)).sum())
        score = (entropy, 1 if g in candidate_set else 0)
        if score > best_score:
            best_index, best_score = g, score
    return best_index


@lru_cache(maxsize=1)
def _best_opener() -> int:
    matrix = pattern_matrix()
    total = matrix.shape[1]
    best_index, best_entropy = -1, -1.0
    for g in range(matrix.shape[0]):
        counts = np.bincount(matrix[g], minlength=243)
        counts = counts[counts > 0]
        probability = counts / total
        entropy = float(-(probability * np.log2(probability)).sum())
        if entropy > best_entropy:
            best_index, best_entropy = g, entropy
    return best_index


def solve(answer: str, max_guesses: int = 6) -> list[str]:
    """Play a full game against a known answer. Returns the guesses made."""
    words = answers()
    matrix = pattern_matrix()
    target = words.index(answer)
    candidates = np.arange(len(words))
    played: list[str] = []

    for turn in range(max_guesses):
        pick = best_guess(candidates, first=(turn == 0))
        played.append(words[pick])
        if pick == target:
            return played
        observed = matrix[pick, target]
        candidates = candidates[matrix[pick, candidates] == observed]
    return played


@lru_cache(maxsize=4096)
def reference_depth(answer: str) -> int:
    """How many guesses the shipped solver needs for this answer.

    This is the `economy` denominator. Cached because the probe and the reward
    both ask for it on the same seeds.
    """
    return len(solve(answer))


def entropy_of(guess: str, candidates: list[str]) -> float:
    """Expected bits gained by playing `guess` against a candidate set."""
    words = answers()
    matrix = pattern_matrix()
    g = words.index(guess)
    idx = np.array([words.index(c) for c in candidates])
    counts = np.bincount(matrix[g, idx], minlength=243)
    counts = counts[counts > 0]
    probability = counts / len(candidates)
    return float(-(probability * np.log2(probability)).sum())


def consistent_candidates(history: list[tuple[str, str]]) -> list[str]:
    """Every answer still viable given the feedback so far."""
    words = answers()
    matrix = pattern_matrix()
    alive = np.arange(len(words))
    for guess, pattern in history:
        if guess not in words:
            # A guess outside the answer list still constrains, but is not a row
            # of this matrix; fall back to the scalar engine for it.
            alive = np.array([i for i in alive if score_guess(guess, words[i]) == pattern])
            continue
        g = words.index(guess)
        alive = alive[matrix[g, alive] == pattern_code(pattern)]
    return [words[i] for i in alive]


def optimal_depth_report(seeds: list[int]) -> dict[str, float]:
    """Published diagnostics for the shipped solver over a set of seeds."""
    from .engine import answer_for_seed

    depths = [reference_depth(answer_for_seed(s)) for s in seeds]
    return {
        "mean_guesses": sum(depths) / len(depths),
        "worst_case": max(depths),
        "solved_within_six": sum(1 for d in depths if d <= 6) / len(depths),
    }