game/brain.py

"""The swappable model interface (SPEC.md §6).

A single ``Brain`` protocol with three implementations. This is what lets the
entire ledger/world/challenge loop be built and tested with ZERO GPU spend
(SPEC §0). Select via the ``BRAIN`` env var: ``stub`` (default) | ``local`` |
``modal``.

torch / transformers / spaces / modal are imported LAZILY inside the
implementation that needs them, so ``BRAIN=stub`` pulls in none of them and the
package imports cleanly for tests (SPEC §6/§7: "if BRAIN != modal, no Modal
dependency should be imported").
"""

from __future__ import annotations

import json
import os
import re
from typing import Optional, Protocol

from .prompt import PLAYER_MARKER


class Brain(Protocol):
    def respond(self, prompt: str) -> str:   # returns raw model text
        ...


# --------------------------------------------------------------------------- #
# 1) StubBrain — deterministic, no GPU (SPEC §6)
# --------------------------------------------------------------------------- #
# Concept proposals the heuristic offers when it detects a teaching moment.
_CANDIDATES = {
    "hidden_info": {
        "id": "hidden_info",
        "label": "hiding information",
        "understanding": "one mind can keep a thing another mind does not have, on purpose",
    },
    "gift": {
        "id": "gift",
        "label": "a gift",
        "understanding": "a thing passed to another to make their inside feel warm",
    },
    "surprise": {
        "id": "surprise",
        "label": "a surprise",
        "understanding": "a gift kept hidden until it is given — hiding and giving, together",
        "built_from": ["hidden_info", "gift"],
    },
    "secret": {
        "id": "secret",
        "label": "a secret",
        "understanding": "a thing one mind hides from another on purpose — hiding, but about knowing",
        "built_from": ["hidden_info"],
    },
}

_HIDE_WORDS = ("hide", "hidden", "conceal", "out of sight", "can't see", "cannot see", "not see")
_GIFT_WORDS = ("give", "gift", "present", "hand it", "hand the", "for the other", "for them")
_PUT_WORDS = ("put", "place", "drop", "into", "in the basket", "inside")
_POINT_WORDS = ("point", "show me", "which")
_MOVE_WORDS = ("move", "go to", "walk", "come")


class StubBrain:
    """No-GPU brain. Two modes:

    * **scripted** — constructed with a list of raw JSON strings, returned FIFO
      (then the safe-fallback JSON). Used by tests for exact control.
    * **heuristic** — default. Reads the player's line + object ids out of the
      prompt and returns a sensible valid-JSON action, so the whole arc is
      playable offline. Drives every §8 win from a natural utterance.
    """

    def __init__(self, scripted: Optional[list[str]] = None):
        self._scripted: Optional[list[str]] = list(scripted) if scripted is not None else None

    def respond(self, prompt: str) -> str:
        if self._scripted is not None:
            if self._scripted:
                return self._scripted.pop(0)
            return _fallback_json()
        return self._heuristic(prompt)

    # -- heuristic helpers --------------------------------------------------- #
    @staticmethod
    def _player_line(prompt: str) -> str:
        m = re.search(rf'{re.escape(PLAYER_MARKER)}\s*"(.*?)"', prompt, re.DOTALL)
        return (m.group(1) if m else "").lower()

    @staticmethod
    def _object_ids(prompt: str) -> list[str]:
        return re.findall(r"id:\s*([a-z_]+)\)", prompt)

    def _pick_stone(self, said: str, obj_ids: list[str]) -> Optional[str]:
        if "blue" in said and "blue_stone" in obj_ids:
            return "blue_stone"
        if "red" in said and "red_stone" in obj_ids:
            return "red_stone"
        stones = [o for o in obj_ids if o.endswith("_stone")]
        return stones[0] if stones else (obj_ids[0] if obj_ids else None)

    def _heuristic(self, prompt: str) -> str:
        said = self._player_line(prompt)
        obj_ids = self._object_ids(prompt)
        has_basket = "basket" in obj_ids or "(id: basket)" in prompt
        stone = self._pick_stone(said, obj_ids)

        def has_any(words):
            return any(w in said for w in words)

        # Order matters: conceal-words win over give/put because "surprise" and
        # "secret" prompts must route to hiding (the generalization beat).
        if has_any(_HIDE_WORDS) or "surprise" in said or "secret" in said:
            if stone and has_basket:
                if "surprise" in said:
                    cand, line, gap = _CANDIDATES["surprise"], (
                        "I take a thing and hide it for the other one. They do not "
                        "know yet. This is the hiding you taught me, but for giving. "
                        "A 'surprise'?"
                    ), None
                elif "secret" in said:
                    cand, line, gap = _CANDIDATES["secret"], (
                        "I keep the thing where the other one cannot find it in their "
                        "mind. I hold what they do not. A 'secret' is a hiding of knowing?"
                    ), None
                else:
                    cand, line, gap = _CANDIDATES["hidden_info"], (
                        "I place the thing inside the holder. The other one cannot see "
                        "it now. One mind has it; the other does not?"
                    ), "I do not understand why you want the other to not-see."
                return _json(
                    {"verb": "put_in", "args": {"obj_id": stone, "container_id": "basket"}},
                    line, gap, cand,
                )

        if has_any(_GIFT_WORDS):
            if stone:
                return _json(
                    {"verb": "give", "args": {"obj_id": stone, "agent_id": "other"}},
                    "I pass the thing to the other one. For them to hold. This makes a "
                    "warmth inside them? A 'present'.",
                    None,
                    _CANDIDATES["gift"],
                )

        if has_any(_PUT_WORDS):
            if stone and has_basket:
                return _json(
                    {"verb": "put_in", "args": {"obj_id": stone, "container_id": "basket"}},
                    "I put the thing into the holder. It rests there now.",
                    None,
                    None,
                )

        if has_any(_POINT_WORDS) and stone:
            return _json(
                {"verb": "point_at", "args": {"target": stone}},
                "I direct your attention. This thing. Here.",
                None,
                None,
            )

        if has_any(_MOVE_WORDS):
            return _json(
                {"verb": "move_to", "args": {"target": "other"}},
                "I change where I am. I come closer.",
                None,
                None,
            )

        # Nothing matched — the alien, honestly, does not understand.
        return _json(
            {"verb": "wait", "args": {}},
            "The alien holds still, watching your mouth make the shapes.",
            "I have no concept for what you ask.",
            None,
        )


