core/scheduler.py

"""
core.scheduler
==============

Simulation scheduler for WorldSmithAI.

This module defines how a World advances by one simulation step. The scheduler
coordinates event processing, agent activation, resource updates, and metric
collection without containing domain-specific logic.

Schedulers do not know what "wolf", "scientist", "dragon", "startup",
"kingdom", "food", "mana", or "knowledge" means. Those semantics are supplied
by the DSL, behaviors, policies, events, and world state.

Minimal usage example
---------------------

    from core.agent import Agent
    from core.scheduler import Scheduler, ActivationMode
    from core.world import World

    world = World(
        scheduler=Scheduler(activation_strategy=ActivationMode.DETERMINISTIC)
    )

    world.add_agent(
        Agent(
            id="agent_1",
            type="researcher",
            position=[0.0, 0.0],
            state={"priority": 10.0},
        )
    )

    result = world.step()
    print(result.to_dict())

Activation strategies
---------------------

    Scheduler(activation_strategy=ActivationMode.DETERMINISTIC)
    Scheduler(activation_strategy=ActivationMode.RANDOM)
    Scheduler(activation_strategy=PriorityActivationStrategy(priority_key="urgency"))
"""

from __future__ import annotations

import logging
from abc import ABC, abstractmethod
from collections.abc import Iterable, Mapping, Sequence
from dataclasses import dataclass, field
from enum import StrEnum
from numbers import Real
from typing import Any

import numpy as np

from core.agent import Agent, AgentStepResult
from core.event import EventExecutionResult
from core.resource import ResourceOperationResult
from core.world import World, WorldStepResult

logger = logging.getLogger(__name__)


class ActivationMode(StrEnum):
    """
    Built-in agent activation modes.

    Attributes
    ----------
    DETERMINISTIC:
        Activate alive agents in world insertion order.
    RANDOM:
        Activate alive agents in a random order using ``world.rng``.
    PRIORITY:
        Activate alive agents by descending numeric priority.
    SIMULTANEOUS:
        Freeze the alive-agent activation set before execution. This prepares
        the API for true simultaneous scheduling but does not yet provide
        transactional behavior updates.
    """

    DETERMINISTIC = "deterministic"
    RANDOM = "random"
    PRIORITY = "priority"
    SIMULTANEOUS = "simultaneous"


class EventPhase(StrEnum):
    """
    Event processing phase.

    Attributes
    ----------
    NONE:
        Do not process events in this scheduler step.
    BEFORE_AGENTS:
        Process events before agent activation.
    AFTER_AGENTS:
        Process events after agent activation.
    """

    NONE = "none"
    BEFORE_AGENTS = "before_agents"
    AFTER_AGENTS = "after_agents"


class ResourcePhase(StrEnum):
    """
    Resource update phase.

    Attributes
    ----------
    NONE:
        Do not update resources in this scheduler step.
    BEFORE_AGENTS:
        Update resources before agent activation.
    AFTER_AGENTS:
        Update resources after agent activation.
    """

    NONE = "none"
    BEFORE_AGENTS = "before_agents"
    AFTER_AGENTS = "after_agents"


class AgentActivationStrategy(ABC):
    """
    Abstract base class for agent activation ordering.

    Activation strategies decide which agents should act and in what order.
    They should not mutate world state. They only return an ordered sequence of
    candidate agents.

    Concrete strategies can implement deterministic, random, priority,
    location-based, faction-based, batched, or learned scheduling.
    """

    name: str

    @abstractmethod
    def order_agents(self, world: World) -> tuple[Agent, ...]:
        """
        Return agents in activation order.

        Parameters
        ----------
        world:
            World whose agents should be ordered.

        Returns
        -------
        tuple[Agent, ...]
            Agents selected for activation.
        """
        raise NotImplementedError

    @staticmethod
    def alive_agents(world: World) -> tuple[Agent, ...]:
        """
        Return alive agents in world insertion order.

        Parameters
        ----------
        world:
            World whose agents should be inspected.

        Returns
        -------
        tuple[Agent, ...]
            Alive agents in deterministic insertion order.
        """
        return tuple(agent for agent in world.agents.values() if agent.alive)


