【ピザラ人狼】9人村をPythonで全パターン計算したら4096万通りに爆発した

"""ピザラ人狼9人村の簡略ランダムモデルを厳密計算する。

昼の処刑先、夜の襲撃先、騎士の護衛先を、それぞれ合法候補から
等確率で選ぶ。会話、投票、占い・霊媒結果、役職COはモデル化しない。
"""

from __future__ import annotations

import argparse
import json
from collections import defaultdict
from dataclasses import dataclass
from fractions import Fraction
from itertools import product
from pathlib import Path
from typing import Iterable, Optional


MODEL_VERSION = "1.0.0"

PLAYERS = (
    "citizen_1",
    "citizen_2",
    "citizen_3",
    "seer",
    "medium",
    "knight",
    "werewolf_1",
    "werewolf_2",
    "madman",
)

CITIZEN_1, CITIZEN_2, CITIZEN_3, SEER, MEDIUM, KNIGHT, WEREWOLF_1, WEREWOLF_2, MADMAN = range(9)
WEREWOLVES = frozenset((WEREWOLF_1, WEREWOLF_2))
WEREWOLF_MASK = sum(1 << player for player in WEREWOLVES)
ALL_ALIVE = (1 << len(PLAYERS)) - 1

VILLAGE = "village"
WOLF = "wolf"
PP = "pp"
RPP = "rpp"


@dataclass(frozen=True)
class State:
    """同じ盤面へ合流した履歴をまとめるための状態。"""

    alive: int
    previous_guard: Optional[int] = None
    seen_gj: bool = False
    seen_pp: bool = False
    seen_rpp: bool = False


@dataclass
class Mass:
    """ある状態へ到達した履歴数と確率質量。"""

    paths: int = 0
    probability: Fraction = Fraction(0, 1)

    def add(self, paths: int, probability: Fraction) -> None:
        self.paths += paths
        self.probability += probability


@dataclass(frozen=True)
class Analysis:
    terminal_paths: dict[str, int]
    terminal_probability: dict[str, Fraction]
    finish_paths: dict[tuple[int, str, str], int]
    finish_probability: dict[tuple[int, str, str], Fraction]
    event_paths: dict[str, int]
    event_probability: dict[str, Fraction]
    board_paths_by_day: dict[tuple[int, str], int]
    board_probability_by_day: dict[tuple[int, str], Fraction]
    day_state_count: dict[int, int]
    night_state_count: dict[int, int]


def is_alive(alive: int, player: int) -> bool:
    return bool(alive & (1 << player))


def living_players(alive: int) -> tuple[int, ...]:
    return tuple(player for player in range(len(PLAYERS)) if is_alive(alive, player))


def remove_player(alive: int, player: int) -> int:
    return alive & ~(1 << player)


def living_werewolf_count(alive: int) -> int:
    return (alive & WEREWOLF_MASK).bit_count()


def winner(alive: int) -> Optional[str]:
    """人狼全滅、または人狼数が非人狼数以上なら勝敗を返す。"""

    werewolves = living_werewolf_count(alive)
    if werewolves == 0:
        return VILLAGE

    non_werewolves = alive.bit_count() - werewolves
    if werewolves >= non_werewolves:
        return WOLF
    return None


def power_play_board(alive: int) -> Optional[str]:
    """昼開始時の投票人数からPP/RPPを判定する。"""

    if winner(alive) is not None:
        return None
    werewolves = living_werewolf_count(alive)
    if werewolves == 0 or not is_alive(alive, MADMAN):
        return None

    wolf_team_votes = werewolves + 1  # 生存している狂人
    village_team_votes = alive.bit_count() - wolf_team_votes
    if wolf_team_votes > village_team_votes:
        return PP
    if wolf_team_votes == village_team_votes:
        return RPP
    return None


def mark_day_events(state: State) -> State:
    board = power_play_board(state.alive)
    return State(
        alive=state.alive,
        previous_guard=state.previous_guard,
        seen_gj=state.seen_gj,
        seen_pp=state.seen_pp or board == PP,
        seen_rpp=state.seen_rpp or board == RPP,
    )


def execution_targets(state: State) -> tuple[int, ...]:
    return living_players(state.alive)


def attack_targets(state: State) -> tuple[int, ...]:
    return tuple(player for player in living_players(state.alive) if player not in WEREWOLVES)


def guard_targets(state: State) -> tuple[Optional[int], ...]:
    if not is_alive(state.alive, KNIGHT):
        return (None,)
    return tuple(
        player
        for player in living_players(state.alive)
        if player != KNIGHT and player != state.previous_guard
    )


def add_mass(frontier: dict[State, Mass], state: State, paths: int, probability: Fraction) -> None:
    frontier.setdefault(state, Mass()).add(paths, probability)


