server/environment.py

"""
Annotation QA Environment — Core Environment Logic.

Implements the OpenEnv 3-method interface:
- reset(task_id) → Observation
- step(action) → Observation
- state → State

The agent reviews intentionally-flawed annotations on real COCO val2017 images
and must correct bounding boxes, fix class labels, add missing annotations,
or remove spurious ones. Dense reward is provided at every step.
"""

import copy
import json
import os
import random
from pathlib import Path
from typing import Any, Dict, List, Optional
from uuid import uuid4

try:
    from openenv.core.env_server.types import Action, Observation, State
except ImportError:
    # Fallback for standalone
    pass

try:
    from ..models import (
        Annotation,
        AnnotationQAAction,
        AnnotationQAObservation,
        AnnotationQAState,
    )
except ImportError:
    from models import (
        Annotation,
        AnnotationQAAction,
        AnnotationQAObservation,
        AnnotationQAState,
    )
from .corruption import ALL_CLASSES, corrupt_annotations
from .grader import (
    compute_annotation_quality,
    compute_step_reward,
    grade_episode,
)


# ──────────────────────────────────────────────
# Task definitions
# ──────────────────────────────────────────────

TASK_CONFIGS = {
    "remove_spurious": {
        "description": (
            "Spurious Box Removal Task. Fake bounding boxes have been randomly drawn. "
            "Identify and remove any annotations that do not strictly bound a real object."
        ),
        "difficulty": "spurious",
        "max_steps": 15,
        "data_file": "task1_remove_spurious/samples.json",
    },
    "fix_classes": {
        "description": (
            "Class Identification Task. Some bounding boxes have incorrect class labels, "
            "and some are completely fake (spurious). Fix class labels using "
            "CHANGE_CLASS and REMOVE spurious labels."
        ),
        "difficulty": "classes",
        "max_steps": 20,
        "data_file": "task2_fix_classes/samples.json",
    },
    "find_missing": {
        "description": (
            "Contextual Object Detection Task. Bounding boxes for key objects have been "
            "entirely removed from the image. You must meticulously identify what object classes "
            "are completely missing from the drawn bounding boxes and flag them."
        ),
        "difficulty": "missing",
        "max_steps": 30,
        "data_file": "task3_find_missing/samples.json",
    },
}