@dataclass(frozen=True, slots=True)
class DeterministicActivationStrategy(AgentActivationStrategy):
    """
    Deterministic insertion-order activation strategy.

    Agents act in the order they were inserted into ``world.agents``. This is
    the default strategy and is ideal for reproducible tests and deterministic
    simulations.
    """

    name: str = "deterministic"

    def order_agents(self, world: World) -> tuple[Agent, ...]:
        """
        Return alive agents in insertion order.

        Parameters
        ----------
        world:
            World whose agents should be ordered.

        Returns
        -------
        tuple[Agent, ...]
            Alive agents in insertion order.
        """
        return self.alive_agents(world)


@dataclass(frozen=True, slots=True)
class RandomActivationStrategy(AgentActivationStrategy):
    """
    Random activation strategy.

    Alive agents are randomly permuted using the world's NumPy random generator.
    This keeps randomness centralized and reproducible when the world has a
    fixed seed.
    """

    name: str = "random"

    def order_agents(self, world: World) -> tuple[Agent, ...]:
        """
        Return alive agents in random order.

        Parameters
        ----------
        world:
            World whose agents should be ordered.

        Returns
        -------
        tuple[Agent, ...]
            Alive agents in randomized order.
        """
        agents = list(self.alive_agents(world))

        if len(agents) <= 1:
            return tuple(agents)

        rng = getattr(world, "rng", None)

        if not isinstance(rng, np.random.Generator):
            logger.warning(
                "World does not expose a NumPy Generator as rng; "
                "falling back to an unseeded generator."
            )
            rng = np.random.default_rng()

        indices = rng.permutation(len(agents))
        return tuple(agents[int(index)] for index in indices)


@dataclass(frozen=True, slots=True)
class PriorityActivationStrategy(AgentActivationStrategy):
    """
    Priority-based activation strategy.

    Agents act according to a numeric priority value stored in their state.

    Examples
    --------
    If ``priority_key="urgency"``, then each agent may define:

        agent.state["urgency"] = 5.0

    Agents with higher urgency act first when ``descending`` is True.
    """

    priority_key: str = "priority"
    descending: bool = True
    default_priority: float = 0.0
    name: str = "priority"

    def __post_init__(self) -> None:
        """
        Validate priority activation configuration.

        Raises
        ------
        ValueError
            If priority_key is invalid.
        TypeError
            If descending or default_priority are invalid.
        """
        if not isinstance(self.priority_key, str) or not self.priority_key.strip():
            raise ValueError("priority_key must be a non-empty string.")

        if not isinstance(self.descending, bool):
            raise TypeError("descending must be a boolean.")

        if isinstance(self.default_priority, bool) or not isinstance(
            self.default_priority,
            Real,
        ):
            raise TypeError("default_priority must be numeric.")

        normalized_default = float(self.default_priority)

        if not np.isfinite(normalized_default):
            raise ValueError("default_priority must be finite.")

        object.__setattr__(self, "priority_key", self.priority_key.strip())
        object.__setattr__(self, "default_priority", normalized_default)

    def order_agents(self, world: World) -> tuple[Agent, ...]:
        """
        Return alive agents ordered by priority.

        Parameters
        ----------
        world:
            World whose agents should be ordered.

        Returns
        -------
        tuple[Agent, ...]
            Alive agents ordered by numeric priority.
        """
        rows: list[tuple[float, int, Agent]] = []

        for index, agent in enumerate(self.alive_agents(world)):
            rows.append((self._priority(agent), index, agent))

        if self.descending:
            rows.sort(key=lambda row: (-row[0], row[1], row[2].id))
        else:
            rows.sort(key=lambda row: (row[0], row[1], row[2].id))

        return tuple(row[2] for row in rows)

    def _priority(self, agent: Agent) -> float:
        """
        Read an agent's numeric priority.

        Parameters
        ----------
        agent:
            Agent whose priority should be read.

        Returns
        -------
        float
            Normalized priority value.
        """
        value = agent.state.get(self.priority_key, self.default_priority)

        if isinstance(value, bool) or not isinstance(value, Real):
            logger.warning(
                "Agent priority is non-numeric; using default priority. "
                "agent_id=%s priority_key=%s",
                agent.id,
                self.priority_key,
            )
            return self.default_priority

        normalized = float(value)

        if not np.isfinite(normalized):
            logger.warning(
                "Agent priority is non-finite; using default priority. "
                "agent_id=%s priority_key=%s",
                agent.id,
                self.priority_key,
            )
            return self.default_priority

        return normalized