def _json(action: dict, utterance: str, gap, candidate) -> str:
    return json.dumps(
        {
            "action": action,
            "utterance": utterance,
            "gap": gap,
            "candidate_concept": candidate,
        }
    )


def _fallback_json() -> str:
    return _json(
        {"verb": "wait", "args": {}},
        "The alien looks at you, not understanding.",
        "I could not grasp that.",
        None,
    )


# --------------------------------------------------------------------------- #
# 2) LocalBrain — transformers on 'cuda', the real ZeroGPU path (SPEC §6)
# --------------------------------------------------------------------------- #
class LocalBrain:
    """Loads a ≤32B instruct model onto 'cuda' at construction time (which
    app.py runs at MODULE level, i.e. Space startup — NOT lazily inside the
    @spaces.GPU function, per SPEC §0/§13).

    The actual generation lives in ``_generate``; app.py is responsible for
    wrapping the call site with ``@spaces.GPU(duration=...)`` so only inference
    holds the GPU — state mutation happens outside it (SPEC §6).
    """

    # Chosen by the 2026-06 bake-off (bakeoff.py --temps sweep, 55 draws/temp/model;
    # numbers in bakeoff_results.json, untracked). JSON-1st was 100% at EVERY temp
    # for every candidate, so the pick fell to arc-completion + concept invention:
    # 14B was the only model strong at both (3/5 arc, 19-24 unique proposals warm).
    # 7B executed nothing (0/5 arc, ~0 proposals); phi-4 won 4/5 but proposes
    # almost nothing (the ledger never grows); Mistral-24B has the richest voice
    # but 0/5 arc and ACTION slips when hot.
    DEFAULT_MODEL = os.environ.get("MODEL_ID", "Qwen/Qwen2.5-14B-Instruct")
    # 0.9: the same sweep showed the JSON envelope holding 100% even at T=1.0, so
    # we run near-peak voice/invention (UTTER-DIV ~0.8, UNIQ-CC ~20+) at zero
    # measured reliability cost. Mutable so a sweep retunes without reloading.
    DEFAULT_TEMPERATURE = float(os.environ.get("LOCALBRAIN_TEMPERATURE", "0.9"))

    def __init__(self, model_id: Optional[str] = None, temperature: Optional[float] = None):
        import torch  # noqa: F401  (lazy: only when BRAIN=local)
        from transformers import AutoModelForCausalLM, AutoTokenizer

        self.model_id = model_id or self.DEFAULT_MODEL
        self.temperature = self.DEFAULT_TEMPERATURE if temperature is None else float(temperature)
        self.max_new_tokens = 300
        self.tokenizer = AutoTokenizer.from_pretrained(self.model_id)
        self.model = AutoModelForCausalLM.from_pretrained(
            self.model_id,
            torch_dtype="auto",
            device_map="cuda",           # module-level CUDA placement (SPEC §0)
        )

    def respond(self, prompt: str) -> str:
        return self._generate(prompt)

    @staticmethod
    def _sampler_kwargs(temperature: float) -> dict:
        """temperature -> generate() sampler kwargs. <=0 => greedy (no temperature
        passed; transformers requires do_sample=False *without* a temperature). >0
        => sample at that heat. Pure + side-effect-free so the temperature wiring
        is unit-testable with no model load (tests/test_brain.py)."""
        if temperature and temperature > 0:
            return {"do_sample": True, "temperature": float(temperature), "top_p": 0.9}
        return {"do_sample": False}

    def _generate(self, prompt: str) -> str:
        import torch

        messages = [{"role": "user", "content": prompt}]
        text = self.tokenizer.apply_chat_template(
            messages, tokenize=False, add_generation_prompt=True
        )
        inputs = self.tokenizer(text, return_tensors="pt").to(self.model.device)
        # >>> Temperature threads through HERE — reconcile if your decode differs. <<<
        gen_kwargs = dict(
            max_new_tokens=self.max_new_tokens,
            pad_token_id=self.tokenizer.eos_token_id,
            **self._sampler_kwargs(self.temperature),
        )
        with torch.no_grad():
            out = self.model.generate(**inputs, **gen_kwargs)
        gen = out[0][inputs["input_ids"].shape[-1]:]
        return self.tokenizer.decode(gen, skip_special_tokens=True)