class AnnotationQAEnvironment:
    """
    Annotation QA Environment following the OpenEnv pattern.

    The agent reviews real COCO val2017 image annotations that contain
    intentional errors and must correct them through a series of actions.
    A VLM is used to visually inspect the images.
    """

    SUPPORTS_CONCURRENT_SESSIONS = True

    def __init__(self):
        self._state = AnnotationQAState()
        self._gold_annotations: List[Dict] = []
        self._initial_annotations: List[Dict] = []
        self._current_annotations: List[Dict] = []
        self._scene_data: Dict[str, Any] = {}
        self._task_config: Dict[str, Any] = {}
        self._corrections_made: int = 0
        self._done: bool = False
        self._data_cache: Dict[str, Any] = {}
        self._next_ann_id: int = 0

        # Load data directory
        self._data_dir = Path(__file__).parent.parent / "data" / "tasks"

    def _load_task_data(self, task_id: str) -> List[Dict]:
        """Load and cache task data from disk."""
        if task_id in self._data_cache:
            return self._data_cache[task_id]

        config = TASK_CONFIGS[task_id]
        data_file = self._data_dir / config["data_file"]

        if not data_file.exists():
            raise FileNotFoundError(
                f"Task data file not found: {data_file}. "
                f"Run 'python -m data.prepare_coco' to generate the COCO dataset."
            )

        with open(data_file, "r") as f:
            data = json.load(f)

        self._data_cache[task_id] = data
        return data

    def reset(
        self,
        seed: Optional[int] = None,
        episode_id: Optional[str] = None,
        task: Optional[str] = None,
        **kwargs: Any,
    ) -> AnnotationQAObservation:
        """
        Start a new episode.

        Args:
            seed: Random seed for reproducibility
            episode_id: Optional episode ID
            task: Task ID — one of "fix_bboxes", "fix_classes", "batch_audit"
        """
        task_id = task or kwargs.get("task_id", "remove_spurious")
        if task_id not in TASK_CONFIGS:
            task_id = "remove_spurious"

        self._task_config = TASK_CONFIGS[task_id]
        data = self._load_task_data(task_id)

        # Select a random sample
        rng = random.Random(seed) if seed is not None else random.Random()

        scene = rng.choice(data)
        sample_seed = scene.get("seed", rng.randint(0, 99999))

        # Store gold annotations
        self._gold_annotations = copy.deepcopy(scene["gold_annotations"])
        self._scene_data = scene

        # Create corrupted annotations
        corrupted, corruption_log = corrupt_annotations(
            self._gold_annotations,
            self._task_config["difficulty"],
            sample_seed,
        )
        self._initial_annotations = copy.deepcopy(corrupted)
        self._current_annotations = copy.deepcopy(corrupted)
        self._corrections_made = 0
        self._done = False

        # Track next annotation ID
        self._next_ann_id = max((a["id"] for a in self._current_annotations), default=-1) + 1

        # Compute initial quality
        initial_quality = compute_annotation_quality(
            self._initial_annotations, self._gold_annotations
        )

        self._state = AnnotationQAState(
            episode_id=episode_id or str(uuid4()),
            step_count=0,
            task_id=task_id,
            sample_id=scene.get("scene_id", "unknown"),
            initial_quality=round(initial_quality, 4),
            current_quality=round(initial_quality, 4),
            corrections_made=0,
        )

        return self._build_observation(
            reward=None,
            message=(
                f"Review the annotations for this COCO image. "
                f"There are {len(self._current_annotations)} annotations. "
                f"Some may have incorrect bounding boxes, wrong class labels, "
                f"or be entirely spurious. Some objects may be missing annotations. "
                f"You have {self._task_config['max_steps']} steps to fix them."
            ),
        )

    def step(
        self,
        action: AnnotationQAAction,
        timeout_s: Optional[float] = None,
        **kwargs: Any,
    ) -> AnnotationQAObservation:
        """Execute a correction action and return updated observation with reward."""
        if self._done:
            return self._build_observation(
                reward=0.0,
                message="Episode is already done. Call reset() to start a new episode.",
            )

        self._state.step_count += 1
        error_msg = None

        # Save pre-action state for reward computation
        old_annotations = copy.deepcopy(self._current_annotations)

        # Process action
        try:
            if action.action_type == "adjust_bbox":
                error_msg = self._handle_adjust_bbox(action)
            elif action.action_type == "change_class":
                error_msg = self._handle_change_class(action)
            elif action.action_type == "add_annotation":
                error_msg = self._handle_add_annotation(action)
            elif action.action_type == "remove_annotation":
                error_msg = self._handle_remove_annotation(action)
            elif action.action_type == "submit":
                return self._handle_submit()
            elif action.action_type == "flag_safety":
                error_msg = self._handle_flag_safety(action)
            elif action.action_type == "change_attribute":
                error_msg = self._handle_change_attribute(action)
            elif action.action_type == "flag_missing":
                error_msg = self._handle_flag_missing(action)
            else:
                error_msg = f"Unknown action_type: {action.action_type}"
        except Exception as e:
            error_msg = f"Error processing action: {str(e)}"

        if error_msg is None:
            self._corrections_made += 1
            self._state.corrections_made = self._corrections_made

        # Compute reward
        if action.action_type == "flag_safety" and not error_msg:
            reward = 0.20
        elif action.action_type == "change_attribute" and not error_msg:
            reward = 0.15
        elif action.action_type == "flag_missing" and not error_msg:
            reward = 0.25
        else:
            reward = compute_step_reward(
                old_annotations,
                self._current_annotations,
                self._gold_annotations,
                action.action_type,
            )

        # Update quality tracking
        current_quality = compute_annotation_quality(
            self._current_annotations, self._gold_annotations
        )
        self._state.current_quality = round(current_quality, 4)

        # Check if max steps reached
        if self._state.step_count >= self._task_config["max_steps"]:
            self._done = True
            final_score = grade_episode(
                self._initial_annotations,
                self._current_annotations,
                self._gold_annotations,
            )
            return self._build_observation(
                reward=final_score,
                message=f"Max steps reached. Final score: {final_score:.3f}",
                error=error_msg,
            )

        return self._build_observation(
            reward=reward,
            message=(
                f"{'Error: ' + error_msg if error_msg else 'Correction applied.'} "
                f"Quality: {current_quality:.3f} "
                f"(was {self._state.initial_quality:.3f}). "
                f"Steps remaining: {self._task_config['max_steps'] - self._state.step_count}"
            ),
            error=error_msg,
        )

    @property
    def state(self) -> AnnotationQAState:
        """Get current episode state."""
        return self._state

    def close(self) -> None:
        """Clean up environment resources."""
        pass

    async def reset_async(self, **kwargs) -> AnnotationQAObservation:
        """Async wrapper for reset (required by OpenEnv server interface)."""
        return self.reset(**kwargs)

    async def step_async(self, action: AnnotationQAAction, **kwargs) -> AnnotationQAObservation:
        """Async wrapper for step (required by OpenEnv server interface)."""
        return self.step(action, **kwargs)

    # ──────────────────────────────────────────
    # Action handlers
    # ──────────────────────────────────────────

    def _handle_adjust_bbox(self, action: AnnotationQAAction) -> Optional[str]:
        """Adjust the bounding box of an existing annotation."""
        if action.annotation_id is None:
            return "annotation_id is required for adjust_bbox"
        if action.new_bbox is None:
            return "new_bbox is required for adjust_bbox"
        if len(action.new_bbox) != 4:
            return "new_bbox must have exactly 4 values [x, y, w, h]"

        ann = self._find_annotation(action.annotation_id)
        if ann is None:
            return f"Annotation {action.annotation_id} not found"

        # Validate bbox values
        for v in action.new_bbox:
            if not (0.0 <= v <= 1.0):
                return "All bbox values must be between 0.0 and 1.0"

        ann["bbox"] = [round(v, 4) for v in action.new_bbox]
        return None

    def _handle_change_class(self, action: AnnotationQAAction) -> Optional[str]:
        """Change the class label of an existing annotation."""
        if action.annotation_id is None:
            return "annotation_id is required for change_class"
        if action.new_class is None:
            return "new_class is required for change_class"
        if action.new_class not in ALL_CLASSES:
            return f"Invalid class '{action.new_class}'. Valid: {ALL_CLASSES}"

        ann = self._find_annotation(action.annotation_id)
        if ann is None:
            return f"Annotation {action.annotation_id} not found"

        ann["class_label"] = action.new_class
        return None

    def _handle_add_annotation(self, action: AnnotationQAAction) -> Optional[str]:
        """Add a new annotation."""
        if action.new_bbox is None:
            return "new_bbox is required for add_annotation"
        if action.new_class is None:
            return "new_class is required for add_annotation"
        if len(action.new_bbox) != 4:
            return "new_bbox must have exactly 4 values [x, y, w, h]"
        if action.new_class not in ALL_CLASSES:
            return f"Invalid class '{action.new_class}'. Valid: {ALL_CLASSES}"

        for v in action.new_bbox:
            if not (0.0 <= v <= 1.0):
                return "All bbox values must be between 0.0 and 1.0"

        new_ann = {
            "id": self._next_ann_id,
            "bbox": [round(v, 4) for v in action.new_bbox],
            "class_label": action.new_class,
        }
        self._current_annotations.append(new_ann)
        self._next_ann_id += 1
        return None

    def _handle_remove_annotation(self, action: AnnotationQAAction) -> Optional[str]:
        """Remove an annotation."""
        if action.annotation_id is None:
            return "annotation_id is required for remove_annotation"

        idx = self._find_annotation_index(action.annotation_id)
        if idx is None:
            return f"Annotation {action.annotation_id} not found"

        self._current_annotations.pop(idx)
        return None

    def _handle_submit(self) -> AnnotationQAObservation:
        """Submit corrections and compute final grade."""
        self._done = True
        final_score = grade_episode(
            self._initial_annotations,
            self._current_annotations,
            self._gold_annotations,
        )

        return self._build_observation(
            reward=final_score,
            message=(
                f"Corrections submitted! "
                f"Final score: {final_score:.3f}. "
                f"Quality went from {self._state.initial_quality:.3f} "
                f"to {self._state.current_quality:.3f} over "
                f"{self._state.step_count} steps."
            ),
        )

    def _handle_flag_safety(self, action: AnnotationQAAction) -> Optional[str]:
        if action.annotation_id is None:
            return "annotation_id is required for flag_safety"
        ann = self._find_annotation(action.annotation_id)
        if ann is None: return "Annotation not found"
        # We don't change state, just append tracking metadata for the grader
        ann["safety_flagged"] = True
        return None

    def _handle_change_attribute(self, action: AnnotationQAAction) -> Optional[str]:
        if action.annotation_id is None:
            return "annotation_id is required for change_attribute"
        if not action.new_attribute:
            return "new_attribute is required"
        ann = self._find_annotation(action.annotation_id)
        if ann is None: return "Annotation not found"
        ann["class_label"] = action.new_attribute
        return None

    def _handle_flag_missing(self, action: AnnotationQAAction) -> Optional[str]:
        if not action.missing_class:
            return "missing_class is required for flag_missing"
        # Flagging missing class adds a placeholder marker
        self._current_annotations.append({
            "id": self._next_ann_id,
            "bbox": [0,0,0,0],
            "class_label": f"missing_{action.missing_class}"
        })
        self._next_ann_id += 1
        return None

    # ──────────────────────────────────────────
    # Helpers
    # ──────────────────────────────────────────

    def _find_annotation(self, ann_id: int) -> Optional[Dict]:
        for ann in self._current_annotations:
            if ann["id"] == ann_id:
                return ann
        return None

    def _find_annotation_index(self, ann_id: int) -> Optional[int]:
        for i, ann in enumerate(self._current_annotations):
            if ann["id"] == ann_id:
                return i
        return None

    def _build_observation(
        self,
        reward: Optional[float],
        message: str,
        error: Optional[str] = None,
    ) -> AnnotationQAObservation:
        """Build an observation from current state."""
        return AnnotationQAObservation(
            done=self._done,
            reward=reward,
            # Image info from COCO
            image_url=self._scene_data.get("image_url"),
            image_width=self._scene_data.get("image_width", 0),
            image_height=self._scene_data.get("image_height", 0),
            # Scene info
            scene_description=self._scene_data.get("scene_description", ""),
            scene_objects=[
                {
                    "id": obj["id"],
                    "class_label": obj["class_label"],
                    "position": obj.get("position", ""),
                    "bbox": obj["bbox"],
                }
                for obj in self._scene_data.get("objects", [])
            ],
            annotations=[
                Annotation(
                    id=ann["id"],
                    bbox=ann["bbox"],
                    class_label=ann["class_label"],
                )
                for ann in self._current_annotations
            ],
            available_classes=ALL_CLASSES,
            task_id=self._state.task_id,
            task_description=self._task_config.get("description", ""),
            corrections_made=self._corrections_made,
            step_count=self._state.step_count,
            max_steps=self._task_config.get("max_steps", 20),
            message=message,
            last_action_error=error,
        )