agents/caption.py

"""
Caption Child Agent — scene captioning powered by ArcisVLM.

Responsibilities:
- Process frames with <caption> task token prefix
- Generate rich scene descriptions and summaries
- Parse model output into structured captioning results
- Support varying levels of detail (brief, standard, detailed)
"""

from __future__ import annotations

import re
import logging
from pathlib import Path
from typing import Any

import torch
from torchvision import transforms

from agents.base import BaseAgent, Task, Result, EXPERT_CAPTION
from agents.postprocess import enrich_result
from model.vlm import VLJEPAModel
from model.tokenizer import BPETokenizer

logger = logging.getLogger(__name__)

_DEFAULT_TRANSFORM = transforms.Compose([
    transforms.Resize((448, 448)),
    transforms.ToTensor(),
    transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
])

# Scene attributes vocabulary for parsing captions
_SCENE_TYPES = [
    "indoor", "outdoor", "street", "road", "parking", "office", "warehouse",
    "lobby", "hallway", "entrance", "garden", "park", "intersection",
    "building", "store", "shop", "restaurant", "residential",
]

_LIGHTING_TERMS = [
    "bright", "dark", "dim", "daylight", "nighttime", "overcast",
    "sunny", "cloudy", "shadow", "illuminated", "fluorescent",
]

_WEATHER_TERMS = [
    "rain", "snow", "fog", "mist", "clear", "cloudy", "windy", "storm",
]


