src/environment.py

"""
FusionOps Environment
OpenEnv-compatible RL environment for ML graph scheduling.
"""

from __future__ import annotations

import json
import re
from dataclasses import dataclass, field
from typing import Optional

from .models import (
    Action, Config, Graph, ScheduleState, SubgraphEntry,
)
from .cost_model import compute_subgraph_latency, compute_naive_latency
from .validator import validate_action
from .observation import format_observation


@dataclass
class StepResult:
    observation: str
    reward: float
    done: bool
    info: dict = field(default_factory=dict)


class FusionOpsEnv:
    """
    RL environment for ML graph scheduling.
    The agent builds an execution schedule step by step.
    """

    def __init__(
        self,
        graph: Graph,
        max_steps: int = 20,
        allow_recomputation: bool = True,
    ):
        self.graph = graph
        self.max_steps = max_steps
        self.allow_recomputation = allow_recomputation

        # Precompute naive baseline for grading
        self.naive_latency = compute_naive_latency(graph)

        self.state: Optional[ScheduleState] = None
        self._last_action_result = ""

    def reset(self) -> StepResult:
        """Initialize the environment. Returns initial observation."""
        self.state = ScheduleState()
        self._last_action_result = ""

        obs = format_observation(
            graph=self.graph,
            state=self.state,
            last_action_result="",
            last_action_error=None,
            max_steps=self.max_steps,
            naive_latency=self.naive_latency,
        )

        return StepResult(
            observation=obs,
            reward=0.0,
            done=False,
            info={"naive_latency": self.naive_latency},
        )

    def _classify_error(self, error_msg: str) -> dict:
        """Classify an error message into type + fix hint for the LLM."""
        msg = (error_msg or "").lower()
        ready_ops = self._compute_ready_ops()

        if "oom" in msg or "working set" in msg or "exceeds" in msg:
            return {
                "type": "Memory Error (working set too large)",
                "reason": error_msg,
                "fix_hint": "Reduce tile size (e.g., config=[64,64,1]) or use split-K (smaller k for MatMul). Working set must fit in fast memory.",
            }
        if "dependency" in msg or "needs tensor" in msg or "not scheduled" in msg:
            return {
                "type": "Dependency Error (op not ready)",
                "reason": error_msg,
                "fix_hint": f"Choose ops only from READY OPS: {ready_ops}",
            }
        if "retain" in msg and "not produced" in msg:
            return {
                "type": "Retention Error (tensor not produced by current subgraph)",
                "reason": error_msg,
                "fix_hint": "You can only retain tensors that the CURRENT subgraph produces (its outputs).",
            }
        if "connected" in msg or "subgraph" in msg:
            return {
                "type": "Connectivity Error (ops do not form connected subgraph)",
                "reason": error_msg,
                "fix_hint": "Ops in a subgraph must form a connected DAG (one op produces a tensor consumed by another in the same subgraph).",
            }
        if "parse" in msg or "format" in msg:
            return {
                "type": "Parse Error (invalid action format)",
                "reason": error_msg,
                "fix_hint": "Use exact format: SCHEDULE ops=[0,1] config=[128,128,1] retain=[]",
            }
        if "already scheduled" in msg or "duplicate" in msg:
            return {
                "type": "Already Scheduled (op already executed)",
                "reason": error_msg,
                "fix_hint": f"This op was already executed. Choose from READY OPS: {ready_ops}",
            }
        # Default
        return {
            "type": "Invalid Action",
            "reason": error_msg,
            "fix_hint": f"Choose ops from READY OPS: {ready_ops}. Use format: SCHEDULE ops=[op_id] config=[128,128,1] retain=[]",
        }

    def _compute_ready_ops(self) -> list:
        """Find ops whose predecessors are all scheduled."""
        if self.state is None:
            return []
        ready = []
        for op in self.graph.operations:
            if op.id in self.state.scheduled_op_ids:
                continue
            preds = self.graph.op_predecessors.get(op.id, set())
            if preds.issubset(self.state.scheduled_op_ids):
                ready.append(op.id)
        return ready

    def _penalty_for_error(self, error_type: str) -> float:
        """Graduated penalties based on error type."""
        # Match by prefix to handle the descriptive tags
        if error_type.startswith("Parse Error"):
            return -0.10
        if error_type.startswith("Memory Error"):
            return -0.05  # Close to valid, just wrong size
        if error_type.startswith("Dependency Error"):
            return -0.20  # Clearly wrong order
        if error_type.startswith("Retention Error"):
            return -0.15
        if error_type.startswith("Connectivity Error"):
            return -0.15
        if error_type.startswith("Already Scheduled"):
            return -0.10
        return -0.10

    def step(self, action: Action) -> StepResult:
        """Execute one scheduling step."""
        assert self.state is not None, "Must call reset() before step()"

        # Handle recomputation: temporarily allow re-scheduling of ops
        state_for_validation = self.state
        if self.allow_recomputation:
            state_for_validation = self.state.clone()
            for oid in action.operation_ids:
                state_for_validation.scheduled_op_ids.discard(oid)

        # Validate
        validation = validate_action(self.graph, action, state_for_validation)
        if not validation.is_valid:
            self.state.step_count += 1
            error_info = self._classify_error(validation.error)
            penalty = self._penalty_for_error(error_info["type"])

            done = self.state.step_count >= self.max_steps
            obs = format_observation(
                graph=self.graph,
                state=self.state,
                last_action_result="",
                last_action_error=error_info,
                max_steps=self.max_steps,
                naive_latency=self.naive_latency,
            )

            return StepResult(
                observation=obs,
                reward=penalty,
                done=done,
                info={"error": validation.error, "error_type": error_info["type"]},
            )

        # Compute latency
        result = compute_subgraph_latency(
            self.graph, action, state_for_validation
        )

        if not result.is_valid:
            self.state.step_count += 1
            error_info = self._classify_error(result.error)
            penalty = self._penalty_for_error(error_info["type"])

            done = self.state.step_count >= self.max_steps
            obs = format_observation(
                graph=self.graph,
                state=self.state,
                last_action_result="",
                last_action_error=error_info,
                max_steps=self.max_steps,
                naive_latency=self.naive_latency,
            )

            return StepResult(
                observation=obs,
                reward=penalty,
                done=done,
                info={"error": result.error, "error_type": error_info["type"]},
            )

        # Update state
        self._update_state(action, result.total_latency)

        # Check if done
        all_ops_covered = self._all_ops_covered()
        step_limit = self.state.step_count >= self.max_steps
        done = all_ops_covered or step_limit

        # Compute reward
        reward = self._compute_reward(result.total_latency, all_ops_covered, done)

        last_result_str = (
            f"VALID. Latency={result.total_latency:.1f}, "
            f"working_set={result.working_set:.0f}"
        )

        obs = format_observation(
            graph=self.graph,
            state=self.state,
            last_action_result=last_result_str,
            last_action_error=None,
            max_steps=self.max_steps,
            naive_latency=self.naive_latency,
        )

        info = {
            "latency": result.total_latency,
            "working_set": result.working_set,
            "tile_latencies": result.tile_latencies,
            "total_latency": self.state.total_latency,
            "all_ops_covered": all_ops_covered,
        }

        return StepResult(
            observation=obs,
            reward=reward,
            done=done,
            info=info,
        )

    def get_state(self) -> dict:
        """Return full environment state for serialization."""
        assert self.state is not None
        return {
            "scheduled_op_ids": list(self.state.scheduled_op_ids),
            "tensors_in_fast_memory": list(self.state.tensors_in_fast_memory),
            "total_latency": self.state.total_latency,
            "step_count": self.state.step_count,
            "naive_latency": self.naive_latency,
            "schedule": [
                {
                    "ops": e.operation_ids,
                    "config": [e.config.w, e.config.h, e.config.k],
                    "retain": e.tensors_to_retain,
                    "latency": e.latency,
                }
                for e in self.state.schedule_history
            ],
        }

    def get_score(self) -> float:
        """
        Compute final score in [0, 1].
        Score = naive_latency / (naive_latency + agent_latency)
        This gives:
          - 0.5 if agent matches naive (no improvement)
          - approaching 1.0 as agent latency approaches 0
          - below 0.5 if agent is worse than naive
        Then we remap [0.5, 1.0] -> [0.0, 1.0] so that:
          - matching naive = 0.0
          - optimal (much better than naive) = approaching 1.0
        """
        if self.state is None or self.state.total_latency <= 0:
            return 0.0
        if not self._all_ops_covered():
            # Penalty: not all ops covered
            covered = len(self._covered_ops())
            total = len(self.graph.operations)
            coverage_ratio = covered / total if total > 0 else 0
            return coverage_ratio * 0.1  # max 0.1 if incomplete

        ratio = self.naive_latency / (self.naive_latency + self.state.total_latency)
        # ratio is in (0, 1). naive match gives 0.5. Better gives > 0.5.
        # Remap: score = (ratio - 0.5) * 2, clamped to [0, 1]
        # But also handle case where agent is worse than naive (ratio < 0.5)
        score = max(0.0, (ratio - 0.5) * 2.0)
        return min(score, 1.0)

    def _update_state(self, action: Action, latency: float):
        """Update schedule state after a valid action."""
        # Mark ops as scheduled
        for oid in action.operation_ids:
            self.state.scheduled_op_ids.add(oid)

        # Update fast memory:
        # 1. Remove all previously retained tensors that aren't in the new retain list
        #    (they get evicted when a new subgraph runs)
        # Actually, the spec says tensors_to_retain controls what stays AFTER this subgraph.
        # Everything else is evicted. So we clear fast memory and only keep retained tensors.

        # Clear fast memory
        self.state.tensors_in_fast_memory.clear()

        # Add retained tensors
        for tid in action.tensors_to_retain:
            self.state.tensors_in_fast_memory.add(tid)

        # Record in history
        entry = SubgraphEntry(
            operation_ids=list(action.operation_ids),
            config=action.config,
            tensors_to_retain=list(action.tensors_to_retain),
            traversal_order=action.traversal_order,
            latency=latency,
        )
        self.state.schedule_history.append(entry)

        # Update counters
        self.state.total_latency += latency
        self.state.step_count += 1

    def _covered_ops(self) -> set[int]:
        """All ops that have been scheduled at least once."""
        return set(self.state.scheduled_op_ids)

    def _all_ops_covered(self) -> bool:
        """Check if every operation has been scheduled at least once."""
        all_op_ids = set(op.id for op in self.graph.operations)
        return all_op_ids.issubset(self._covered_ops())

    def _compute_reward(
        self, step_latency: float, all_ops_covered: bool, done: bool
    ) -> float:
        """
        Compute per-step reward. Higher is better.
        Combines efficiency signal with completion incentive.
        """
        # Small positive baseline for any valid step (encourages exploration)
        valid_step_bonus = 0.02

        # Base reward: negative normalized latency (lower latency = higher reward)
        avg_naive_per_op = self.naive_latency / len(self.graph.operations)
        latency_reward = -step_latency / avg_naive_per_op * 0.1

        # Bonus for fusion (multiple ops in one subgraph)
        num_ops_in_step = len(
            self.state.schedule_history[-1].operation_ids
        ) if self.state.schedule_history else 1
        fusion_bonus = (num_ops_in_step - 1) * 0.05

        # Completion bonus
        completion_bonus = 0.0
        if done and all_ops_covered:
            score = self.naive_latency / self.state.total_latency
            completion_bonus = score * 0.5

        return valid_step_bonus + latency_reward + fusion_bonus + completion_bonus


