engine/controller.py

from __future__ import annotations

from dataclasses import dataclass, field
from enum import Enum
from typing import Mapping

import numpy as np

from engine.bocpd import run_bocpd
from engine.brain import BrainSignals, coerce_signals


class Action(str, Enum):
    SILENT = "SILENT"
    BACKCHANNEL = "BACKCHANNEL"
    TAKE_FLOOR = "TAKE_FLOOR"
    INTERRUPT = "INTERRUPT"


@dataclass(frozen=True)
class ControllerConfig:
    w_surprise: float = 0.70
    w_change: float = 1.30
    w_readiness: float = 0.80
    w_end: float = 1.20
    w_barge: float = 0.60
    negative_surprise_weight: float = 0.25
    tau: float = 1.60
    backchannel_tau_fraction: float = 0.70
    barge_tau_fraction: float = 0.50
    take_floor_p_end: float = 0.70
    interrupt_p_end_max: float = 0.35
    backchannel_p_end_max: float = 0.35
    min_readiness: float = 0.45
    refractory_steps: int = 2
    surprise_z_cap: float = 3.0
    change_hazard: float = 0.35
    change_prior_kappa: float = 0.20
    change_prior_alpha: float = 0.75
    change_prior_beta: float = 0.20
    change_z_cap: float = 3.0
    change_z_threshold: float = 1.15
    turn_end_tau_discount: float = 0.35

    @property
    def backchannel_tau(self) -> float:
        return self.backchannel_tau_fraction * self.tau

    @property
    def barge_tau(self) -> float:
        return self.barge_tau_fraction * self.tau


@dataclass
class RunningStats:
    n: int = 0
    mean: float = 0.0
    m2: float = 0.0

    @property
    def std(self) -> float:
        if self.n < 2:
            return 1.0
        return max((self.m2 / (self.n - 1)) ** 0.5, 1.0e-6)

    def zscore(self, value: float, cap: float) -> float:
        if self.n < 2:
            return 0.0
        z_value = (float(value) - self.mean) / self.std
        return float(np.clip(z_value, -cap, cap))

    def update(self, value: float) -> None:
        self.n += 1
        delta = float(value) - self.mean
        self.mean += delta / self.n
        self.m2 += delta * (float(value) - self.mean)


@dataclass
class AgentState:
    surprise_stats: RunningStats = field(default_factory=RunningStats)
    previous_hidden: np.ndarray | None = None
    hidden_deltas: list[float] = field(default_factory=list)
    hidden_delta_z: list[float] = field(default_factory=list)
    hidden_delta_stats: RunningStats = field(default_factory=RunningStats)
    previous_map_run_length: int | None = None
    refractory_until: int = 0


@dataclass(frozen=True)
class AgentDecision:
    agent_id: str
    action: Action
    urge: float
    z_surprise: float
    change_score: float
    readiness: float
    p_end: float
    hidden_delta: float
    map_run_length: int
    refractory: bool


@dataclass(frozen=True)
class ControllerTick:
    step: int
    floor_holder: str
    winner: str | None
    decisions: dict[str, AgentDecision]


