src/environment.py

"""
FusionOps V2 - Environment

Sequential step environment:
- Walk graph in topo order
- At each node, agent decides: fuse_with_prev, tile, retain
- Episode ends when all nodes are scheduled
- Reward = improvement over greedy baseline + step signals

Observation: compact JSON (not prose)
Action: JSON with fuse_with_prev, tile, retain
"""

from __future__ import annotations

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

from .models import (
    Graph, Node, Action, FusionGroup, ScheduleState, OpType,
    POINTWISE_OPS, COMPUTE_OPS,
)
from .cost_model import (
    compute_group_cost, compute_greedy_baseline, compute_naive_baseline,
    CostBreakdown,
)


@dataclass
class StepResult:
    observation: str      # compact JSON string
    reward: float
    done: bool
    info: dict = field(default_factory=dict)
    score: Optional[float] = None


class FusionOpsEnv:
    """
    RL environment for ML computation graph scheduling.

    Sequential decisions: walk topo order, at each node decide
    fuse/tile/retain. One step per node.
    """

    def __init__(self, graph: Graph, max_steps: Optional[int] = None):
        self.graph = graph
        self.max_steps = max_steps or len(graph.nodes)
        self.topo_order = graph.topo_order()

        # Precompute baselines
        self.greedy_latency = compute_greedy_baseline(graph)
        self.naive_latency = compute_naive_baseline(graph)

        self.state: Optional[ScheduleState] = None

    def reset(self) -> StepResult:
        """Initialize episode. Returns first observation."""
        self.state = ScheduleState()
        self.state.remaining_uses = {
            n.id: self.graph.future_uses(n.id) for n in self.graph.nodes
        }
        self.state.current_node_idx = 0
        self.state.current_group = FusionGroup(
            node_ids=[], tile=128, retained=[]
        )

        obs = self._make_observation()
        return StepResult(observation=obs, reward=0.0, done=False)

    def step(self, action: Action) -> StepResult:
        """Execute one step: process the current node with the given action."""
        assert self.state is not None, "Must call reset() first"

        if self.state.current_node_idx >= len(self.topo_order):
            return StepResult(
                observation=self._make_observation(),
                reward=0.0,
                done=True,
                score=self.get_score(),
            )

        current_nid = self.topo_order[self.state.current_node_idx]
        current_node = self.graph.nodes[current_nid]
        hw = self.graph.hardware

        # Validate tile
        if action.tile not in Action.VALID_TILES:
            action.tile = 128  # default

        # Validate retain: can only retain nodes whose outputs exist
        valid_retain = []
        for r in action.retain:
            if 0 <= r < len(self.graph.nodes):
                valid_retain.append(r)
        action.retain = valid_retain

        step_reward = 0.0
        error = None

        if action.fuse_with_prev and self.state.current_group.node_ids:
            # Try to fuse with current group
            proposed_ids = self.state.current_group.node_ids + [current_nid]

            # Check connectivity
            connected = any(
                inp in self.state.current_group.node_ids
                for inp in current_node.inputs
            )

            # Check fusion depth limit
            within_limit = len(proposed_ids) <= hw.max_fusion_depth

            if connected and within_limit:
                # Test if the fused group is valid (memory check)
                test_group = FusionGroup(
                    node_ids=proposed_ids,
                    tile=action.tile,
                    retained=action.retain,
                )
                cost = compute_group_cost(self.graph, test_group, self.state)

                if cost.is_valid:
                    # Successful fusion
                    self.state.current_group = test_group
                    step_reward += 0.02  # small reward for valid fusion
                else:
                    # OOM: can't fuse, finalize current group and start new
                    error = cost.error
                    step_reward -= 0.01  # small penalty for OOM attempt
                    self._finalize_current_group()
                    self.state.current_group = FusionGroup(
                        node_ids=[current_nid],
                        tile=action.tile,
                        retained=action.retain,
                    )
            else:
                # Not connected or exceeds depth: finalize and start new
                if not connected:
                    error = "Cannot fuse: not connected to current group"
                else:
                    error = f"Cannot fuse: exceeds max fusion depth {hw.max_fusion_depth}"
                step_reward -= 0.005
                self._finalize_current_group()
                self.state.current_group = FusionGroup(
                    node_ids=[current_nid],
                    tile=action.tile,
                    retained=action.retain,
                )
        else:
            # Don't fuse: finalize current group (if any) and start new
            self._finalize_current_group()
            self.state.current_group = FusionGroup(
                node_ids=[current_nid],
                tile=action.tile,
                retained=action.retain,
            )

        # Update remaining uses for consumed inputs
        for inp_id in current_node.inputs:
            if inp_id in self.state.remaining_uses:
                self.state.remaining_uses[inp_id] -= 1
                # If this was the last use, the tensor can be freed
                if self.state.remaining_uses[inp_id] <= 0:
                    self.state.fast_mem_contents.discard(inp_id)

        # Retention step reward
        for r in action.retain:
            remaining = self.state.remaining_uses.get(r, 0)
            if remaining > 0:
                step_reward += 0.01  # good: retaining something that will be reused
            else:
                step_reward -= 0.01  # bad: retaining dead tensor, wastes memory

        # Reload penalty
        for inp_id in current_node.inputs:
            if inp_id not in self.state.fast_mem_contents and inp_id >= 0:
                # Had to reload from slow memory
                self.state.total_reloads += 1
                step_reward -= 0.005

        # Advance
        self.state.step += 1
        self.state.current_node_idx += 1

        # Check if episode is done
        done = self.state.current_node_idx >= len(self.topo_order)
        if done:
            # Finalize last group
            self._finalize_current_group()

        # Check step limit
        if self.state.step >= self.max_steps and not done:
            done = True
            self._finalize_current_group()

        score = self.get_score() if done else None

        # Final reward
        if done:
            # Episode completion reward based on improvement over greedy
            improvement = (self.greedy_latency - self.state.total_latency) / self.greedy_latency
            step_reward += max(0, improvement) * 0.5  # big bonus for beating greedy
            if improvement < 0:
                step_reward += improvement * 0.1  # smaller penalty for being worse

        obs = self._make_observation(error=error)
        return StepResult(
            observation=obs,
            reward=step_reward,
            done=done,
            info={
                "total_latency": self.state.total_latency,
                "greedy_latency": self.greedy_latency,
                "kernel_launches": self.state.total_kernel_launches,
                "reloads": self.state.total_reloads,
                "error": error,
            },
            score=score,
        )

    def _finalize_current_group(self):
        """Finalize the current fusion group: compute cost, update state."""
        if not self.state.current_group or not self.state.current_group.node_ids:
            return

        group = self.state.current_group
        cost = compute_group_cost(self.graph, group, self.state)

        if cost.is_valid:
            group.latency = cost.total_latency
            self.state.total_latency += cost.total_latency
            self.state.total_kernel_launches += 1

            # Update fast memory
            # Remove tensors not retained
            new_fast_mem = set()
            for r in group.retained:
                new_fast_mem.add(r)
            # Keep previously retained tensors that still have uses
            for existing in self.state.fast_mem_contents:
                if self.state.remaining_uses.get(existing, 0) > 0:
                    new_fast_mem.add(existing)
            self.state.fast_mem_contents = new_fast_mem

            # Compute fast mem used
            self.state.fast_mem_used = sum(
                self.graph.nodes[nid].output_bytes
                for nid in self.state.fast_mem_contents
            )
        else:
            # Fallback: execute each op individually
            for nid in group.node_ids:
                solo = FusionGroup(node_ids=[nid], tile=128, retained=[])
                sc = compute_group_cost(self.graph, solo, self.state)
                if sc.is_valid:
                    self.state.total_latency += sc.total_latency
                else:
                    self.state.total_latency += 1e6  # catastrophic failure
                self.state.total_kernel_launches += 1

        self.state.fusion_groups.append(group)
        self.state.current_group = FusionGroup(
            node_ids=[], tile=128, retained=[]
        )

    def _make_observation(self, error: Optional[str] = None) -> str:
        """
        Compact JSON observation.
        Designed for LLM consumption: short, structured, stable schema.

        Step 0: full graph (compact per-node: id, op, inputs only -- no shape
                 unless graph has >20 nodes, then only show first 20 + summary)
        Steps 1+: current node + 2-node lookahead only
        """
        state = self.state

        # Current node info
        if state.current_node_idx < len(self.topo_order):
            current_nid = self.topo_order[state.current_node_idx]
            current_node = self.graph.nodes[current_nid]
        else:
            current_nid = -1
            current_node = None

        # Future uses for visible tensors (only non-zero)
        future_uses = {}
        for nid, uses in state.remaining_uses.items():
            if uses > 0:
                future_uses[str(nid)] = uses

        # Current fusion group
        current_group_info = {
            "node_ids": state.current_group.node_ids if state.current_group else [],
            "tile": state.current_group.tile if state.current_group else 128,
        }

        obs = {
            "step": state.step,
            "current_node": current_nid,
            "current_group": current_group_info,
            "fast_mem": sorted(state.fast_mem_contents),
            "fast_mem_used": state.fast_mem_used,
            "capacity": self.graph.hardware.fast_mem_capacity,
            "max_fusion": self.graph.hardware.max_fusion_depth,
            "future_uses": future_uses,
            "total_latency": round(state.total_latency, 1),
            "kernel_launches": state.total_kernel_launches,
            "greedy_baseline": round(self.greedy_latency, 1),
        }

        if error:
            obs["error"] = error

        # Step 0: include graph structure (compact)
        if state.step == 0:
            n_nodes = len(self.graph.nodes)
            # For large graphs (>20 nodes), only include shape for the first node
            # and omit shape for the rest to save tokens
            nodes_data = []
            for node in self.graph.nodes:
                entry = {
                    "id": node.id,
                    "op": node.op.value,
                    "inputs": node.inputs,
                }
                if n_nodes <= 20 or node.id == 0:
                    entry["shape"] = node.shape
                nodes_data.append(entry)
            obs["nodes"] = nodes_data
            obs["tensor_bytes"] = self.graph.nodes[0].output_bytes
        else:
            # Steps 1+: only current node + lookahead
            if current_node:
                obs["current_node_info"] = {
                    "id": current_nid,
                    "op": current_node.op.value,
                    "shape": current_node.shape,
                    "inputs": current_node.inputs,
                    "output_bytes": current_node.output_bytes,
                }
                # Show next 2 nodes for lookahead
                lookahead = []
                for offset in range(1, 3):
                    next_idx = state.current_node_idx + offset
                    if next_idx < len(self.topo_order):
                        next_nid = self.topo_order[next_idx]
                        nn = self.graph.nodes[next_nid]
                        lookahead.append({
                            "id": next_nid,
                            "op": nn.op.value,
                            "inputs": nn.inputs,
                        })
                if lookahead:
                    obs["lookahead"] = lookahead

        return json.dumps(obs, separators=(",", ":"))

    def get_score(self) -> float:
        """
        Score in [0, 1].
        0.0 = same as greedy baseline
        1.0 = 20%+ improvement over greedy (calibrated to achievable ceiling)
        Negative if worse than greedy.
        """
        if self.state is None or self.state.total_latency <= 0:
            return 0.0

        improvement = (self.greedy_latency - self.state.total_latency) / self.greedy_latency

        # Map to [0, 1]: 0% -> 0.0, 20%+ -> 1.0
        # Retention gap tests show 12-18% is achievable, so 20% = perfect
        score = improvement / 0.2
        return max(-1.0, min(1.0, score))

    def get_state(self) -> dict:
        """Full state for debugging."""
        assert self.state is not None
        return {
            "step": self.state.step,
            "current_node_idx": self.state.current_node_idx,
            "total_latency": self.state.total_latency,
            "greedy_latency": self.greedy_latency,
            "naive_latency": self.naive_latency,
            "kernel_launches": self.state.total_kernel_launches,
            "reloads": self.state.total_reloads,
            "fast_mem_contents": sorted(self.state.fast_mem_contents),
            "fast_mem_used": self.state.fast_mem_used,
            "fusion_groups": [
                {
                    "node_ids": g.node_ids,
                    "tile": g.tile,
                    "retained": g.retained,
                    "latency": g.latency,
                }
                for g in self.state.fusion_groups
            ],
            "score": self.get_score(),
        }


