monte_carlo_search.py

import os
import sys
import itertools
import math
from typing import List, Dict, Optional, Tuple

import numpy as np
import torch

# Allow `python src/monte_carlo_search.py` from the repo root without install.
_REPO_ROOT = os.path.abspath(os.path.join(os.path.dirname(__file__), ".."))
if _REPO_ROOT not in sys.path:
    sys.path.insert(0, _REPO_ROOT)

from src.model.transformer import AlphaMonGPT
from src.data.vocab_manager import VocabManager

# OBJECTIVES
TURN = "turn"  # it means we need to know the best move for the turn
TEAM_PREVIEW = "teampreview"  # it means we need to know the best team selection of 4
FAINT = "faint"  # it means we need to know the best switch in after mon fainted

# Tokens that mark the end of a per-action block during rollout parsing.
_ACTION_BLOCK_TERMINATORS = (
    "[CMD_MOVE]", "[TURN]", "[SWITCH]", "[FAINT]", "[TERASTALLIZE]",
)


def _win_score(value: float, as_player: int) -> float:
    """Convert the value head's P1-win-prob into the requested player's win prob."""
    return (1.0 - value) if as_player == 2 else value


# ---------------------------------------------------------------------------
#  Structured action extraction -- robust to duplicated / extra tokens
# ---------------------------------------------------------------------------

def extract_slot_actions(
    decoded_tokens: List[str], player_token: str
) -> Dict[str, dict]:
    """Parse decoded tokens into per-slot actions for *player_token*.

    Returns a dict keyed by slot ("[SIDE_A]" / "[SIDE_B]") with values like:
        {"type": "move", "move": "MOVE:FAKE_OUT",
         "target_player": "[P2]", "target_side": "[SIDE_A]", "target_mon": "MON:X"}
        {"type": "switch", "mon": "MON:INCINEROAR"}

    Tolerant of duplicated player/side tokens the model sometimes generates.
    """
    actions: Dict[str, dict] = {}
    tera_info: Dict[str, str] = {}
    tokens = decoded_tokens
    n = len(tokens)
    i = 0

    while i < n:
        tok = tokens[i]

        # ---- [CMD_MOVE] block ----
        if tok == "[CMD_MOVE]":
            j = i + 1
            actor_player = None
            actor_side = None
            actor_mon = None
            move_name = None
            target_player = None
            target_side = None
            target_mon = None
            phase = "actor"

            while j < n and tokens[j] not in _ACTION_BLOCK_TERMINATORS:
                t = tokens[j]
                if t in ("[P1]", "[P2]"):
                    if phase == "actor":
                        actor_player = t  # overwrite duplicates
                    else:
                        target_player = t
                elif t in ("[SIDE_A]", "[SIDE_B]"):
                    if phase == "actor" and actor_side is None:
                        actor_side = t
                    elif phase == "target":
                        target_side = t
                elif t.startswith("MON:"):
                    if phase == "actor" and actor_mon is None:
                        actor_mon = t
                        phase = "target"
                    elif phase == "target" and target_mon is None:
                        target_mon = t
                elif t.startswith("MOVE:"):
                    if move_name is None:
                        move_name = t
                j += 1

            if actor_player == player_token and actor_side and move_name:
                actions[actor_side] = {
                    "type": "move",
                    "move": move_name,
                    "mon": actor_mon,
                    "target_player": target_player,
                    "target_side": target_side,
                    "target_mon": target_mon,
                }
            i = j
            continue

        # ---- [SWITCH] block ----
        if tok == "[SWITCH]":
            j = i + 1
            sw_player = None
            sw_side = None
            sw_mon = None
            while j < n and tokens[j] not in _ACTION_BLOCK_TERMINATORS:
                t = tokens[j]
                if t in ("[P1]", "[P2]"):
                    sw_player = t
                elif t in ("[SIDE_A]", "[SIDE_B]"):
                    if sw_side is None:
                        sw_side = t
                elif t.startswith("MON:"):
                    if sw_mon is None:
                        sw_mon = t
                j += 1

            if sw_player == player_token and sw_side and sw_mon:
                actions[sw_side] = {"type": "switch", "mon": sw_mon}
            i = j
            continue

        # ---- [TERASTALLIZE] block ----
        if tok == "[TERASTALLIZE]":
            j = i + 1
            tera_player = None
            tera_side = None
            tera_type = None
            while j < n and tokens[j] not in _ACTION_BLOCK_TERMINATORS:
                t = tokens[j]
                if t in ("[P1]", "[P2]"):
                    tera_player = t
                elif t in ("[SIDE_A]", "[SIDE_B]"):
                    if tera_side is None:
                        tera_side = t
                elif t.startswith("TYPE:"):
                    tera_type = t
                j += 1

            if tera_player == player_token and tera_side and tera_type:
                tera_info[tera_side] = tera_type
            i = j
            continue

        i += 1

    # Merge tera info into actions
    for slot, ttype in tera_info.items():
        if slot in actions:
            actions[slot]["tera"] = ttype

    return actions