def parse_action(text: str, graph: Graph) -> Optional[Action]:
    """
    Parse an action from LLM text output.
    Expected format: SCHEDULE ops=[0,1] config=[128,128,1] retain=[2]
    Flexible parsing to handle variations.
    """
    text = text.strip()

    # Extract operation IDs
    ops_match = re.search(r'ops\s*=\s*\[([^\]]*)\]', text)
    if not ops_match:
        # Try simpler format: just numbers
        ops_match = re.search(r'ops\s*=\s*(\d[\d,\s]*)', text)
    if not ops_match:
        return None

    try:
        ops_str = ops_match.group(1)
        op_ids = [int(x.strip()) for x in ops_str.split(',') if x.strip()]
    except (ValueError, IndexError):
        return None

    # Extract config [w, h, k]
    config_match = re.search(r'config\s*=\s*\[([^\]]*)\]', text)
    if not config_match:
        # Default to native granularity
        nw, nh = graph.hardware.native_granularity
        config = Config(nw, nh, 1)
    else:
        try:
            parts = [int(x.strip()) for x in config_match.group(1).split(',')]
            if len(parts) == 3:
                config = Config(parts[0], parts[1], parts[2])
            elif len(parts) == 2:
                config = Config(parts[0], parts[1], 1)
            else:
                nw, nh = graph.hardware.native_granularity
                config = Config(nw, nh, 1)
        except (ValueError, IndexError):
            nw, nh = graph.hardware.native_granularity
            config = Config(nw, nh, 1)

    # Extract retain list
    retain_match = re.search(r'retain\s*=\s*\[([^\]]*)\]', text)
    if retain_match:
        try:
            retain_str = retain_match.group(1).strip()
            if retain_str:
                retain = [int(x.strip()) for x in retain_str.split(',') if x.strip()]
            else:
                retain = []
        except (ValueError, IndexError):
            retain = []
    else:
        retain = []

    # Extract traversal order (optional)
    trav_match = re.search(r'traversal\s*=\s*\[([^\]]*)\]', text)
    traversal = None
    if trav_match:
        try:
            trav_str = trav_match.group(1).strip()
            if trav_str:
                traversal = [int(x.strip()) for x in trav_str.split(',')]
        except (ValueError, IndexError):
            traversal = None

    return Action(
        operation_ids=op_ids,
        config=config,
        tensors_to_retain=retain,
        traversal_order=traversal,
    )