class CaptionAgent(BaseAgent):
    """
    Scene captioning agent.

    Uses the VL-JEPA + MoE model with the <caption> task token to produce
    descriptive scene summaries. Supports different verbosity levels through
    query phrasing.

    Args:
        agent_id: Unique identifier.
        model: Pre-loaded VLJEPAModel.
        tokenizer: BPETokenizer instance.
        device: Torch device.
        max_new_tokens: Maximum tokens to generate per query.
        temperature: Sampling temperature.
        transform: Image preprocessing transform.
    """

    def __init__(
        self,
        agent_id: str,
        model: VLJEPAModel,
        tokenizer: BPETokenizer,
        device: str = "cpu",
        max_new_tokens: int = 256,
        temperature: float = 0.0,
        transform: transforms.Compose | None = None,
    ) -> None:
        super().__init__(agent_id=agent_id, expert_type=EXPERT_CAPTION)
        self.model = model
        self.tokenizer = tokenizer
        self.device = torch.device(device)
        self.max_new_tokens = max_new_tokens
        self.temperature = temperature
        self.transform = transform or _DEFAULT_TRANSFORM

        # (task tokens removed — training uses plain BOS/EOS from tokenizer)

    # -- Lifecycle -----------------------------------------------------------

    def on_start(self) -> None:
        self.model.to(self.device)
        self.model.eval()
        logger.info("CaptionAgent %s: model on %s", self.agent_id, self.device)

    def on_stop(self) -> None:
        self.model.cpu()

    # -- Image loading -------------------------------------------------------

    def _load_image(self, image_ref: str) -> torch.Tensor:
        path = Path(image_ref)
        if path.exists() and path.is_file():
            from PIL import Image
            img = Image.open(path).convert("RGB")
            tensor = self.transform(img)
            return tensor.unsqueeze(0).to(self.device)
        raise FileNotFoundError(f"Image not found: {image_ref}")

    # -- Query preparation ---------------------------------------------------

    def _prepare_query(self, query: str) -> tuple[torch.Tensor, torch.Tensor]:
        """Tokenize query with BOS/EOS and pad to max_q=64 to match training format."""
        max_q = 64  # Must match training
        ids = self.tokenizer.encode(query, add_special=True)  # BOS + tokens + EOS
        pad_id = self.tokenizer.pad_id
        if len(ids) > max_q:
            ids = ids[:max_q]
        else:
            ids = ids + [pad_id] * (max_q - len(ids))
        query_ids = torch.tensor([ids], dtype=torch.long, device=self.device)
        padding_mask = (query_ids != pad_id).long()
        return query_ids, padding_mask

    # -- Parsing model output ------------------------------------------------

    @staticmethod
    def _parse_caption(answer: str) -> dict[str, Any]:
        """
        Parse the model's free-form caption into structured scene attributes.

        Extracts:
        - scene_type: detected scene category (indoor, outdoor, street, etc.)
        - lighting: lighting conditions mentioned
        - weather: weather conditions if mentioned
        - objects_mentioned: list of notable objects referenced in the caption
        - sentence_count: number of sentences in the caption
        """
        answer_lower = answer.lower()
        result: dict[str, Any] = {
            "scene_type": "unknown",
            "lighting": [],
            "weather": [],
            "objects_mentioned": [],
            "sentence_count": 0,
        }

        # Detect scene type
        for scene in _SCENE_TYPES:
            if re.search(rf"\b{re.escape(scene)}\b", answer_lower):
                result["scene_type"] = scene
                break

        # Detect lighting conditions
        for term in _LIGHTING_TERMS:
            if re.search(rf"\b{re.escape(term)}\b", answer_lower):
                result["lighting"].append(term)

        # Detect weather
        for term in _WEATHER_TERMS:
            if re.search(rf"\b{re.escape(term)}\b", answer_lower):
                result["weather"].append(term)

        # Extract mentioned objects
        object_pattern = re.compile(
            r"\b(person|people|car|vehicle|truck|bus|bicycle|motorcycle|"
            r"dog|cat|bird|chair|table|tree|building|sign|pedestrian|"
            r"traffic light|bench|fence|wall|door|window)\b",
            re.IGNORECASE,
        )
        for match in object_pattern.finditer(answer_lower):
            obj = match.group(1).lower()
            if obj == "people":
                obj = "person"
            if obj not in result["objects_mentioned"]:
                result["objects_mentioned"].append(obj)

        # Count sentences
        sentences = re.split(r"[.!?]+", answer.strip())
        result["sentence_count"] = len([s for s in sentences if s.strip()])

        return result

    # -- Core processing -----------------------------------------------------

    def process(self, task: Task) -> Result:
        """Run captioning inference for a single task."""
        image_ref = task.image_ref
        query = task.query or "Describe this scene in detail."

        if not image_ref:
            return Result(answer="", error="No image_ref provided", expert_used=EXPERT_CAPTION)

        # Load image
        payload_tensor = task.payload.get("image_tensor")
        if payload_tensor is not None and isinstance(payload_tensor, torch.Tensor):
            if payload_tensor.dim() == 3:
                payload_tensor = payload_tensor.unsqueeze(0)
            images = payload_tensor.to(self.device)
        else:
            images = self._load_image(image_ref)

        # Prepare query with <caption> token
        query_ids, query_mask = self._prepare_query(query)

        # Generate
        with torch.no_grad():
            generated_ids = self.model.generate(
                images=images,
                query_ids=query_ids,
                query_padding_mask=query_mask,
                max_new_tokens=self.max_new_tokens,
                temperature=self.temperature,
            )

        answer = self.tokenizer.decode(generated_ids[0].tolist())

        # Estimate confidence
        with torch.no_grad():
            embed = self.model.get_embedding(images, query_ids, query_mask)
            norm = embed.norm(dim=-1).mean().item()
            confidence = min(1.0, max(0.0, norm / (norm + 1.0)))

        # Parse caption attributes
        caption_info = self._parse_caption(answer)

        result = Result(
            answer=answer,
            confidence=round(confidence, 4),
            expert_used=EXPERT_CAPTION,
            metadata={
                "image_ref": image_ref,
                "query": query,
                "scene_type": caption_info["scene_type"],
                "lighting": caption_info["lighting"],
                "weather": caption_info["weather"],
                "objects_mentioned": caption_info["objects_mentioned"],
                "sentence_count": caption_info["sentence_count"],
                "tokens_generated": generated_ids.shape[1],
            },
        )
        return enrich_result(result, image_path=image_ref)

    # -- Health check --------------------------------------------------------

    def health_check(self) -> bool:
        try:
            return self._started
        except Exception:
            return False