@dataclass(frozen=True, slots=True)
class FrozenOrderSimultaneousActivationStrategy(AgentActivationStrategy):
    """
    Frozen-order simultaneous-compatible activation strategy.

    This strategy snapshots the alive-agent activation set before execution.
    That means agents spawned during the current agent phase do not act until a
    later step, and agents removed before their turn are skipped.

    Important
    ---------
    This is not full transactional simultaneous scheduling. True simultaneous
    updates require a future two-phase behavior protocol where behaviors first
    produce intents and the scheduler commits them together.
    """

    name: str = "simultaneous"

    def order_agents(self, world: World) -> tuple[Agent, ...]:
        """
        Return a frozen alive-agent activation set.

        Parameters
        ----------
        world:
            World whose agents should be ordered.

        Returns
        -------
        tuple[Agent, ...]
            Alive agents captured before the agent execution phase.
        """
        return self.alive_agents(world)


def make_activation_strategy(
    mode: ActivationMode | str,
) -> AgentActivationStrategy:
    """
    Create a built-in activation strategy from a mode.

    Parameters
    ----------
    mode:
        Activation mode string or enum.

    Returns
    -------
    AgentActivationStrategy
        Concrete activation strategy.

    Raises
    ------
    ValueError
        If the mode is unknown.
    """
    normalized_mode = ActivationMode(str(mode))

    if normalized_mode is ActivationMode.DETERMINISTIC:
        return DeterministicActivationStrategy()

    if normalized_mode is ActivationMode.RANDOM:
        return RandomActivationStrategy()

    if normalized_mode is ActivationMode.PRIORITY:
        return PriorityActivationStrategy()

    if normalized_mode is ActivationMode.SIMULTANEOUS:
        return FrozenOrderSimultaneousActivationStrategy()

    raise ValueError(f"Unsupported activation mode: {mode}")