class WhenToSpeakController:
    """Training-free multi-agent timing controller."""

    def __init__(self, agent_ids: list[str], config: ControllerConfig | None = None) -> None:
        if not agent_ids:
            raise ValueError("agent_ids must not be empty")
        self.agent_ids = list(agent_ids)
        self.config = config or ControllerConfig()
        self.states = {agent_id: AgentState() for agent_id in self.agent_ids}
        self.step = 0
        self.floor_holder = "human"

    def reset(self) -> None:
        self.states = {agent_id: AgentState() for agent_id in self.agent_ids}
        self.step = 0
        self.floor_holder = "human"

    def tick(
        self,
        signals_by_agent: Mapping[str, BrainSignals | dict[str, object]],
        *,
        floor_holder: str | None = None,
    ) -> ControllerTick:
        self.step += 1
        if floor_holder is not None:
            self.floor_holder = floor_holder

        scored: dict[str, AgentDecision] = {}
        proposed: dict[str, Action] = {}
        for agent_id in self.agent_ids:
            if agent_id not in signals_by_agent:
                raise KeyError(f"missing signals for agent {agent_id!r}")
            signal = coerce_signals(signals_by_agent[agent_id])
            decision = self._score_agent(agent_id, signal)
            scored[agent_id] = decision
            proposed[agent_id] = decision.action

        winner = self._floor_winner(scored)
        final_decisions: dict[str, AgentDecision] = {}
        for agent_id, decision in scored.items():
            action = decision.action
            if action in {Action.TAKE_FLOOR, Action.INTERRUPT} and agent_id != winner:
                action = Action.BACKCHANNEL if self._may_backchannel(decision) else Action.SILENT
            final_decisions[agent_id] = AgentDecision(
                agent_id=decision.agent_id,
                action=action,
                urge=decision.urge,
                z_surprise=decision.z_surprise,
                change_score=decision.change_score,
                readiness=decision.readiness,
                p_end=decision.p_end,
                hidden_delta=decision.hidden_delta,
                map_run_length=decision.map_run_length,
                refractory=decision.refractory,
            )

        for agent_id, decision in final_decisions.items():
            if decision.action in {Action.TAKE_FLOOR, Action.INTERRUPT}:
                self.states[agent_id].refractory_until = self.step + self.config.refractory_steps

        if winner is not None:
            self.floor_holder = winner

        return ControllerTick(
            step=self.step,
            floor_holder=self.floor_holder,
            winner=winner,
            decisions=final_decisions,
        )

    def _score_agent(self, agent_id: str, signal: BrainSignals) -> AgentDecision:
        state = self.states[agent_id]
        z_surprise = state.surprise_stats.zscore(signal.surprise, self.config.surprise_z_cap)
        hidden_delta, change_score, map_run_length = self._change_features(state, signal.hidden)
        barge = self.config.w_barge * max(z_surprise, 0.0) * signal.readiness * (1.0 - signal.p_end)
        surprise_term = z_surprise if z_surprise >= 0.0 else self.config.negative_surprise_weight * z_surprise
        urge = (
            self.config.w_surprise * surprise_term
            + self.config.w_change * change_score
            + self.config.w_readiness * signal.readiness
            + self.config.w_end * signal.p_end
            + barge
        )

        refractory = self.step <= state.refractory_until
        action = self._classify(urge, z_surprise, change_score, signal, refractory)

        state.surprise_stats.update(signal.surprise)
        state.previous_hidden = signal.hidden.astype(np.float32, copy=True)
        return AgentDecision(
            agent_id=agent_id,
            action=action,
            urge=float(urge),
            z_surprise=float(z_surprise),
            change_score=float(change_score),
            readiness=signal.readiness,
            p_end=signal.p_end,
            hidden_delta=float(hidden_delta),
            map_run_length=int(map_run_length),
            refractory=refractory,
        )

    def _change_features(self, state: AgentState, hidden: np.ndarray) -> tuple[float, float, int]:
        if state.previous_hidden is None:
            return 0.0, 0.0, 0

        hidden_delta = cosine_distance(state.previous_hidden, hidden)
        delta_z = state.hidden_delta_stats.zscore(hidden_delta, self.config.change_z_cap)
        state.hidden_delta_stats.update(hidden_delta)
        state.hidden_deltas.append(hidden_delta)
        state.hidden_delta_z.append(delta_z)
        results = run_bocpd(
            np.asarray(state.hidden_delta_z, dtype=np.float64),
            hazard=self.config.change_hazard,
            prior_kappa=self.config.change_prior_kappa,
            prior_alpha=self.config.change_prior_alpha,
            prior_beta=self.config.change_prior_beta,
        )
        latest = results[-1]
        previous_map = state.previous_map_run_length
        state.previous_map_run_length = latest.map_run_length
        if previous_map is None:
            return hidden_delta, 0.0, latest.map_run_length

        collapsed = latest.map_run_length < previous_map
        collapse_ratio = (previous_map - latest.map_run_length) / max(previous_map, 1)
        collapse_score = max(1.0, collapse_ratio) if collapsed else 0.0
        posterior_score = max(0.0, (latest.cp_prob - self.config.change_hazard) / max(1.0 - self.config.change_hazard, 1.0e-9))
        z_score = max(0.0, abs(delta_z) - self.config.change_z_threshold) / max(
            self.config.change_z_cap - self.config.change_z_threshold,
            1.0e-9,
        )
        change_score = max(collapse_score, posterior_score, z_score)
        return hidden_delta, float(change_score), latest.map_run_length

    def _classify(
        self,
        urge: float,
        z_surprise: float,
        change_score: float,
        signal: BrainSignals,
        refractory: bool,
    ) -> Action:
        if refractory:
            return Action.SILENT

        ready = signal.readiness >= self.config.min_readiness
        human_has_floor = self.floor_holder == "human"
        barge_signal = max(z_surprise, 0.0) * signal.readiness * (1.0 - signal.p_end)
        floor_tau = self._effective_tau(signal.p_end)

        if human_has_floor and ready and signal.p_end >= self.config.take_floor_p_end and urge >= floor_tau:
            return Action.TAKE_FLOOR
        if (
            human_has_floor
            and ready
            and signal.p_end <= self.config.interrupt_p_end_max
            and urge >= floor_tau
            and barge_signal >= self.config.barge_tau
        ):
            return Action.INTERRUPT
        if (
            human_has_floor
            and ready
            and signal.p_end <= self.config.backchannel_p_end_max
            and change_score > 0.0
            and urge >= self.config.backchannel_tau
        ):
            return Action.BACKCHANNEL
        if (
            human_has_floor
            and ready
            and urge >= self.config.backchannel_tau
            and signal.p_end <= self.config.backchannel_p_end_max
        ):
            return Action.BACKCHANNEL
        return Action.SILENT

    def _floor_winner(self, decisions: Mapping[str, AgentDecision]) -> str | None:
        contenders = [
            decision
            for decision in decisions.values()
            if decision.action in {Action.TAKE_FLOOR, Action.INTERRUPT}
            and decision.urge >= self._effective_tau(decision.p_end)
        ]
        if not contenders:
            return None
        return max(contenders, key=lambda decision: (decision.urge, -self.agent_ids.index(decision.agent_id))).agent_id

    def _may_backchannel(self, decision: AgentDecision) -> bool:
        return (
            not decision.refractory
            and self.floor_holder == "human"
            and decision.urge >= self.config.backchannel_tau
            and decision.p_end <= self.config.backchannel_p_end_max
        )

    def _effective_tau(self, p_end: float) -> float:
        discount = self.config.turn_end_tau_discount * float(np.clip(p_end, 0.0, 1.0))
        return self.config.tau * max(0.20, 1.0 - discount)


def cosine_distance(left: np.ndarray, right: np.ndarray) -> float:
    left_vec = np.asarray(left, dtype=np.float32)
    right_vec = np.asarray(right, dtype=np.float32)
    denom = float(np.linalg.norm(left_vec) * np.linalg.norm(right_vec))
    if denom <= 1.0e-12:
        return 0.0
    similarity = float(np.dot(left_vec, right_vec) / denom)
    return float(np.clip(1.0 - similarity, 0.0, 2.0))