Add Brain + intrigue on location support (#17)

* Add brain ability

* Support for intrigue on locations

* Fix for 3.11

Author: wnojopra

app/models/api_models.py

@@ -36,7 +36,7 @@ def from_game_state(game_id: str, game_state: GameState) -> "GameStateResponse":
             characters=[
                 CharacterState(
                     name=c.name,
-                    location=c.location.name,
+                    location=c.location.location_type,
                     paranoia=c.paranoia,
                     goodwill=c.goodwill,
                     intrigue=c.intrigue,

engine/ability_effects.py

@@ -1,18 +1,24 @@
-from .models import Character, RoleType
+from .models import Character, Location, RoleType
 from .state import GameState
 
-def resolve_conspiracy_theorist_ability(source_char: Character, target_char: Character, game_state: GameState):
-    """
-    The Conspiracy Theorist can use ADD_PARANOIA on anyone at their location.
-    """
+def resolve_conspiracy_theorist_ability(source: Character, target: Character | Location):
+    assert isinstance(target, Character)
     # Find characters at the same location (excluding themself)
-    if source_char.location != target_char.location:
+    if source.location != target.location:
         raise ValueError(f"Error resolving conspiracy theorist ability: target at different location")
     else:
-        target_char.paranoia += 1
+        target.paranoia += 1
+
+def resolve_brain_ability(source: Character, target: Character | Location):
+    if isinstance(target, Character) and source.location != target.location:
+        raise ValueError(f"Error resolving conspiracy theorist ability: target at different location")
+    else:
+        target.intrigue += 1
+
 
 # The main registry for all ability logic
 ABILITY_LOGIC = {
     RoleType.CONSPIRACY_THEORIST: resolve_conspiracy_theorist_ability,
+    RoleType.BRAIN: resolve_brain_ability,
     # Add other roles with abilities here
 }

engine/engine.py

@@ -2,18 +2,18 @@
 
 from engine.ability_effects import ABILITY_LOGIC
 from .incident_effects import INCIDENT_EFFECTS
-from .models import Action, Character, Location, ActionType, ALL_LOCATIONS, RoleType, TurnData
+from .models import Action, Character, Location, ActionType, ALL_LOCATIONS, LocationType, RoleType, TurnData
 from .role_effects import ROLE_EFFECTS
 from .state import GameState
 from collections import defaultdict
 
 # Define the layout as a 2x2 grid
 LOCATION_GRID = [
-    [Location.HOSPITAL, Location.SHRINE],   # Top row
-    [Location.CITY, Location.SCHOOL]        # Bottom row
+    [LocationType.HOSPITAL, LocationType.SHRINE],   # Top row
+    [LocationType.CITY, LocationType.SCHOOL]        # Bottom row
 ]
 
-def find_location_coords(location: Location) -> tuple[int, int]:
+def find_location_coords(location: LocationType) -> tuple[int, int]:
     for row_idx, row in enumerate(LOCATION_GRID):
         for col_idx, loc in enumerate(row):
             if loc == location:
@@ -80,7 +80,7 @@ def resolve_actions(game_state: GameState, turn_data: TurnData):
     # Resolve actions targeting characters
     for char, action_list in char_to_actions.items():
         for action in action_list:
-            resolve_action(char, action.type) # type: ignore
+            resolve_action(char, action.type)
 
     # TODO: Resolve location-based actions (not yet implemented)
     # for loc, action_list in loc_to_actions.items():
@@ -98,18 +98,22 @@ def resolve_abilities(game_state: GameState, turn_data: TurnData):
 
     for choice in turn_data.ability_actions:
         source_char = next((c for c in game_state.characters if c.role == choice.source), None)
-        target_char = next((c for c in game_state.characters if c.name == choice.target), None)
+        if choice.target.lower() in {l.value for l in LocationType}:
+            target_loc_enum = LocationType(choice.target.lower())
+            target_obj = game_state.location_states[target_loc_enum]
+        else:
+            target_obj = next((c for c in game_state.characters if c.name == choice.target), None)
 
-        if not source_char or not target_char:
-            print(f"Warning: Invalid character in ability choice, skipping.")
+        if not source_char or not target_obj:
+            print(f"Warning: Invalid character or location in ability choice, skipping.")
             continue
 
         ability_fn = ABILITY_LOGIC.get(source_char.role)
         if not ability_fn:
             print(f"Warning: {source_char.name} ({source_char.role.name}) has no defined ability, skipping.")
             continue
 
-        ability_fn(source_char, target_char, game_state)
+        ability_fn(source_char, target_obj)
 
 
 def resolve_roles(game_state: GameState):

engine/models.py

@@ -20,14 +20,14 @@ class RoleType(str, Enum):
     WITCH = "witch"
 
 
-class Location(str, Enum):
+class LocationType(str, Enum):
     CITY = "city"
     HOSPITAL = "hospital"
     SCHOOL = "school"
     SHRINE = "shrine"
 
 
-ALL_LOCATIONS = list(Location)
+ALL_LOCATIONS = list(LocationType)
 
 
 class ActionType(str, Enum):
@@ -58,9 +58,9 @@ class IncidentType(str, Enum):
 
 class Character(BaseModel):
     name: str
-    location: Location
-    starting_location: Location
-    disallowed_locations: List[Location]
+    location: LocationType
+    starting_location: LocationType
+    disallowed_locations: List[LocationType]
     paranoia: int = 0
     paranoia_limit: int = 0
     goodwill: int = 0
@@ -77,6 +77,11 @@ def __eq__(self, other):
         return self.name == other.name
 
 
+class Location(BaseModel):
+    location_type: LocationType
+    intrigue: int = 0
+
+
 class Incident(BaseModel):
     type: IncidentType
     day: int
@@ -117,4 +122,4 @@ class TurnData(BaseModel):
     incident_choices: Optional[List[IncidentChoice]] = None
     ability_actions: Optional[List[AbilityChoice]] = None
 
-AllActionsByDay = Dict[int, TurnData]
+AllActionsByDay = Dict[int, TurnData]

engine/simulation.py

@@ -1,4 +1,6 @@
 import argparse
+
+from engine.models import Location, LocationType
 from .state import GameState
 from .victory_checker import check_victory
 from .yaml_loader import load_actions_from_yaml, load_script_from_yaml
@@ -8,6 +10,9 @@ def create_starting_game_state(script_path: str, actions_path: str) -> GameState
     day = 1
     script = load_script_from_yaml(script_path)
     actions = load_actions_from_yaml(actions_path)
+    location_states = {
+        location_type: Location(location_type=location_type) for location_type in LocationType
+    }
 
     return GameState(
         day=day,
@@ -16,7 +21,8 @@ def create_starting_game_state(script_path: str, actions_path: str) -> GameState
         max_loops=script.max_loops,
         characters=script.characters,
         incidents=script.incidents,
-        actions=actions
+        actions=actions,
+        location_states=location_states
     )
 
 def simulate_day(game_state: GameState):

engine/state.py

@@ -1,7 +1,7 @@
 from dataclasses import dataclass, field
-from typing import List, Optional, Set
+from typing import Dict, List, Optional, Set
 
-from .models import Character, Incident, AllActionsByDay, RoleType
+from .models import Character, Incident, AllActionsByDay, Location, LocationType, RoleType
 
 
 @dataclass
@@ -13,13 +13,14 @@ class GameState:
     characters: List[Character]
     incidents: List[Incident]
     actions: AllActionsByDay
+    location_states: Dict[LocationType, Location]
     game_result: Optional[str] = None  # "protagonists_win", "mastermind_win", or None
     revealed_roles: Set[RoleType] = field(default_factory=set)
 
 
     def print_characters(self):
         for char in self.characters:
-            message = f"{char.name} is at {char.location}, paranoia={char.paranoia}, goodwill={char.goodwill}, intrigue={char.intrigue}"
+            message = f"{char.name} is at {char.location.name}, paranoia={char.paranoia}, goodwill={char.goodwill}, intrigue={char.intrigue}"
             if not char.alive:
                 message += ", status=DEAD"
             print(message)

engine/yaml_loader.py

@@ -1,7 +1,7 @@
 import yaml
 from .models import (
     AbilityChoice, Action, ActionSource, ActionTargetType, ActionType,
-    AllActionsByDay, Character, IncidentChoice, IncidentType, Location,
+    AllActionsByDay, Character, IncidentChoice, IncidentType, Location, LocationType,
     RoleType, Incident, Script, TurnData
 )
 
@@ -13,9 +13,9 @@ def load_script_from_yaml(path: str) -> Script:
     for entry in data["characters"]:
         character = Character(
             name=entry["name"],
-            location=Location[entry["starting_location"]],
-            starting_location=Location[entry["starting_location"]],
-            disallowed_locations=[Location[loc] for loc in entry["disallowed_locations"]],
+            location=LocationType[entry["starting_location"]],
+            starting_location=LocationType[entry["starting_location"]],
+            disallowed_locations=[LocationType[loc] for loc in entry["disallowed_locations"]],
             paranoia_limit=entry["paranoia_limit"],
             role=RoleType[entry["role"]]
         )
@@ -56,7 +56,7 @@ def parse_actions(source: ActionSource, actions_raw: list, day: int) -> list[Act
 
             target_str = action_dict["target"]
             # Infer target_type based on name
-            if target_str.lower() in {loc.value for loc in Location}:
+            if target_str.lower() in {loc.value for loc in LocationType}:
                 target_type = ActionTargetType.LOCATION
             else:
                 target_type = ActionTargetType.CHARACTER

scripts/the_first_script/actions_4.yaml

@@ -8,7 +8,7 @@
       type: ADD_INTRIGUE
 
     - target: Doctor
-      type: ADD_INTRIGUE
+      type: MOVE_VERTICAL
 
   protagonist:
     - target: Shrine Maiden
@@ -20,6 +20,10 @@
     - target: Doctor
       type: ADD_GOODWILL
 
+  ability_actions:
+    - source: BRAIN
+      target: Office Worker
+
 - day: 2
   mastermind:
     - target: Girl Student
@@ -41,6 +45,10 @@
     - target: Doctor
       type: ADD_GOODWILL
 
+  ability_actions:
+    - source: BRAIN
+      target: Office Worker
+
 - day: 3
   mastermind:
     - target: Girl Student
@@ -62,6 +70,10 @@
     - target: Doctor
       type: ADD_GOODWILL
 
+  ability_actions:
+    - source: BRAIN
+      target: Office Worker
+
 - day: 4
   mastermind:
     - target: Girl Student
@@ -82,3 +94,7 @@
 
     - target: Doctor
       type: ADD_GOODWILL
+
+  ability_actions:
+    - source: BRAIN
+      target: CITY

tests/test_scenarios.py

@@ -1,5 +1,5 @@
 from engine.simulation import run_full_simulation
-from engine.models import RoleType
+from engine.models import LocationType, RoleType
 
 def test_friend_dies_at_end_of_loop():
     """
@@ -58,4 +58,14 @@ def test_conspiracy_theorist_ability():
     final_state = run_full_simulation(script_path, actions_path)
     assert final_state.game_result == "mastermind_win"
     key_person = next(c for c in final_state.characters if c.role == RoleType.KEY_PERSON)
-    assert not key_person.alive
+    assert not key_person.alive
+
+def test_brain_ability():
+    script_path = "scripts/the_first_script/script.yaml"
+    actions_path = "scripts/the_first_script/actions_4.yaml"
+    final_state = run_full_simulation(script_path, actions_path)
+    assert final_state.game_result == "protagonist_win"
+    assert all(char.alive for char in final_state.characters)
+    office_worker = next(c for c in final_state.characters if c.name == "Office Worker")
+    assert office_worker.intrigue == 7
+    assert final_state.location_states[LocationType.CITY].intrigue == 1