@dataclass(slots=True)
class Scheduler:
    """
    Generic simulation scheduler.

    Scheduler coordinates events, agents, resources, and metrics for one world
    step. It is domain-agnostic and configurable through activation strategies
    and phase settings.

    Attributes
    ----------
    activation_strategy:
        Strategy used to order agent activation. May be a concrete strategy,
        an ActivationMode, or a mode string.
    event_phase:
        When due events should execute.
    resource_phase:
        When resource regeneration updates should execute.
    collect_step_metrics:
        Whether to collect world metrics during each scheduler step.
    stop_on_failure:
        Whether later phases should be skipped after a failed phase.
    metadata:
        Optional scheduler metadata included in step results.
    """

    activation_strategy: AgentActivationStrategy | ActivationMode | str = field(
        default_factory=DeterministicActivationStrategy
    )
    event_phase: EventPhase | str = EventPhase.BEFORE_AGENTS
    resource_phase: ResourcePhase | str = ResourcePhase.AFTER_AGENTS
    collect_step_metrics: bool = True
    stop_on_failure: bool = False
    metadata: dict[str, Any] = field(default_factory=dict)

    def __post_init__(self) -> None:
        """
        Validate and normalize scheduler configuration.

        Raises
        ------
        TypeError
            If configuration fields have invalid types.
        ValueError
            If phase or activation mode values are invalid.
        """
        self.activation_strategy = self._normalize_activation_strategy(
            self.activation_strategy
        )
        self.event_phase = EventPhase(str(self.event_phase))
        self.resource_phase = ResourcePhase(str(self.resource_phase))

        if not isinstance(self.collect_step_metrics, bool):
            raise TypeError("collect_step_metrics must be a boolean.")

        if not isinstance(self.stop_on_failure, bool):
            raise TypeError("stop_on_failure must be a boolean.")

        if not isinstance(self.metadata, Mapping):
            raise TypeError("metadata must be a mapping.")

        self.metadata = dict(self.metadata)
        self._validate_string_keys(self.metadata, "metadata")

    @classmethod
    def from_activation_mode(
        cls,
        mode: ActivationMode | str,
        *,
        event_phase: EventPhase | str = EventPhase.BEFORE_AGENTS,
        resource_phase: ResourcePhase | str = ResourcePhase.AFTER_AGENTS,
        collect_step_metrics: bool = True,
        stop_on_failure: bool = False,
        metadata: Mapping[str, Any] | None = None,
    ) -> Scheduler:
        """
        Build a scheduler from a built-in activation mode.

        Parameters
        ----------
        mode:
            Built-in activation mode.
        event_phase:
            Event execution phase.
        resource_phase:
            Resource update phase.
        collect_step_metrics:
            Whether to collect metrics.
        stop_on_failure:
            Whether to stop later phases after a failure.
        metadata:
            Optional scheduler metadata.

        Returns
        -------
        Scheduler
            Configured scheduler instance.
        """
        return cls(
            activation_strategy=make_activation_strategy(mode),
            event_phase=event_phase,
            resource_phase=resource_phase,
            collect_step_metrics=collect_step_metrics,
            stop_on_failure=stop_on_failure,
            metadata=dict(metadata or {}),
        )

    def step(self, world: World) -> WorldStepResult:
        """
        Execute one scheduled world step.

        The scheduler does not increment ``world.step_count``. The owning
        ``World.step()`` method increments the step count after receiving this
        result.

        Parameters
        ----------
        world:
            World to advance.

        Returns
        -------
        WorldStepResult
            Structured result of the scheduled step.
        """
        self._validate_world(world)

        started_step = world.step_count
        event_results: tuple[EventExecutionResult, ...] = tuple()
        agent_results: tuple[AgentStepResult, ...] = tuple()
        resource_results: tuple[ResourceOperationResult, ...] = tuple()
        errors: list[str] = []
        phases_executed: list[str] = []
        activation_order_ids: list[str] = []

        try:
            if self.event_phase is EventPhase.BEFORE_AGENTS:
                event_results = self.process_events(world)
                phases_executed.append("events_before_agents")

                if self._should_stop_after(event_results):
                    phases_executed.append("stopped_after_events")
                    return self._build_step_result(
                        world=world,
                        started_step=started_step,
                        event_results=event_results,
                        agent_results=agent_results,
                        resource_results=resource_results,
                        phases_executed=phases_executed,
                        activation_order_ids=activation_order_ids,
                        errors=errors,
                    )

            if self.resource_phase is ResourcePhase.BEFORE_AGENTS:
                resource_results = self.update_resources(world)
                phases_executed.append("resources_before_agents")

                if self._should_stop_after(resource_results):
                    phases_executed.append("stopped_after_resources")
                    return self._build_step_result(
                        world=world,
                        started_step=started_step,
                        event_results=event_results,
                        agent_results=agent_results,
                        resource_results=resource_results,
                        phases_executed=phases_executed,
                        activation_order_ids=activation_order_ids,
                        errors=errors,
                    )

            ordered_agents = self.order_agents(world)
            activation_order_ids = [agent.id for agent in ordered_agents]
            agent_results = self.execute_agents(
                world,
                ordered_agents=ordered_agents,
            )
            phases_executed.append("agents")

            if self._should_stop_after(agent_results):
                phases_executed.append("stopped_after_agents")
                return self._build_step_result(
                    world=world,
                    started_step=started_step,
                    event_results=event_results,
                    agent_results=agent_results,
                    resource_results=resource_results,
                    phases_executed=phases_executed,
                    activation_order_ids=activation_order_ids,
                    errors=errors,
                )

            if self.event_phase is EventPhase.AFTER_AGENTS:
                event_results = self.process_events(world)
                phases_executed.append("events_after_agents")

                if self._should_stop_after(event_results):
                    phases_executed.append("stopped_after_events")
                    return self._build_step_result(
                        world=world,
                        started_step=started_step,
                        event_results=event_results,
                        agent_results=agent_results,
                        resource_results=resource_results,
                        phases_executed=phases_executed,
                        activation_order_ids=activation_order_ids,
                        errors=errors,
                    )

            if self.resource_phase is ResourcePhase.AFTER_AGENTS:
                resource_results = self.update_resources(world)
                phases_executed.append("resources_after_agents")

        except Exception as exc:
            logger.exception(
                "Scheduler step raised an exception at world step %s.",
                started_step,
            )
            errors.append(repr(exc))

        return self._build_step_result(
            world=world,
            started_step=started_step,
            event_results=event_results,
            agent_results=agent_results,
            resource_results=resource_results,
            phases_executed=phases_executed,
            activation_order_ids=activation_order_ids,
            errors=errors,
        )

    def process_events(self, world: World) -> tuple[EventExecutionResult, ...]:
        """
        Process all events due at the world's current step.

        Parameters
        ----------
        world:
            World whose due events should execute.

        Returns
        -------
        tuple[EventExecutionResult, ...]
            Event execution results.
        """
        self._validate_world(world)
        return world.process_events()

    def order_agents(self, world: World) -> tuple[Agent, ...]:
        """
        Return agents in activation order.

        Parameters
        ----------
        world:
            World whose agents should be ordered.

        Returns
        -------
        tuple[Agent, ...]
            Ordered activation candidates.

        Raises
        ------
        TypeError
            If the activation strategy returns invalid objects.
        """
        self._validate_world(world)

        ordered_agents = self.activation_strategy.order_agents(world)

        if not isinstance(ordered_agents, tuple):
            ordered_agents = tuple(ordered_agents)

        seen_ids: set[str] = set()

        for agent in ordered_agents:
            if not isinstance(agent, Agent):
                raise TypeError(
                    "Activation strategies must return Agent objects only."
                )

            if agent.id in seen_ids:
                raise ValueError(
                    f"Activation strategy returned duplicate agent id: {agent.id}"
                )

            seen_ids.add(agent.id)

        return ordered_agents

    def execute_agents(
        self,
        world: World,
        *,
        ordered_agents: Sequence[Agent] | None = None,
    ) -> tuple[AgentStepResult, ...]:
        """
        Execute one step for selected agents.

        Agents removed before their turn are skipped. Agents marked dead before
        their turn are skipped. This makes activation robust when earlier
        behaviors mutate the world.

        Parameters
        ----------
        world:
            World in which agents execute.
        ordered_agents:
            Optional precomputed activation order.

        Returns
        -------
        tuple[AgentStepResult, ...]
            Per-agent step results.
        """
        self._validate_world(world)

        agents = tuple(ordered_agents) if ordered_agents is not None else self.order_agents(world)
        results: list[AgentStepResult] = []

        for scheduled_agent in agents:
            if not isinstance(scheduled_agent, Agent):
                raise TypeError("ordered_agents must contain Agent objects.")

            current_agent = world.agents.get(scheduled_agent.id)

            if current_agent is None:
                logger.debug(
                    "Skipping scheduled agent because it no longer exists: agent_id=%s",
                    scheduled_agent.id,
                )
                continue

            if not current_agent.alive:
                logger.debug(
                    "Skipping scheduled agent because it is not alive: agent_id=%s",
                    current_agent.id,
                )
                continue

            try:
                result = current_agent.step(world)
            except Exception as exc:
                logger.exception(
                    "Agent execution raised an exception: agent_id=%s",
                    current_agent.id,
                )
                result = AgentStepResult(
                    agent_id=current_agent.id,
                    agent_type=current_agent.type,
                    step_count=world.step_count,
                    behavior_id=None,
                    success=False,
                    reward=0.0,
                    message="Agent execution raised an exception.",
                    metadata={"error": repr(exc)},
                )

            results.append(result)

            if self.stop_on_failure and not result.success:
                logger.info(
                    "Stopping agent phase after failed agent step: agent_id=%s",
                    result.agent_id,
                )
                break

        return tuple(results)

    def update_resources(self, world: World) -> tuple[ResourceOperationResult, ...]:
        """
        Update world resources.

        The default implementation delegates to ``World.update_resources()``,
        which applies generic regeneration based on each resource's metadata.

        Parameters
        ----------
        world:
            World whose resources should update.

        Returns
        -------
        tuple[ResourceOperationResult, ...]
            Resource operation results.
        """
        self._validate_world(world)
        return world.update_resources()

    def _build_step_result(
        self,
        *,
        world: World,
        started_step: int,
        event_results: tuple[EventExecutionResult, ...],
        agent_results: tuple[AgentStepResult, ...],
        resource_results: tuple[ResourceOperationResult, ...],
        phases_executed: list[str],
        activation_order_ids: list[str],
        errors: list[str],
    ) -> WorldStepResult:
        """
        Build a WorldStepResult from scheduler phase outputs.

        Parameters
        ----------
        world:
            World being scheduled.
        started_step:
            Step count at the beginning of the step.
        event_results:
            Event phase results.
        agent_results:
            Agent phase results.
        resource_results:
            Resource phase results.
        phases_executed:
            Names of phases executed by the scheduler.
        activation_order_ids:
            Agent ids selected for activation.
        errors:
            Scheduler-level errors.

        Returns
        -------
        WorldStepResult
            Structured world step result.
        """
        metrics = world.collect_metrics() if self.collect_step_metrics else {}

        success = (
            not errors
            and self._all_successful(event_results)
            and self._all_successful(agent_results)
            and self._all_successful(resource_results)
        )

        metadata = {
            "scheduler": self.__class__.__name__,
            "activation_strategy": self.activation_strategy.name,
            "event_phase": str(self.event_phase),
            "resource_phase": str(self.resource_phase),
            "phases_executed": list(phases_executed),
            "activation_order_ids": list(activation_order_ids),
            "stop_on_failure": self.stop_on_failure,
        }
        metadata.update(self.metadata)

        return WorldStepResult(
            started_step=started_step,
            ended_step=started_step + 1,
            success=success,
            agent_results=agent_results,
            event_results=event_results,
            resource_results=resource_results,
            metrics=metrics,
            message=(
                "Scheduled world step completed successfully."
                if success
                else "Scheduled world step completed with failures."
            ),
            metadata=metadata,
            errors=tuple(errors),
        )

    def _should_stop_after(self, results: Iterable[Any]) -> bool:
        """
        Return whether scheduling should stop after a phase.

        Parameters
        ----------
        results:
            Phase result objects with optional ``success`` attributes.

        Returns
        -------
        bool
            True if stop_on_failure is enabled and any result failed.
        """
        return self.stop_on_failure and not self._all_successful(results)

    @staticmethod
    def _all_successful(results: Iterable[Any]) -> bool:
        """
        Return whether all result objects are successful.

        Parameters
        ----------
        results:
            Iterable of result objects exposing ``success``.

        Returns
        -------
        bool
            True if all results are successful or if the iterable is empty.
        """
        for result in results:
            success = getattr(result, "success", True)
            if not bool(success):
                return False

        return True

    @staticmethod
    def _normalize_activation_strategy(
        value: AgentActivationStrategy | ActivationMode | str,
    ) -> AgentActivationStrategy:
        """
        Normalize an activation strategy configuration.

        Parameters
        ----------
        value:
            Concrete strategy, ActivationMode, or mode string.

        Returns
        -------
        AgentActivationStrategy
            Concrete activation strategy.

        Raises
        ------
        TypeError
            If the value is not a valid strategy or mode.
        """
        if isinstance(value, AgentActivationStrategy):
            return value

        if isinstance(value, (str, ActivationMode)):
            return make_activation_strategy(value)

        raise TypeError(
            "activation_strategy must be an AgentActivationStrategy, "
            "ActivationMode, or mode string."
        )

    @staticmethod
    def _validate_world(world: World) -> None:
        """
        Validate that a value is a World.

        Parameters
        ----------
        world:
            Candidate world.

        Raises
        ------
        TypeError
            If world is not a World.
        """
        if not isinstance(world, World):
            raise TypeError("world must be a World instance.")

    @staticmethod
    def _validate_string_keys(mapping: Mapping[str, Any], label: str) -> None:
        """
        Validate that all mapping keys are strings.

        Parameters
        ----------
        mapping:
            Mapping to validate.
        label:
            Human-readable mapping label.

        Raises
        ------
        TypeError
            If any key is not a string.
        """
        for key in mapping:
            if not isinstance(key, str):
                raise TypeError(f"Scheduler.{label} keys must be strings.")


__all__ = [
    "ActivationMode",
    "AgentActivationStrategy",
    "DeterministicActivationStrategy",
    "EventPhase",
    "FrozenOrderSimultaneousActivationStrategy",
    "PriorityActivationStrategy",
    "RandomActivationStrategy",
    "ResourcePhase",
    "Scheduler",
    "make_activation_strategy",
]