def _action_signature_key(actions: Dict[str, dict]) -> str:
    """Build a canonical, order-independent signature from per-slot actions.

    Two rollouts that choose the same moves/targets on the same slots produce
    the same key even if the raw token streams differ.
    """
    parts = []
    for slot in sorted(actions.keys()):
        a = actions[slot]
        if a["type"] == "move":
            tera = a.get("tera", "")
            parts.append(
                f"{slot}|move|{a['move']}|{a.get('target_player', '')}|"
                f"{a.get('target_side', '')}|{tera}"
            )
        elif a["type"] == "switch":
            parts.append(f"{slot}|switch|{a['mon']}")
    return ";;".join(parts) if parts else ""


# ---------------------------------------------------------------------------
#  Human-readable formatting
# ---------------------------------------------------------------------------

def _pretty_name(token: str) -> str:
    """MON:CHIEN_PAO -> Chien Pao"""
    if not token:
        return ""
    return token.split(":", 1)[-1].replace("_", " ").title()


def format_move_for_display(actions: Dict[str, dict], player_token: str) -> str:
    """Produce a clean HTML display string from structured per-slot actions."""
    lines = []
    for slot in ("[SIDE_A]", "[SIDE_B]"):
        a = actions.get(slot)
        if a is None:
            continue
        tera_prefix = ""
        if a.get("tera"):
            tera_prefix = f"<i>(Tera {_pretty_name(a['tera'])})</i> "

        if a["type"] == "move":
            mon = _pretty_name(a.get("mon", ""))
            move = _pretty_name(a["move"])
            target = _pretty_name(a.get("target_mon", ""))
            tgt_side = a.get("target_side", "")
            line = f"{tera_prefix}<b>{mon}</b> uses <b>{move}</b>"
            if target:
                side_label = "A" if tgt_side == "[SIDE_A]" else "B" if tgt_side == "[SIDE_B]" else ""
                t_player = a.get("target_player", "")
                if side_label and t_player:
                    line += f" -> {target} ({t_player} slot {side_label})"
                elif target:
                    line += f" -> {target}"
            lines.append(line)
        elif a["type"] == "switch":
            mon = _pretty_name(a["mon"])
            lines.append(f"{tera_prefix}Switch to <b>{mon}</b>")
    return " <br/> ".join(lines) if lines else "Unknown action"


def _format_from_raw_tokens(decoded_tokens: List[str], player_token: str) -> str:
    """Fallback formatter when structured parsing yields nothing."""
    actions = extract_slot_actions(decoded_tokens, player_token)
    if actions:
        return format_move_for_display(actions, player_token)
    raw = " ".join(decoded_tokens)
    raw = raw.replace("[CMD_MOVE]", " | ").replace("[P1]", "").replace("[P2]", "")
    for prefix in ("MON:", "MOVE:", "TYPE:", "SWAP:", "ABIL:", "ITEM:"):
        raw = raw.replace(prefix, "")
    return raw.replace("_", " ").strip() or "Unknown"


# ---------------------------------------------------------------------------
#  Confidence-adjusted scoring
# ---------------------------------------------------------------------------

def lower_confidence_bound(avg: float, n: int, z: float = 1.0) -> float:
    """Pessimistic score that penalises low sample counts.

    Uses LCB: avg - z / sqrt(n).  A single sample at 0.90 scores 0.0,
    while 31 samples at 0.80 score ~0.62.
    """
    if n <= 0:
        return 0.0
    return avg - z / math.sqrt(n)


# ---------------------------------------------------------------------------
#  MonteCarloSearcher
# ---------------------------------------------------------------------------

