environment/consequences.py

from __future__ import annotations

from dataclasses import replace
import random
from typing import Dict

from .state import EpisodeState, RelationshipState, clamp01


def _copy_state(state: EpisodeState) -> EpisodeState:
    return replace(
        state,
        relationships=replace(state.relationships),
        decisions=[*state.decisions],
        history=[*state.history],
    )


def apply_action(state: EpisodeState, event_id: str, action_id: str, action_data: Dict) -> EpisodeState:
    new_state = _copy_state(state)

    new_state.energy = clamp01(state.energy + float(action_data.get("energy_delta", 0.0)))
    new_state.sprint_health = clamp01(state.sprint_health + float(action_data.get("sprint_delta", 0.0)))

    rel = new_state.relationships
    for person, delta in action_data.get("relationship_delta", {}).items():
        if hasattr(rel, person):
            setattr(rel, person, clamp01(getattr(rel, person) + float(delta)))
    new_state.relationships = RelationshipState(**vars(rel)).clamped()

    if event_id == "E1_staging" and action_id == "fix_directly":
        new_state.staging_fixed = True
    if event_id == "E4_leave":
        if action_id == "ping_sundar":
            new_state.leave_status = "approved" if random.random() > 0.20 else "pending"
        elif action_id == "cancel_leave":
            new_state.leave_status = "cancelled"
    if event_id == "E5_appraisal" and action_id == "block_wednesday":
        new_state.appraisal_done = True
    if event_id == "E6_oncall" and action_id == "yes_people_pleaser":
        new_state.oncall_accepted = True

    new_state.decisions.append(
        {
            "day": state.day,
            "event_id": event_id,
            "action_id": action_id,
            "energy_before": state.energy,
            "energy_after": new_state.energy,
        }
    )

    new_state.history.append(f"{event_id}:{action_id}")
    return new_state.normalized()