Connections

The prompt

Build a connections puzzle. Sixteen words, and I have to find four groups of four. The groups have a difficulty, from obvious to tricky, and when I find one it should lock in as a coloured bar, easiest colour first in the usual way. I click four words and submit. I can unselect a word before I submit.

A correct group stays solved. A wrong guess costs one of four mistakes, and if I was one word away, tell me that. Four mistakes ends the game and shows the groups I missed. Finding all four is a perfect round. Harder groups should be worth more than easier ones. Mistakes should cost a little. Show the tiles, the mistakes left, and the score. Several puzzles. The same puzzle should present the words in the same order.

# AI : Connections — find four groups of four words. Same task always lays the same board.
from __future__ import annotations

from typing import Any, Dict, List, Optional, Tuple

from helpers import base_observation, clamp01, pick_task, rng_for, task_ids

GAME_CLASS = "Connections"

DIFFICULTY = ["yellow", "green", "blue", "purple"]
POINTS = {"yellow": 10, "green": 20, "blue": 30, "purple": 40}
MISTAKE_COST = 5
TOTAL_MISTAKES = 4
PERFECT_BONUS = 10
MAX_SCORE = sum(POINTS.values()) + PERFECT_BONUS


def _g(key: str, label: str, words: List[str]) -> Dict[str, Any]:
    return {"key": key, "label": label, "words": list(words), "points": POINTS[key]}


PUZZLES = [
    {
        "id": "s1",
        "groups": [
            _g("yellow", "Ocean Animals", ["WHALE", "SHARK", "DOLPHIN", "OCTOPUS"]),
            _g("green", "Trees", ["OAK", "PINE", "MAPLE", "BIRCH"]),
            _g("blue", "Roman Numerals", ["FIVE", "TEN", "FIFTY", "HUNDRED"]),
            _g("purple", "___ Cheese", ["BLUE", "COTTAGE", "CREAM", "GOAT"]),
        ],
    },
    {
        "id": "s2",
        "groups": [
            _g("yellow", "Footwear", ["SHOE", "BOOT", "SANDAL", "SNEAKER"]),
            _g("green", "Instruments", ["PIANO", "GUITAR", "VIOLIN", "FLUTE"]),
            _g("blue", "Planets", ["MARS", "VENUS", "JUPITER", "SATURN"]),
            _g("purple", "Silent First Letter", ["KNIFE", "LAMB", "HOUR", "WRITE"]),
        ],
    },
    {
        "id": "s3",
        "groups": [
            _g("yellow", "Breakfast Foods", ["PANCAKE", "WAFFLE", "BACON", "OMELET"]),
            _g("green", "Dog Breeds", ["LABRADOR", "POODLE", "BEAGLE", "HUSKY"]),
            _g("blue", "Noble Gases", ["HELIUM", "NEON", "ARGON", "KRYPTON"]),
            _g("purple", "NATO Alphabet", ["ALPHA", "BRAVO", "CHARLIE", "DELTA"]),
        ],
    },
    {
        "id": "s4",
        "groups": [
            _g("yellow", "Body Parts", ["ARM", "LEG", "HEAD", "HAND"]),
            _g("green", "Spanish Numbers", ["UNO", "DOS", "TRES", "CUATRO"]),
            _g("blue", "U.S. States", ["TEXAS", "OHIO", "MAINE", "IOWA"]),
            _g("purple", "Famous Parks", ["CENTRAL", "YELLOWSTONE", "JURASSIC", "HYDE"]),
        ],
    },
    {
        "id": "s5",
        "groups": [
            _g("yellow", "Weather", ["RAIN", "SNOW", "WIND", "FOG"]),
            _g("green", "Cooking Methods", ["BAKE", "FRY", "GRILL", "BOIL"]),
            _g("blue", "Greek Letters", ["ALPHA", "BETA", "GAMMA", "DELTA"]),
            _g("purple", "Card Suits", ["HEARTS", "SPADES", "CLUBS", "DIAMONDS"]),
        ],
    },
    {
        "id": "s6",
        "groups": [
            _g("yellow", "Furniture", ["CHAIR", "TABLE", "SOFA", "BED"]),
            _g("green", "Sports", ["SOCCER", "TENNIS", "GOLF", "HOCKEY"]),
            _g("blue", "Programming Languages", ["PYTHON", "JAVA", "RUBY", "SWIFT"]),
            _g("purple", "Constellations", ["ORION", "LEO", "LYRA", "CYGNUS"]),
        ],
    },
]

PUZZLE_IDS = task_ids(len(PUZZLES))
_BY_ID = {p["id"]: p for p in PUZZLES}


