env/environment.py

"""
SYNAPSE-X Environment

OpenEnv-style environment for predictive task scheduling under risk,
uncertainty, deadlines, and resource constraints.
"""

from copy import deepcopy
from typing import Any, Optional

from env.echo import ECHO
from env.models import Action, Observation, StepResult, Task
from env.prism import PRISM
from env.reward import RewardEngine


class SynapseXEnvironment:
    """
    SYNAPSE-X predictive decision intelligence environment.
    """

    MAX_STEPS = 30
    MAX_RESOURCES = 1.0
    MAX_DEADLINE = 20.0
    REWARD_MIN = -2.0
    REWARD_MAX = 2.0
    COMPLETION_BONUS = 2.0
    CASCADE_DELAY_ALPHA = 0.03
    CASCADE_DEPENDENCY_BETA = 0.12
    CASCADE_PRESSURE_GAMMA = 0.05
    CASCADE_HIDDEN_DELAY = 0.06
    CASCADE_MEMORY_DECAY = 0.85
    CASCADE_MEMORY_FEEDBACK = 0.05
    CASCADE_PHASE_THRESHOLD = 0.75
    CASCADE_PHASE_MULTIPLIER = 1.2
    CASCADE_PRESSURE_DEADLINE_ACCEL = 0.25
    CASCADE_RESOURCE_CONTENTION = 0.12

    def __init__(self, task_config: Optional[list[dict[str, Any]]] = None, seed: int = 42):
        self.seed = seed
        self.task_config = task_config or self._default_task_config()
        self.echo = ECHO(max_deadline=self.MAX_DEADLINE)
        self.prism = PRISM(seed=seed)
        self.reward_engine = RewardEngine()

        self._tasks: list[Task] = []
        self._time = 0
        self._resources = self.MAX_RESOURCES
        self._total_reward = 0.0
        self._history: list[dict[str, Any]] = []
        self._done = False
        self._dependency_graph = {
            cfg["id"]: list(cfg.get("dependencies", []))
            for cfg in self.task_config
        }
        self._hidden_penalties = {cfg["id"]: 0.0 for cfg in self.task_config}
        self._cascade_mode = any(self._dependency_graph.values())

    def reset(self) -> Observation:
        self.prism = PRISM(seed=self.seed)
        self._time = 0
        self._resources = self.MAX_RESOURCES
        self._total_reward = 0.0
        self._history = []
        self._done = False
        self._hidden_penalties = {cfg["id"]: 0.0 for cfg in self.task_config}

        self._tasks = []
        for cfg in sorted(self.task_config, key=lambda item: item["id"]):
            task = Task(**cfg)
            task.released = task.release_time <= self._time
            self.echo.predict(task, self._time)
            self._tasks.append(task)

        return self._build_observation()

    def step(self, action: Action) -> StepResult:
        if self._done:
            raise RuntimeError("Episode is done. Call reset() first.")

        action_type = getattr(action, "action_type", None)
        task_id = getattr(action, "task_id", None)
        if action_type not in {"execute", "delay", "reallocate"}:
            return StepResult(
                observation=self._build_observation(),
                reward=-0.2,
                done=self._done,
                info={
                    "error": "invalid_action",
                    "task_id": task_id,
                    "action": action_type,
                },
            )

        task = self._get_task(task_id)
        if task is None:
            return StepResult(
                observation=self._build_observation(),
                reward=-0.2,
                done=self._done,
                info={
                    "error": "invalid_task",
                    "task_id": task_id,
                    "action": action_type,
                },
            )
        if not task.released:
            return StepResult(
                observation=self._build_observation(),
                reward=-0.15,
                done=self._done,
                info={
                    "error": "task_unavailable",
                    "task_id": task_id,
                    "action": action_type,
                    "release_time": task.release_time,
                },
            )

        reward = 0.0
        info: dict[str, Any] = {"action": action_type, "task_id": task_id}

        if action_type == "execute":
            if task.is_terminal:
                reward = -0.1
                info["result"] = "already_terminal"
            elif self._resources < task.resources_required:
                reward = -0.3
                info["result"] = "insufficient_resources"
            else:
                if self._cascade_mode:
                    self._apply_execution_pressure(task)
                success = self.prism.execution_succeeds(task)
                self._resources = max(0.0, self._resources - task.resources_required)

                if success:
                    task.completed = True
                    info["result"] = "success"
                else:
                    task.failed = True
                    info["result"] = "failure"

                reward = self.reward_engine.compute(
                    task=task,
                    success=success,
                    current_time=self._time,
                    resources=self._resources,
                    action_type="execute",
                )

        elif action_type == "delay":
            if task.is_terminal:
                reward = -0.1
                info["result"] = "already_terminal"
            else:
                task.deadline = max(0.0, task.deadline - 1.0)
                task.delay_count += 1
                reward = self.reward_engine.compute(
                    task=task,
                    success=False,
                    current_time=self._time,
                    resources=self._resources,
                    action_type="delay",
                )
                info["result"] = "delayed"

        elif action_type == "reallocate":
            boost = min(0.2, self.MAX_RESOURCES - self._resources)
            self._resources = min(self.MAX_RESOURCES, self._resources + boost)

            if boost <= 0.0:
                reward = -0.05
                info["result"] = "resources_full"
            else:
                reward = self.reward_engine.compute(
                    task=task,
                    success=False,
                    current_time=self._time,
                    resources=self._resources,
                    action_type="reallocate",
                )
                info["result"] = "reallocated"

            info["resource_boost"] = round(boost, 4)

        self._time += 1

        self._release_available_tasks()
        system_pressure = self._apply_cascade_dynamics(trigger_task=task, action_type=action_type)

        for current_task in self._tasks:
            if current_task.is_active:
                deadline_decay = 1.0
                if self._cascade_mode:
                    deadline_decay += self.CASCADE_PRESSURE_DEADLINE_ACCEL * system_pressure
                current_task.deadline = max(0.0, current_task.deadline - deadline_decay)
                if current_task.deadline <= 0.0:
                    current_task.failed = True

        for current_task in self._tasks:
            self.echo.predict(current_task, self._time)

        all_terminal = all(current_task.is_terminal for current_task in self._tasks)
        all_completed = all(current_task.completed for current_task in self._tasks)
        time_up = self._time >= self.MAX_STEPS
        self._done = all_terminal or time_up

        if self._cascade_mode:
            reward += self._cascade_reward_adjustment(task, action_type, system_pressure)

        if all_completed:
            reward += self.COMPLETION_BONUS
            info["terminal_bonus"] = self.COMPLETION_BONUS

        reward = self._clamp_reward(reward)
        self._total_reward += reward

        if self._done:
            info["done_reason"] = "all_tasks_terminal" if all_terminal else "time_limit"
            info["all_tasks_completed"] = all_completed

        observation = self._build_observation(done=self._done)
        self._history.append({"step": self._time, "reward": reward, "info": info})
        return StepResult(observation=observation, reward=reward, done=self._done, info=info)

    def state(self) -> dict[str, Any]:
        return {
            "time": self._time,
            "resources": round(self._resources, 4),
            "total_reward": round(self._total_reward, 4),
            "done": self._done,
            "tasks": [task.model_dump() for task in self._sorted_tasks()],
            "history": self._history,
            "seed": self.seed,
        }

    def _build_observation(self, done: bool = False) -> Observation:
        return Observation(
            tasks=deepcopy(self._sorted_tasks()),
            time=self._time,
            resources=round(self._resources, 4),
            episode_done=done,
        )

    def _sorted_tasks(self) -> list[Task]:
        return sorted(self._tasks, key=lambda task: task.id)

    def _clamp_reward(self, reward: float) -> float:
        return round(max(self.REWARD_MIN, min(self.REWARD_MAX, reward)), 4)

    def _get_task(self, task_id: int) -> Optional[Task]:
        for task in self._tasks:
            if task.id == task_id:
                return task
        return None

    def _release_available_tasks(self) -> None:
        for task in self._tasks:
            if not task.released and task.release_time <= self._time:
                task.released = True

    def _apply_execution_pressure(self, task: Task) -> None:
        pressure = self._compute_system_pressure()
        dependency_penalty = self._dependency_penalty(task)
        hidden_penalty = self._hidden_penalties.get(task.id, 0.0)
        contention_penalty = max(0.0, task.resources_required - self._resources) * self.CASCADE_RESOURCE_CONTENTION
        pressure_penalty = self.CASCADE_PRESSURE_GAMMA * (pressure**2)
        task.risk = min(
            1.0,
            task.risk + dependency_penalty + hidden_penalty + pressure_penalty + contention_penalty,
        )

    def _apply_cascade_dynamics(self, trigger_task: Optional[Task], action_type: str) -> float:
        if not self._cascade_mode:
            return 0.0

        if action_type == "delay" and trigger_task is not None and trigger_task.is_active:
            self._hidden_penalties[trigger_task.id] += self.CASCADE_HIDDEN_DELAY

        active_tasks = [task for task in self._tasks if task.is_active]
        if not active_tasks:
            return 0.0

        system_pressure = self._compute_system_pressure(active_tasks)
        pressure_field = system_pressure**2

        for task in active_tasks:
            delay_penalty = self.CASCADE_DELAY_ALPHA * task.delay_count
            dependency_penalty = self._dependency_penalty(task)
            hidden_penalty = self._hidden_penalties.get(task.id, 0.0) * system_pressure
            task.risk = min(
                1.0,
                task.risk + delay_penalty + dependency_penalty + self.CASCADE_PRESSURE_GAMMA * pressure_field + hidden_penalty,
            )
            task._history_risk = self.CASCADE_MEMORY_DECAY * task._history_risk + task.risk
            task.risk = min(1.0, task.risk + self.CASCADE_MEMORY_FEEDBACK * task._history_risk)

        system_pressure = self._compute_system_pressure(active_tasks)
        if system_pressure > self.CASCADE_PHASE_THRESHOLD:
            for task in active_tasks:
                task.uncertainty = min(1.0, task.uncertainty * self.CASCADE_PHASE_MULTIPLIER)

        return system_pressure

    def _compute_system_pressure(self, tasks: Optional[list[Task]] = None) -> float:
        active_tasks = tasks or [task for task in self._tasks if task.is_active]
        if not active_tasks:
            return 0.0
        weighted_pressure = sum(task.risk * (1.0 + task.deadline_pressure) for task in active_tasks)
        raw_pressure = weighted_pressure / len(active_tasks)
        return raw_pressure / (1.0 + raw_pressure)

    def _dependency_penalty(self, task: Task) -> float:
        if not self._cascade_mode:
            return 0.0
        parents = self._dependency_graph.get(task.id, [])
        if not parents:
            return 0.0

        penalty = 0.0
        for parent_id in parents:
            parent = self._get_task(parent_id)
            if parent is None:
                continue
            penalty += 0.05 * parent.delay_count
            if parent.failed:
                penalty += self.CASCADE_DEPENDENCY_BETA
        return penalty

    def _cascade_reward_adjustment(self, task: Task, action_type: str, system_pressure: float) -> float:
        failed_dependencies = 0
        for parent_id in self._dependency_graph.get(task.id, []):
            parent = self._get_task(parent_id)
            if parent is not None and parent.failed:
                failed_dependencies += 1

        cascade_penalty = 0.5 * failed_dependencies
        pressure_penalty = 0.1 * system_pressure
        hidden_delay_penalty = 0.05 * self._hidden_penalties.get(task.id, 0.0)

        adjustment = -(cascade_penalty + pressure_penalty + hidden_delay_penalty)
        if action_type == "delay":
            adjustment -= 0.1 * task.delay_count
        return adjustment

    def _default_task_config(self) -> list[dict[str, Any]]:
        return [
            {
                "id": 0,
                "name": "Low-Risk Report Generation",
                "priority": 0.5,
                "risk": 0.1,
                "uncertainty": 0.1,
                "deadline": 15.0,
                "resources_required": 0.1,
            }
        ]

    @property
    def total_reward(self) -> float:
        return self._total_reward

    @property
    def time(self) -> int:
        return self._time

    @property
    def done(self) -> bool:
        return self._done