class MonteCarloSearcher:
    def __init__(self, model: AlphaMonGPT, vm: VocabManager, device="cpu"):
        self.model = model.to(device)
        self.vm = vm
        self.device = device
        self.model.eval()

    def _get_roster(self, context_tokens: List[int], as_player: int) -> List[str]:
        player_token = self.vm.PLAYER_1 if as_player == 1 else self.vm.PLAYER_2
        roster = []
        i = 0
        while i < len(context_tokens):
            t_str = self.vm.decode(context_tokens[i])
            # Roster declaration always finishes before the first [START]; bail
            # out early instead of walking the entire battle log.
            if t_str == self.vm.START:
                break
            if t_str == "[POKE]":
                if i + 1 < len(context_tokens):
                    p_ref = self.vm.decode(context_tokens[i + 1])
                    if p_ref == player_token:
                        if i + 2 < len(context_tokens):
                            mon_token = self.vm.decode(context_tokens[i + 2])
                            if mon_token.startswith("MON:"):
                                roster.append(mon_token)
            i += 1
        return roster

    # ------------------------------------------------------------------
    #  Team Preview (exhaustive)
    # ------------------------------------------------------------------

    def _evaluate_team_permutations(
        self, context_tokens: List[int], roster: List[str], as_player: int
    ) -> List[Dict]:
        n_mons = len(roster)
        if n_mons == 0:
            return []

        if n_mons < 4:
            permutations = [tuple(range(1, n_mons + 1))]
        else:
            permutations = list(itertools.permutations(range(1, n_mons + 1), 4))

        player_token_id = self.vm.vocab[
            self.vm.PLAYER_1 if as_player == 1 else self.vm.PLAYER_2
        ]

        candidates = []
        for p in permutations:
            seq = [player_token_id]
            for idx in p:
                token_str = self.vm.formatted_token("SWAP", str(idx))
                if token_str in self.vm.vocab:
                    seq.append(self.vm.vocab[token_str])
            if len(seq) == len(p) + 1:
                candidates.append((p, seq))

        if not candidates:
            return []

        context_tensor = torch.tensor(
            context_tokens, dtype=torch.long, device=self.device
        )

        batch_size = 64
        results = []

        for i in range(0, len(candidates), batch_size):
            batch_candidates = candidates[i : i + batch_size]
            # Build the per-permutation suffix as a single CPU tensor and move it
            # to the device once (instead of one tiny .to(device) per perm).
            seq_array = np.array([seq for _, seq in batch_candidates], dtype=np.int64)
            suffix_tensor = torch.from_numpy(seq_array).to(self.device)
            ctx_repeat = context_tensor.unsqueeze(0).expand(suffix_tensor.size(0), -1)
            input_tensor = torch.cat([ctx_repeat, suffix_tensor], dim=1)
            with torch.no_grad():
                _, values, _, _ = self.model(input_tensor)
                batch_values = values[:, -1, 0].cpu().numpy()
            del input_tensor, values
            if self.device == "cuda":
                torch.cuda.empty_cache()

            for j, (p, _) in enumerate(batch_candidates):
                val = float(batch_values[j])
                team_names = []
                for idx in p:
                    if 0 <= idx - 1 < len(roster):
                        team_names.append(_pretty_name(roster[idx - 1]))

                leads = ", ".join(team_names[:2])
                back = ", ".join(team_names[2:])
                action_str = f"Leads: {leads} | Back: {back}"

                results.append({
                    "action_str": action_str,
                    "sig_sequence": list(p),
                    "value": val,
                    "p2_win_prob": 1.0 - val,
                })

        # Aggregate (lead-order and back-order independent)
        del context_tensor
        if self.device == "cuda":
            torch.cuda.empty_cache()
        aggregated: Dict[tuple, dict] = {}
        for r in results:
            p = r["sig_sequence"]
            leads = tuple(sorted(p[:2]))
            back = tuple(sorted(p[2:]))
            key = (leads, back)

            if key not in aggregated:
                aggregated[key] = {
                    "count": 0, "sum_value": 0.0, "sum_p2_prob": 0.0,
                    "example_p": p,
                }
            aggregated[key]["count"] += 1
            aggregated[key]["sum_value"] += r["value"]
            aggregated[key]["sum_p2_prob"] += r["p2_win_prob"]

        final_list = []
        for (leads, back), data in aggregated.items():
            avg_val = data["sum_value"] / data["count"]
            avg_p2 = data["sum_p2_prob"] / data["count"]

            team_names_leads = [
                _pretty_name(roster[idx - 1])
                for idx in leads if 0 <= idx - 1 < len(roster)
            ]
            team_names_back = [
                _pretty_name(roster[idx - 1])
                for idx in back if 0 <= idx - 1 < len(roster)
            ]
            action_str = f"Leads: {', '.join(team_names_leads)} | Back: {', '.join(team_names_back)}"

            final_list.append({
                "action_str": action_str,
                "sig_sequence": list(data["example_p"]),
                "value": avg_val,
                "p2_win_prob": avg_p2,
                "count": data["count"],
            })

        return final_list

    # ------------------------------------------------------------------
    #  Faint / forced switch (exhaustive)
    # ------------------------------------------------------------------

    def _evaluate_faint_switches(
        self, context_tokens: List[int], roster: List[str], as_player: int
    ) -> List[Dict]:
        player_token_id = self.vm.vocab[
            self.vm.PLAYER_1 if as_player == 1 else self.vm.PLAYER_2
        ]
        switch_token_id = self.vm.vocab[self.vm.SWITCH]
        context_tensor = torch.tensor(
            context_tokens, dtype=torch.long, device=self.device
        )

        batch_input_ids = []
        valid_mons = []

        for mon_entry in roster:
            mon_token_id = self.vm.vocab.get(mon_entry)
            if mon_token_id is None:
                continue
            valid_mons.append(mon_entry)
            next_tokens = torch.tensor(
                [player_token_id, switch_token_id, mon_token_id],
                dtype=torch.long, device=self.device,
            )
            batch_input_ids.append(torch.cat([context_tensor, next_tokens]))

        results = []
        if batch_input_ids:
            input_batch = torch.stack(batch_input_ids)
            with torch.no_grad():
                _, values, _, _ = self.model(input_batch)
                batch_values = values[:, -1, 0].cpu().numpy()
            del input_batch, values
            if self.device == "cuda":
                torch.cuda.empty_cache()

            for i, mon_entry in enumerate(valid_mons):
                mon_name = _pretty_name(mon_entry)
                val = float(batch_values[i])
                results.append({
                    "action_str": f"Switch to {mon_name}",
                    "value": val,
                    "p2_win_prob": 1.0 - val,
                    "slot_actions": {
                        "[SIDE_A]": {"type": "switch", "mon": mon_entry},
                    },
                })
        return results

    # ------------------------------------------------------------------
    #  Main entry point
    # ------------------------------------------------------------------

    def search_move(
        self,
        context_tokens: List[int],
        num_rollouts: int = 128,
        rollout_len: int = 64,
        temperature: float = 1.0,
        top_k: int = 64,
        as_player: int = 2,
        objective: str = TURN,
        logit_bias: Optional[torch.Tensor] = None,
    ) -> List[Dict]:
        """Performs Monte Carlo rollouts to find the best next move.

        Returns a list of dicts sorted by confidence-adjusted score (LCB),
        each containing: action, action_signature, slot_actions, count,
        avg_score, lcb_score, max_score.
        """
        player_token = self.vm.PLAYER_1 if as_player == 1 else self.vm.PLAYER_2

        # ========== TEAM PREVIEW (exhaustive) ==========
        if objective == TEAM_PREVIEW:
            roster = self._get_roster(context_tokens, as_player)
            perm_results = self._evaluate_team_permutations(
                context_tokens, roster, as_player
            )
            summary = []
            for r in perm_results:
                score = _win_score(r["value"], as_player)
                summary.append({
                    "action": r["action_str"],
                    "action_signature": r["action_str"],
                    "slot_actions": {},
                    "count": r.get("count", 1),
                    "avg_score": score,
                    "lcb_score": score,
                    "max_score": score,
                })
            summary.sort(key=lambda x: x["avg_score"], reverse=True)
            return summary[:top_k]

        # ========== FAINT (exhaustive) ==========
        if objective == FAINT:
            roster = self._get_roster(context_tokens, as_player)
            faint_results = self._evaluate_faint_switches(
                context_tokens, roster, as_player
            )
            for r in faint_results:
                r["score"] = _win_score(r["value"], as_player)
            faint_results.sort(key=lambda x: x["score"], reverse=True)
            summary = []
            for r in faint_results:
                s = r["score"]
                summary.append({
                    "action": r["action_str"],
                    "action_signature": r["action_str"],
                    "slot_actions": r.get("slot_actions", {}),
                    "count": 1,
                    "avg_score": s,
                    "lcb_score": s,
                    "max_score": s,
                })
            return summary

        # ========== TURN (Monte Carlo rollouts) ==========
        context_len = len(context_tokens)
        input_tensor = torch.tensor(
            [context_tokens], dtype=torch.long, device=self.device
        ).repeat(num_rollouts, 1)

        with torch.no_grad():
            generated_ids = self.model.generate(
                input_tensor,
                max_new_tokens=rollout_len,
                temperature=temperature,
                top_k=top_k,
                logit_bias=logit_bias,
            )
            _, values, _, _ = self.model(generated_ids)
            # (B, T, 1) -> (B, T): keep the full per-position win-prob curve so
            # we can pick the value at the *first terminator* of each rollout
            # (i.e. the position the value head was actually trained at).
            values_per_pos = values[:, :, 0].cpu()
            all_new_tokens = generated_ids[:, context_len:].cpu().tolist()

        del input_tensor, generated_ids, values
        if self.device == "cuda":
            torch.cuda.empty_cache()

        # ---------- group rollouts by structured signature ----------
        grouped: Dict[str, dict] = {}

        for i in range(len(all_new_tokens)):
            new_tokens = all_new_tokens[i]

            decoded: List[str] = []
            terminator_offset = None  # index into new_tokens where rollout ends
            for j, t in enumerate(new_tokens):
                token_str = self.vm.decode(t)
                if token_str in (self.vm.EOS_TOKEN, "[TURN]"):
                    terminator_offset = j
                    break
                decoded.append(token_str)

            # Score position: the first terminator (clamped to last token if
            # no terminator was generated). value index in the full sequence
            # is context_len + terminator_offset.
            if terminator_offset is None:
                terminator_offset = len(new_tokens) - 1
            score_idx = context_len + terminator_offset
            score = float(values_per_pos[i, score_idx])

            slot_actions = extract_slot_actions(decoded, player_token)
            sig = _action_signature_key(slot_actions)

            if not sig:
                sig = " ".join(decoded)

            if sig not in grouped:
                grouped[sig] = {
                    "slot_actions": slot_actions,
                    "decoded_samples": [],
                    "scores": [],
                }
            grouped[sig]["decoded_samples"].append(decoded)
            grouped[sig]["scores"].append(score)

        # ---------- aggregate & rank ----------
        summary = []
        for sig, data in grouped.items():
            scores = np.array(data["scores"])
            win_scores = (1.0 - scores) if as_player == 2 else scores

            avg = float(np.mean(win_scores))
            count = len(win_scores)
            lcb = lower_confidence_bound(avg, count)

            slot_actions = data["slot_actions"]
            if slot_actions:
                display = format_move_for_display(slot_actions, player_token)
            else:
                display = _format_from_raw_tokens(
                    data["decoded_samples"][0], player_token
                )

            summary.append({
                "action_signature": sig,
                "action": display,
                "slot_actions": slot_actions,
                "raw_tokens": data["decoded_samples"][0],
                "count": count,
                "avg_score": avg,
                "lcb_score": lcb,
                "max_score": float(np.max(win_scores)),
            })

        # ---------- secondary merge: same display string ----------
        merged: Dict[str, dict] = {}
        for item in summary:
            key = item["action"]
            if key not in merged:
                merged[key] = {
                    "action_signature": item["action_signature"],
                    "action": key,
                    "slot_actions": item["slot_actions"],
                    "raw_tokens": item.get("raw_tokens", []),
                    "count": 0,
                    "total_score": 0.0,
                    "max_score": item["max_score"],
                }
            m = merged[key]
            m["count"] += item["count"]
            m["total_score"] += item["avg_score"] * item["count"]
            m["max_score"] = max(m["max_score"], item["max_score"])

        final = []
        for v in merged.values():
            v["avg_score"] = v["total_score"] / v["count"] if v["count"] else 0.0
            v["lcb_score"] = lower_confidence_bound(v["avg_score"], v["count"])
            del v["total_score"]
            final.append(v)

        # Sort by LCB -- high sample + high avg wins
        final.sort(key=lambda x: x["lcb_score"], reverse=True)
        return final