def analyze() -> Analysis:
    """全ての合法な行動履歴を動的計画法で集計する。"""

    terminal_paths = {VILLAGE: 0, WOLF: 0}
    terminal_probability = {VILLAGE: Fraction(0, 1), WOLF: Fraction(0, 1)}
    finish_paths: dict[tuple[int, str, str], int] = defaultdict(int)
    finish_probability: dict[tuple[int, str, str], Fraction] = defaultdict(Fraction)
    event_paths = {"gj": 0, PP: 0, RPP: 0}
    event_probability = {"gj": Fraction(0, 1), PP: Fraction(0, 1), RPP: Fraction(0, 1)}
    board_paths_by_day: dict[tuple[int, str], int] = defaultdict(int)
    board_probability_by_day: dict[tuple[int, str], Fraction] = defaultdict(Fraction)
    day_state_count: dict[int, int] = {}
    night_state_count: dict[int, int] = {}

    initial = mark_day_events(State(alive=ALL_ALIVE))
    day_frontier: dict[State, Mass] = {initial: Mass(paths=1, probability=Fraction(1, 1))}
    day = 1

    def record_terminal(
        outcome: str,
        stage: str,
        state: State,
        paths: int,
        probability: Fraction,
    ) -> None:
        terminal_paths[outcome] += paths
        terminal_probability[outcome] += probability
        finish_paths[(day, stage, outcome)] += paths
        finish_probability[(day, stage, outcome)] += probability
        for event, seen in (("gj", state.seen_gj), (PP, state.seen_pp), (RPP, state.seen_rpp)):
            if seen:
                event_paths[event] += paths
                event_probability[event] += probability

    while day_frontier:
        if day > len(PLAYERS):
            raise RuntimeError("解析が想定日数内に収束しませんでした")

        day_state_count[day] = len(day_frontier)
        night_frontier: dict[State, Mass] = {}

        for state, mass in day_frontier.items():
            board = power_play_board(state.alive)
            if board is not None:
                board_paths_by_day[(day, board)] += mass.paths
                board_probability_by_day[(day, board)] += mass.probability

            targets = execution_targets(state)
            action_probability = Fraction(1, len(targets))
            for target in targets:
                child_alive = remove_player(state.alive, target)
                child = State(
                    alive=child_alive,
                    previous_guard=state.previous_guard if is_alive(child_alive, KNIGHT) else None,
                    seen_gj=state.seen_gj,
                    seen_pp=state.seen_pp,
                    seen_rpp=state.seen_rpp,
                )
                branch_probability = mass.probability * action_probability
                outcome = winner(child_alive)
                if outcome is not None:
                    record_terminal(outcome, "after_execution", child, mass.paths, branch_probability)
                else:
                    add_mass(night_frontier, child, mass.paths, branch_probability)

        night_state_count[day] = len(night_frontier)
        next_day_frontier: dict[State, Mass] = {}

        for state, mass in night_frontier.items():
            attacks = attack_targets(state)
            guards = guard_targets(state)
            actions: Iterable[tuple[int, Optional[int]]] = product(attacks, guards)
            action_count = len(attacks) * len(guards)
            action_probability = Fraction(1, action_count)

            for attack, guard in actions:
                guarded = guard is not None and attack == guard
                child_alive = state.alive if guarded else remove_player(state.alive, attack)
                child = State(
                    alive=child_alive,
                    previous_guard=guard if is_alive(child_alive, KNIGHT) else None,
                    seen_gj=state.seen_gj or guarded,
                    seen_pp=state.seen_pp,
                    seen_rpp=state.seen_rpp,
                )
                branch_probability = mass.probability * action_probability
                outcome = winner(child_alive)
                if outcome is not None:
                    record_terminal(outcome, "after_night", child, mass.paths, branch_probability)
                else:
                    add_mass(next_day_frontier, mark_day_events(child), mass.paths, branch_probability)

        day_frontier = next_day_frontier
        day += 1

    if terminal_probability[VILLAGE] + terminal_probability[WOLF] != 1:
        raise AssertionError("終局確率の合計が1ではありません")
    if sum(finish_paths.values()) != sum(terminal_paths.values()):
        raise AssertionError("終局履歴数の内訳が一致しません")
    if sum(finish_probability.values(), Fraction(0, 1)) != 1:
        raise AssertionError("日別終局確率の合計が1ではありません")

    return Analysis(
        terminal_paths=terminal_paths,
        terminal_probability=terminal_probability,
        finish_paths=dict(finish_paths),
        finish_probability=dict(finish_probability),
        event_paths=event_paths,
        event_probability=event_probability,
        board_paths_by_day=dict(board_paths_by_day),
        board_probability_by_day=dict(board_probability_by_day),
        day_state_count=day_state_count,
        night_state_count=night_state_count,
    )


def fraction_json(value: Fraction) -> dict[str, object]:
    return {
        "fraction": f"{value.numerator}/{value.denominator}",
        "decimal": round(float(value), 12),
        "percent": round(float(value) * 100, 6),
    }