class Connections:
    def __init__(self) -> None:
        self.task_id: str = ""
        self.puzzle: Dict[str, Any] = {}
        self.words: List[Dict[str, Any]] = []
        self.word_to_key: Dict[str, str] = {}
        self.selected: List[str] = []
        self.solved: set = set()
        self.mistakes: int = 0
        self.score: int = 0
        self.done: bool = False
        self.perfect: bool = False

    def reset(self, task_id: str = "") -> Dict[str, Any]:
        self.task_id = pick_task(task_id, PUZZLE_IDS, default=PUZZLE_IDS[0])
        self.puzzle = _BY_ID[self.task_id]
        rng = rng_for(self.task_id)
        flat = [(w, g["key"]) for g in self.puzzle["groups"] for w in g["words"]]
        rng.shuffle(flat)
        self.words = []
        self.word_to_key = {}
        for i, (word, key) in enumerate(flat):
            self.words.append({"id": str(i), "word": word, "solved": False})
            self.word_to_key[str(i)] = key
        self.selected = []
        self.solved = set()
        self.mistakes = 0
        self.score = 0
        self.done = False
        self.perfect = False
        index = PUZZLE_IDS.index(self.task_id) + 1
        return self._obs("", f"Find four groups of four words. Puzzle {index}.", "")

    def _group(self, key: str) -> Dict[str, Any]:
        for g in self.puzzle["groups"]:
            if g["key"] == key:
                return {"key": g["key"], "label": g["label"], "words": list(g["words"]), "points": g["points"]}
        return {}

    def _legal(self) -> List[str]:
        moves = [f"pick:{w['id']}" for w in self.words if not w["solved"]]
        if len(self.selected) == 4:
            moves.append("submit")
        if self.selected:
            moves.append("clear")
        return moves

    def _obs(self, move: str, message: str = "", message_kind: str = "") -> Dict[str, Any]:
        words = [{"id": w["id"], "word": w["word"], "solved": w["solved"]} for w in self.words]
        solved_groups = [self._group(k) for k in DIFFICULTY if k in self.solved]
        obs = base_observation(
            reward=clamp01(self.score / MAX_SCORE),
            done=self.done,
            legal_moves=self._legal(),
            last_move=move or None,
            puzzle_id=self.task_id,
            puzzle_label=f"Puzzle {PUZZLE_IDS.index(self.task_id) + 1}",
            words=words,
            selected=list(self.selected),
            solved_groups=solved_groups,
            mistakes_left=max(0, TOTAL_MISTAKES - self.mistakes),
            mistakes_total=TOTAL_MISTAKES,
            score=self.score,
            message=message,
            message_kind=message_kind,
            perfect=self.perfect,
        )
        if self.done:
            obs["missed_groups"] = [self._group(k) for k in DIFFICULTY if k not in self.solved]
        return obs

    def step(self, action: Any) -> Tuple[Dict[str, Any], bool, Optional[str]]:
        move = (action.move or "").strip()
        if self.done:
            return self._obs(move), False, "game over"
        if move.startswith("pick:"):
            wid = move[5:]
            if wid not in self.word_to_key:
                return self._obs(move), False, "unknown word"
            w = self.words[int(wid)]
            if w["solved"]:
                return self._obs(move), False, "already solved"
            if wid in self.selected:
                self.selected.remove(wid)
            elif len(self.selected) < 4:
                self.selected.append(wid)
            return self._obs(move), True, None
        if move == "clear":
            self.selected = []
            return self._obs(move), True, None
        if move == "submit":
            if len(self.selected) != 4:
                return self._obs(move), False, "select four words"
            return self._submit(move)
        return self._obs(move), False, "unknown move"

    def _submit(self, move: str) -> Tuple[Dict[str, Any], bool, Optional[str]]:
        chosen = {self.words[int(i)]["word"] for i in self.selected}
        matched: Optional[str] = None
        for key in DIFFICULTY:
            if key in self.solved:
                continue
            if set(self._group(key)["words"]) == chosen:
                matched = key
                break
        if matched is not None:
            self.solved.add(matched)
            for i in list(self.selected):
                self.words[int(i)]["solved"] = True
            self.score += self._group(matched)["points"]
            self.selected = []
            if len(self.solved) == 4:
                self.done = True
                self.perfect = True
                self.score += PERFECT_BONUS
                return self._obs(move, "Perfect round! All four groups found.", "good"), True, None
            return self._obs(move, "Correct! " + self._group(matched)["label"], "good"), True, None
        one_away = False
        for key in DIFFICULTY:
            if key in self.solved:
                continue
            if len(chosen & set(self._group(key)["words"])) == 3:
                one_away = True
                break
        self.mistakes += 1
        self.score -= MISTAKE_COST
        self.selected = []
        if self.mistakes >= TOTAL_MISTAKES:
            self.done = True
            self.perfect = False
            return self._obs(move, "Out of mistakes. Missed groups are shown above.", "bad"), True, None
        if one_away:
            return self._obs(move, "One away!", "oneaway"), True, None
        return self._obs(move, "Wrong guess", "bad"), True, None