def parse_action(text: str) -> Optional[Action]:
    """
    Parse action from LLM output.
    Accepts JSON: {"fuse_with_prev": true, "tile": 128, "retain": [1, 3]}
    Also accepts natural language fallback.
    """
    text = text.strip()

    # Try JSON parse first
    try:
        # Handle markdown code fences
        if text.startswith("```"):
            lines = text.split("\n")
            text = "\n".join(l for l in lines if not l.startswith("```"))
            text = text.strip()

        data = json.loads(text)
        return Action(
            fuse_with_prev=bool(data.get("fuse_with_prev", False)),
            tile=int(data.get("tile", 128)),
            retain=list(data.get("retain", [])),
        )
    except (json.JSONDecodeError, TypeError, ValueError):
        pass

    # Fallback: parse key=value format
    import re
    fuse = False
    tile = 128
    retain = []

    fuse_match = re.search(r'fuse[_\s]*(?:with[_\s]*prev)?[=:\s]*(true|false|yes|no|1|0)', text, re.I)
    if fuse_match:
        fuse = fuse_match.group(1).lower() in ("true", "yes", "1")

    tile_match = re.search(r'tile[=:\s]*(\d+)', text, re.I)
    if tile_match:
        tile = int(tile_match.group(1))

    retain_match = re.search(r'retain[=:\s]*\[([^\]]*)\]', text, re.I)
    if retain_match:
        retain_str = retain_match.group(1).strip()
        if retain_str:
            retain = [int(x.strip()) for x in retain_str.split(",") if x.strip()]

    return Action(fuse_with_prev=fuse, tile=tile, retain=retain)