simulator.py

"""Simulation engine for Analog Town state transitions."""

import json
import traceback
from datetime import datetime

from schemas import (
    AgentProfile,
    AgentState,
    BroadcastEvent,
    SimulationResult,
    StateTransition,
    Town,
)
from prompts import SYSTEM_PROMPT, STATE_TRANSITION_PROMPT
from model_client import ModelClient


class Simulator:
    """Runs state-transition simulations for all agents in a town."""

    def __init__(self, model_client: ModelClient | None = None):
        self.model_client = model_client or ModelClient()

    def _get_initial_state(self, agent: AgentProfile) -> AgentState:
        """Create initial agent state from profile defaults."""
        return AgentState(
            agent_id=agent.id,
            mood_label="baseline",
            trust=0.5,
            anger=0.2,
            fear=0.3,
            hope=0.5,
            curiosity=0.5,
            social_energy=0.5,
            current_belief="No strong opinion yet.",
            active_memory=None,
            unresolved_tension=None,
        )

    def _clamp(self, value: float, min_val: float = 0.0, max_val: float = 1.0) -> float:
        """Clamp a value between min and max."""
        return max(min_val, min(max_val, value))

    def _apply_emotion_deltas(
        self, old_state: AgentState, deltas: dict[str, float]
    ) -> dict[str, float]:
        """Apply emotion deltas to old state, clamping to [0.0, 1.0]."""
        emotions = {}
        for key in ["trust", "anger", "fear", "hope", "curiosity", "social_energy"]:
            old_val = getattr(old_state, key, 0.5)
            delta = deltas.get(key, 0.0)
            emotions[key] = self._clamp(old_val + delta)
        return emotions

    def _validate_transition(
        self, transition_data: dict, agent: AgentProfile
    ) -> tuple[bool, str]:
        """Validate a raw transition dict before parsing into Pydantic model."""
        errors = []

        # Check agent_id matches
        if transition_data.get("agent_id") != agent.id:
            transition_data["agent_id"] = agent.id

        # Check required fields
        required = [
            "event_summary_from_agent_view",
            "noticed_detail",
            "activated_memory",
            "value_conflict",
            "emotion_delta",
            "updated_state",
            "internal_monologue",
            "likely_private_action",
            "likely_public_action",
            "uncertainty",
        ]
        for field in required:
            if field not in transition_data:
                errors.append(f"Missing required field: {field}")

        # Check emotion values in updated_state
        updated = transition_data.get("updated_state", {})
        for emotion in ["trust", "anger", "fear", "hope", "curiosity", "social_energy"]:
            val = updated.get(emotion)
            if val is not None:
                if not isinstance(val, (int, float)):
                    errors.append(f"Emotion '{emotion}' must be a number, got {type(val)}")
                elif val < 0.0 or val > 1.0:
                    # Clamp instead of error
                    updated[emotion] = self._clamp(float(val))

        # Ensure agent_id in updated_state
        if "updated_state" in transition_data:
            transition_data["updated_state"]["agent_id"] = agent.id

        # Check internal monologue is not empty
        monologue = transition_data.get("internal_monologue", "")
        if not monologue or not monologue.strip():
            errors.append("Internal monologue is empty")

        # Ensure safety note
        if "safety_note" not in transition_data or not transition_data["safety_note"]:
            transition_data["safety_note"] = (
                "This is a fictional perspective rehearsal, not a prediction."
            )

        if errors:
            return False, "; ".join(errors)
        return True, ""

    def run_agent_transition(
        self,
        agent: AgentProfile,
        state: AgentState,
        event: BroadcastEvent,
        day: int = 1,
    ) -> StateTransition:
        """Run a single agent through the state transition.

        Args:
            agent: The agent's profile
            state: The agent's current state
            event: The broadcast event
            day: 1-indexed broadcast day, used to instruct the model to vary follow-up beats

        Returns:
            StateTransition with updated state and monologue

        Raises:
            RuntimeError: If both generation and repair fail
        """
        user_prompt = STATE_TRANSITION_PROMPT.format(
            agent_profile=agent.model_dump_json(indent=2),
            previous_state=state.model_dump_json(indent=2),
            broadcast_event=event.model_dump_json(indent=2),
        )

        if day > 1:
            user_prompt = (
                f"DAY {day} OF THIS SCENARIO. "
                f"The agent has already lived through {day - 1} previous broadcast(s). "
                "The fields current_belief, active_memory, and unresolved_tension in 'Previous State' "
                "capture where the agent ENDED LAST TIME — they are CONTEXT, not your script. "
                "Write a FRESH internal_monologue that is clearly different in wording from any prior beat: "
                "the agent has had time to process, talk to others, sleep on it, or harden their view. "
                "Make at least one emotion_delta non-zero. Do NOT echo any sentence verbatim from current_belief.\n\n"
                + user_prompt
            )

        original_temp = getattr(self.model_client, "temperature", 0.3)
        try:
            if day > 1:
                self.model_client.temperature = min(0.85, original_temp + 0.25)
            transition_data = self.model_client.generate_json(
                system_prompt=SYSTEM_PROMPT,
                user_prompt=user_prompt,
            )
        finally:
            self.model_client.temperature = original_temp

        # Validate and fix
        is_valid, error_msg = self._validate_transition(transition_data, agent)
        if not is_valid:
            raise RuntimeError(f"Transition validation failed: {error_msg}")

        # Parse into Pydantic model
        transition = StateTransition(**transition_data)
        return transition

    def run_town_simulation(
        self,
        town: Town,
        event: BroadcastEvent,
        previous_states: dict[str, AgentState] | None = None,
        progress_callback=None,
        day: int = 1,
    ) -> SimulationResult:
        """Run simulation for all agents in the town.

        Args:
            town: The town with agents
            event: The broadcast event
            previous_states: Optional dict of agent_id -> AgentState seeded from prior run
            progress_callback: Optional callback(agent_name, status, index, total)

        Returns:
            SimulationResult with all transitions (failed agents are skipped)
        """
        transitions = []
        total = len(town.agents)

        for i, agent in enumerate(town.agents):
            agent_name = agent.name
            try:
                if progress_callback:
                    progress_callback(agent_name, "processing", i, total)

                state = (
                    previous_states.get(agent.id, self._get_initial_state(agent))
                    if previous_states
                    else self._get_initial_state(agent)
                )

                transition = self.run_agent_transition(agent, state, event, day=day)
                transitions.append(transition)

                if progress_callback:
                    progress_callback(agent_name, "complete", i, total)

            except Exception as e:
                # Don't crash on individual agent failures
                print(f"⚠ Agent '{agent_name}' failed: {e}")
                traceback.print_exc()
                if progress_callback:
                    progress_callback(agent_name, f"failed: {str(e)[:100]}", i, total)
                continue

        return SimulationResult(
            town_id=town.id,
            event=event,
            transitions=transitions,
            created_at=datetime.now().isoformat(),
        )