# --------------------------------------------------------------------------- #
# 3) ModalBrain — optional dev/serving endpoint, NEVER the submission (SPEC §7)
# --------------------------------------------------------------------------- #
class ModalBrain:
    """Calls a Modal HTTP endpoint. Optional. urllib only (stdlib) so it is never
    a hard dependency of the Space (SPEC §7)."""

    # Kept in lockstep with LocalBrain.DEFAULT_TEMPERATURE (guarded by a test).
    DEFAULT_TEMPERATURE = float(os.environ.get("LOCALBRAIN_TEMPERATURE", "0.9"))

    def __init__(self, endpoint: Optional[str] = None,
                 model_id: Optional[str] = None, temperature: Optional[float] = None):
        self.endpoint = endpoint or os.environ.get("MODAL_ENDPOINT", "")
        if not self.endpoint:
            raise RuntimeError("ModalBrain requires MODAL_ENDPOINT to be set.")
        self.model_id = model_id or os.environ.get("MODEL_ID") or None
        self.temperature = self.DEFAULT_TEMPERATURE if temperature is None else float(temperature)

    def respond(self, prompt: str) -> str:
        import json as _json
        import urllib.request

        payload = {"prompt": prompt, "temperature": self.temperature}
        if self.model_id:
            payload["model_id"] = self.model_id
        req = urllib.request.Request(
            self.endpoint, data=_json.dumps(payload).encode(),
            headers={"Content-Type": "application/json"}, method="POST")
        with urllib.request.urlopen(req, timeout=180) as r:
            data = _json.loads(r.read().decode())
        return data.get("text", "") if isinstance(data, dict) else str(data)


# --------------------------------------------------------------------------- #
# Factory
# --------------------------------------------------------------------------- #
def _default_brain() -> str:
    """Which brain to use when BRAIN is not set explicitly. Hugging Face sets
    SPACE_ID on every Space, so default to the real model there and the zero-GPU
    stub locally — SPEC §6's "stub locally, local on the Space". Without this, an
    unset BRAIN on the Space silently served the StubBrain (canned, repetitive)."""
    explicit = os.environ.get("BRAIN")
    if explicit:
        return explicit
    return "local" if os.environ.get("SPACE_ID") else "stub"


def make_brain(name: Optional[str] = None) -> Brain:
    name = (name or _default_brain()).lower()
    if name == "stub":
        return StubBrain()
    if name == "local":
        return LocalBrain()
    if name == "modal":
        return ModalBrain()
    raise ValueError(f"unknown BRAIN '{name}' (expected stub|local|modal)")