def result_json(analysis: Analysis) -> dict[str, object]:
    total_paths = sum(analysis.terminal_paths.values())
    history_share = {
        outcome: Fraction(paths, total_paths)
        for outcome, paths in analysis.terminal_paths.items()
    }
    finish_days = sorted({key[0] for key in analysis.finish_paths})
    finish_by_day = []
    for day in finish_days:
        row: dict[str, object] = {"day": day}
        for stage in ("after_execution", "after_night"):
            for outcome in (VILLAGE, WOLF):
                key = (day, stage, outcome)
                row[f"{stage}_{outcome}_paths"] = analysis.finish_paths.get(key, 0)
                row[f"{stage}_{outcome}_probability"] = fraction_json(
                    analysis.finish_probability.get(key, Fraction(0, 1))
                )
        finish_by_day.append(row)

    board_days = sorted({key[0] for key in analysis.board_probability_by_day})
    board_by_day = []
    for day in board_days:
        row: dict[str, object] = {"day": day}
        for board in (PP, RPP):
            key = (day, board)
            row[f"{board}_prefix_paths"] = analysis.board_paths_by_day.get(key, 0)
            row[f"{board}_probability"] = fraction_json(
                analysis.board_probability_by_day.get(key, Fraction(0, 1))
            )
        board_by_day.append(row)

    return {
        "model_version": MODEL_VERSION,
        "players": list(PLAYERS),
        "assumptions": {
            "day_execution": "uniform among all living players",
            "night_attack": "uniform among living non-werewolves",
            "knight_guard": "uniform among living players except self and previous target",
            "first_night_attack": False,
            "named_players_are_distinct": True,
            "seer_medium_discussion_and_votes": "not modeled",
        },
        "terminal_histories": {
            "total": total_paths,
            VILLAGE: analysis.terminal_paths[VILLAGE],
            WOLF: analysis.terminal_paths[WOLF],
            f"{VILLAGE}_share": fraction_json(history_share[VILLAGE]),
            f"{WOLF}_share": fraction_json(history_share[WOLF]),
        },
        "random_choice_probability": {
            VILLAGE: fraction_json(analysis.terminal_probability[VILLAGE]),
            WOLF: fraction_json(analysis.terminal_probability[WOLF]),
        },
        "events_at_least_once": {
            event: {
                "terminal_histories": analysis.event_paths[event],
                "probability": fraction_json(analysis.event_probability[event]),
            }
            for event in ("gj", PP, RPP)
        },
        "finish_by_day": finish_by_day,
        "power_play_board_by_day": board_by_day,
        "aggregated_states": {
            "day": analysis.day_state_count,
            "night": analysis.night_state_count,
            "total": sum(analysis.day_state_count.values()) + sum(analysis.night_state_count.values()),
        },
    }


def format_number(value: int) -> str:
    return f"{value:,}"


def print_text(data: dict[str, object]) -> None:
    histories = data["terminal_histories"]
    probability = data["random_choice_probability"]
    events = data["events_at_least_once"]

    print("ピザラ人狼9人村・簡略ランダムモデル")
    print(f"モデル: v{data['model_version']}")
    print()
    print("終局パターン数（行動履歴を1本ずつ数える）")
    print(f"  合計:       {format_number(histories['total'])}")
    print(
        f"  市民陣営:   {format_number(histories[VILLAGE])} "
        f"({histories[f'{VILLAGE}_share']['percent']:.6f}%)"
    )
    print(
        f"  人狼陣営:   {format_number(histories[WOLF])} "
        f"({histories[f'{WOLF}_share']['percent']:.6f}%)"
    )
    print()
    print("各局面で合法手を等確率に選ぶ場合の勝率")
    print(f"  市民陣営:   {probability[VILLAGE]['percent']:.6f}%")
    print(f"  人狼陣営:   {probability[WOLF]['percent']:.6f}%")
    print()
    print("一度以上その状態になる確率")
    print(f"  GJ成功:     {events['gj']['probability']['percent']:.6f}%")
    print(f"  PP盤面:     {events[PP]['probability']['percent']:.6f}%")
    print(f"  RPP盤面:    {events[RPP]['probability']['percent']:.6f}%")
    print()
    print(f"集約して処理した状態数: {format_number(data['aggregated_states']['total'])}")


def main() -> None:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--format", choices=("text", "json"), default="text")
    parser.add_argument("--output", type=Path, help="同じ出力をUTF-8のファイルへ保存する")
    args = parser.parse_args()

    data = result_json(analyze())
    if args.format == "json":
        output = json.dumps(data, ensure_ascii=False, indent=2)
    else:
        output = None

    if args.output:
        args.output.parent.mkdir(parents=True, exist_ok=True)
        if output is None:
            raise SystemExit("--output は --format json と組み合わせてください")
        args.output.write_text(f"{output}\n", encoding="utf-8")

    if args.output:
        print(f"Saved analysis JSON: {args.output}")
    elif output is not None:
        print(output)
    else:
        print_text(data)


if __name__ == "__main__":
    main()