visualization/animation.py

"""
Generic animation utilities for WorldSmithAI.

This module turns rendered world states into GIF or MP4 animations. It is built
to work from a root-level Hugging Face Spaces ``app.py`` and does not assume any
``app/`` package or special folder layout.

The animation layer is domain-agnostic. It does not know about farms,
civilizations, research ecosystems, transport networks, power grids, fantasy
worlds, or social networks. It advances a runtime world, asks
``visualization.renderer`` to render each state, and writes the resulting frames
to an animation file.

Gradio usage:
    from visualization.animation import animate_world_to_gif_path

    def simulate(prompt: str) -> str:
        world = build_world_from_prompt(prompt)
        return animate_world_to_gif_path(world, frame_count=60)

MP4 usage:
    from visualization.animation import animate_world_to_mp4_path

    def simulate_video(prompt: str) -> str:
        world = build_world_from_prompt(prompt)
        return animate_world_to_mp4_path(world, frame_count=90)

Future extensibility:
    - Add event overlays and behavior traces.
    - Add side-by-side charts in animation frames.
    - Add camera tracking for selected agents.
    - Add graph/network animation mode.
    - Add streaming frame generation for long simulations.
    - Add optional WebM writer for browser-native video output.
"""

from __future__ import annotations

import copy
import inspect
import logging
import tempfile
from collections.abc import Callable, Mapping, MutableSequence, Sequence
from dataclasses import dataclass, field
from enum import Enum
from numbers import Real
from pathlib import Path
from types import MappingProxyType
from typing import TYPE_CHECKING, Any, ClassVar

import matplotlib

matplotlib.use("Agg", force=False)

import matplotlib.pyplot as plt
import numpy as np
from matplotlib.animation import FFMpegWriter, FuncAnimation, PillowWriter, writers
from matplotlib.figure import Figure

from visualization.renderer import (
    RenderSnapshot,
    RendererConfig,
    WorldRenderer,
    close_figure,
    figure_to_rgb_array,
)

if TYPE_CHECKING:
    from core.world import World


logger = logging.getLogger(__name__)

_EPSILON = 1.0e-12


class AnimationFormat(str, Enum):
    """Supported animation output formats."""

    GIF = "gif"
    MP4 = "mp4"


class AnimationAdvanceMode(str, Enum):
    """Strategies for advancing a world between frames."""

    WORLD_STEP = "world_step"
    SCHEDULER_STEP = "scheduler_step"
    AUTO = "auto"
    NONE = "none"


@dataclass(frozen=True)
class AnimationFrame:
    """One captured animation frame.

    The image array is intentionally omitted from ``to_dict`` to keep metadata
    small and JSON-friendly.
    """

    index: int
    step: int | None
    image: np.ndarray = field(repr=False, compare=False)
    snapshot: RenderSnapshot | None = None
    metadata: Mapping[str, Any] = field(default_factory=dict)

    @property
    def shape(self) -> tuple[int, ...]:
        """Return the underlying image array shape."""

        return tuple(int(value) for value in self.image.shape)

    def to_dict(self) -> dict[str, Any]:
        """Return a JSON-friendly frame metadata dictionary."""

        return {
            "index": self.index,
            "step": self.step,
            "shape": list(self.shape),
            "snapshot": None if self.snapshot is None else self.snapshot.to_dict(),
            "metadata": copy.deepcopy(dict(self.metadata)),
        }


@dataclass(frozen=True)
class AnimationResult:
    """Result returned after writing an animation."""

    path: str
    format: AnimationFormat
    frame_count: int
    fps: float
    frames: tuple[AnimationFrame, ...] = field(default_factory=tuple, repr=False, compare=False)
    metadata: Mapping[str, Any] = field(default_factory=dict)

    def to_dict(self, *, include_frame_metadata: bool = True) -> dict[str, Any]:
        """Return a JSON-friendly animation result dictionary."""

        return {
            "path": self.path,
            "format": self.format.value,
            "frame_count": self.frame_count,
            "fps": self.fps,
            "frames": [frame.to_dict() for frame in self.frames]
            if include_frame_metadata
            else [],
            "metadata": copy.deepcopy(dict(self.metadata)),
        }


@dataclass
class AnimationConfig:
    """Configuration for ``WorldAnimator``.

    Defaults favor GIF output because GIFs are easy to display in a root-level
    Gradio ``app.py`` and do not require ffmpeg.
    """

    frame_count: int = 60
    steps_per_frame: int = 1
    include_initial_frame: bool = True

    format: AnimationFormat | str = AnimationFormat.GIF
    fps: float | None = None
    interval_ms: int = 150
    repeat: bool = True

    advance_mode: AnimationAdvanceMode | str = AnimationAdvanceMode.AUTO
    ensure_step_count_progress: bool = True

    renderer_config: RendererConfig = field(default_factory=RendererConfig)

    output_dir: str | Path | None = None
    filename_prefix: str = "worldsmithai_animation_"

    dpi: int | None = None
    close_frame_figures: bool = True
    close_animation_figure: bool = True

    title: str | None = None
    show_animation_title: bool = False

    metadata: Mapping[str, Any] = field(default_factory=dict)

    def resolved_fps(self) -> float:
        """Return effective frames per second."""

        if self.fps is not None:
            return max(_EPSILON, float(self.fps))

        interval = max(1, int(self.interval_ms))
        return 1000.0 / float(interval)

    def resolved_format(self) -> AnimationFormat:
        """Return effective animation format."""

        return normalize_animation_format(self.format)


class AnimationError(RuntimeError):
    """Base exception raised by animation utilities."""


class AnimationWriterError(AnimationError):
    """Raised when animation frames cannot be written to the requested format."""


class WorldAdvanceError(AnimationError):
    """Raised when the world cannot be advanced between frames."""


@dataclass
class WorldAnimator:
    """Capture and write animations for generic WorldSmithAI worlds.

    The animator mutates the supplied world by advancing it between frames.
    This is usually what a simulation animation should do. Callers that need
    to preserve the original world should pass in a copied world object.
    """

    config: AnimationConfig = field(default_factory=AnimationConfig)

    name: ClassVar[str] = "world_animator"

    def capture_frames(
        self,
        world: World,
        *,
        frame_count: int | None = None,
        progress_callback: Callable[[int, int], None] | None = None,
    ) -> tuple[AnimationFrame, ...]:
        """Capture rendered frames while advancing the world.

        Args:
            world: Runtime world object.
            frame_count: Optional override for configured frame count.
            progress_callback: Optional callback receiving
                ``(frame_index, total_frames)``. This is useful for app-level
                progress reporting without importing Gradio here.

        Returns:
            Tuple of captured animation frames.
        """

        total_frames = max(1, int(frame_count if frame_count is not None else self.config.frame_count))
        renderer = WorldRenderer(config=self.config.renderer_config)
        frames: list[AnimationFrame] = []

        for frame_index in range(total_frames):
            if frame_index == 0 and self.config.include_initial_frame:
                pass
            else:
                self.advance_world(world, steps=max(0, int(self.config.steps_per_frame)))

            frame = self._capture_single_frame(
                world,
                renderer=renderer,
                index=frame_index,
            )
            frames.append(frame)

            if progress_callback is not None:
                progress_callback(frame_index + 1, total_frames)

        logger.debug(
            "Captured %s animation frame(s) for world at step %s",
            len(frames),
            _world_step(world),
        )

        return tuple(frames)

    def animate(
        self,
        world: World,
        *,
        output_path: str | Path | None = None,
        format: AnimationFormat | str | None = None,
        progress_callback: Callable[[int, int], None] | None = None,
    ) -> AnimationResult:
        """Capture frames, write an animation file, and return the result.

        Args:
            world: Runtime world object.
            output_path: Optional explicit output path.
            format: Optional output format override.
            progress_callback: Optional frame-capture progress callback.

        Returns:
            ``AnimationResult`` with output path and frame metadata.
        """

        frames = self.capture_frames(world, progress_callback=progress_callback)
        return self.write_frames(
            frames,
            output_path=output_path,
            format=format,
        )

    def write_frames(
        self,
        frames: Sequence[AnimationFrame | np.ndarray],
        *,
        output_path: str | Path | None = None,
        format: AnimationFormat | str | None = None,
    ) -> AnimationResult:
        """Write existing frames to an animation file.

        Args:
            frames: Sequence of ``AnimationFrame`` objects or RGB/RGBA arrays.
            output_path: Optional output file path.
            format: Optional output format override.

        Returns:
            ``AnimationResult``.
        """

        normalized_format = (
            normalize_animation_format(format)
            if format is not None
            else self._format_from_path_or_config(output_path)
        )
        normalized_frames = self._normalize_frames(frames)
        if not normalized_frames:
            raise AnimationWriterError("Cannot write animation with zero frames")

        path = (
            str(output_path)
            if output_path is not None
            else temporary_animation_path(
                normalized_format,
                output_dir=self.config.output_dir,
                prefix=self.config.filename_prefix,
            )
        )

        write_animation_arrays(
            [frame.image for frame in normalized_frames],
            output_path=path,
            format=normalized_format,
            fps=self.config.resolved_fps(),
            interval_ms=int(self.config.interval_ms),
            repeat=bool(self.config.repeat),
            dpi=int(self.config.dpi or self.config.renderer_config.dpi),
            title=self.config.title,
            show_title=bool(self.config.show_animation_title),
            close_figure=bool(self.config.close_animation_figure),
        )

        return AnimationResult(
            path=path,
            format=normalized_format,
            frame_count=len(normalized_frames),
            fps=self.config.resolved_fps(),
            frames=tuple(normalized_frames),
            metadata={
                **copy.deepcopy(dict(self.config.metadata)),
                "steps_per_frame": int(self.config.steps_per_frame),
                "include_initial_frame": bool(self.config.include_initial_frame),
                "advance_mode": normalize_advance_mode(self.config.advance_mode).value,
            },
        )

    def animate_to_gif(
        self,
        world: World,
        *,
        output_path: str | Path | None = None,
        progress_callback: Callable[[int, int], None] | None = None,
    ) -> AnimationResult:
        """Create a GIF animation for a world."""

        return self.animate(
            world,
            output_path=output_path,
            format=AnimationFormat.GIF,
            progress_callback=progress_callback,
        )

    def animate_to_mp4(
        self,
        world: World,
        *,
        output_path: str | Path | None = None,
        progress_callback: Callable[[int, int], None] | None = None,
    ) -> AnimationResult:
        """Create an MP4 animation for a world.

        MP4 output requires an available ffmpeg writer.
        """

        return self.animate(
            world,
            output_path=output_path,
            format=AnimationFormat.MP4,
            progress_callback=progress_callback,
        )

    def advance_world(self, world: World, *, steps: int = 1) -> None:
        """Advance the runtime world by a number of simulation steps."""

        for _ in range(max(0, int(steps))):
            self._advance_world_once(world)

    def _advance_world_once(self, world: World) -> None:
        """Advance the world by one step using the configured strategy."""

        mode = normalize_advance_mode(self.config.advance_mode)
        before_step = _world_step(world)

        advanced = False

        if mode in {AnimationAdvanceMode.AUTO, AnimationAdvanceMode.WORLD_STEP}:
            step_method = getattr(world, "step", None)
            if callable(step_method):
                _call_step_method(step_method, world)
                advanced = True

        if not advanced and mode in {AnimationAdvanceMode.AUTO, AnimationAdvanceMode.SCHEDULER_STEP}:
            scheduler = getattr(world, "scheduler", None)
            if scheduler is not None:
                scheduler_step = getattr(scheduler, "step", None)
                if callable(scheduler_step):
                    _call_step_method(scheduler_step, world)
                    advanced = True
                else:
                    advanced = _call_scheduler_phases(scheduler, world)

        if not advanced and mode is AnimationAdvanceMode.NONE:
            return

        if not advanced and mode is not AnimationAdvanceMode.AUTO:
            raise WorldAdvanceError(f"Could not advance world using mode {mode.value!r}")

        after_step = _world_step(world)
        if self.config.ensure_step_count_progress and before_step == after_step:
            _increment_world_step(world)

    def _capture_single_frame(
        self,
        world: World,
        *,
        renderer: WorldRenderer,
        index: int,
    ) -> AnimationFrame:
        """Capture a single rendered frame and metadata snapshot."""

        render_result = renderer.render_result(world)
        image = figure_to_rgb_array(render_result.figure)

        if self.config.close_frame_figures:
            close_figure(render_result.figure)

        return AnimationFrame(
            index=index,
            step=_world_step(world),
            image=image,
            snapshot=render_result.snapshot,
            metadata={
                "world_step": _world_step(world),
                "object_count": render_result.snapshot.object_count,
            },
        )

    def _normalize_frames(
        self,
        frames: Sequence[AnimationFrame | np.ndarray],
    ) -> tuple[AnimationFrame, ...]:
        """Normalize frame inputs to ``AnimationFrame`` objects."""

        normalized: list[AnimationFrame] = []

        for index, frame in enumerate(frames):
            if isinstance(frame, AnimationFrame):
                normalized.append(frame)
                continue

            array = normalize_image_array(frame)
            normalized.append(
                AnimationFrame(
                    index=index,
                    step=None,
                    image=array,
                    snapshot=None,
                    metadata={"source": "array"},
                )
            )

        validate_frame_shapes([frame.image for frame in normalized])
        return tuple(normalized)

    def _format_from_path_or_config(self, output_path: str | Path | None) -> AnimationFormat:
        """Infer animation format from output path or config."""

        if output_path is None:
            return self.config.resolved_format()

        suffix = Path(output_path).suffix.lower().lstrip(".")
        if suffix in {"gif", "mp4"}:
            return normalize_animation_format(suffix)

        return self.config.resolved_format()


def animate_world(
    world: World,
    *,
    output_path: str | Path | None = None,
    frame_count: int = 60,
    steps_per_frame: int = 1,
    format: AnimationFormat | str = AnimationFormat.GIF,
    renderer_config: RendererConfig | None = None,
    output_dir: str | Path | None = None,
) -> AnimationResult:
    """Convenience function that animates a world to GIF or MP4."""

    config = AnimationConfig(
        frame_count=frame_count,
        steps_per_frame=steps_per_frame,
        format=format,
        renderer_config=renderer_config or RendererConfig(),
        output_dir=output_dir,
    )
    return WorldAnimator(config=config).animate(world, output_path=output_path, format=format)


def animate_world_to_gif_path(
    world: World,
    *,
    output_path: str | Path | None = None,
    frame_count: int = 60,
    steps_per_frame: int = 1,
    renderer_config: RendererConfig | None = None,
    output_dir: str | Path | None = None,
) -> str:
    """Animate a world to a GIF path.

    This is the safest default for a root-level Gradio ``app.py`` because GIF
    output does not require ffmpeg.
    """

    result = animate_world(
        world,
        output_path=output_path,
        frame_count=frame_count,
        steps_per_frame=steps_per_frame,
        format=AnimationFormat.GIF,
        renderer_config=renderer_config,
        output_dir=output_dir,
    )
    return result.path


def animate_world_to_mp4_path(
    world: World,
    *,
    output_path: str | Path | None = None,
    frame_count: int = 60,
    steps_per_frame: int = 1,
    renderer_config: RendererConfig | None = None,
    output_dir: str | Path | None = None,
) -> str:
    """Animate a world to an MP4 path.

    MP4 output requires ffmpeg to be available in the runtime environment.
    """

    result = animate_world(
        world,
        output_path=output_path,
        frame_count=frame_count,
        steps_per_frame=steps_per_frame,
        format=AnimationFormat.MP4,
        renderer_config=renderer_config,
        output_dir=output_dir,
    )
    return result.path


def capture_world_frames(
    world: World,
    *,
    frame_count: int = 60,
    steps_per_frame: int = 1,
    renderer_config: RendererConfig | None = None,
) -> tuple[AnimationFrame, ...]:
    """Capture animation frames without writing a file."""

    config = AnimationConfig(
        frame_count=frame_count,
        steps_per_frame=steps_per_frame,
        renderer_config=renderer_config or RendererConfig(),
    )
    return WorldAnimator(config=config).capture_frames(world)


def capture_world_frame_arrays(
    world: World,
    *,
    frame_count: int = 60,
    steps_per_frame: int = 1,
    renderer_config: RendererConfig | None = None,
) -> tuple[np.ndarray, ...]:
    """Capture animation frames as RGB arrays without writing a file."""

    frames = capture_world_frames(
        world,
        frame_count=frame_count,
        steps_per_frame=steps_per_frame,
        renderer_config=renderer_config,
    )
    return tuple(frame.image for frame in frames)


def write_animation_arrays(
    frames: Sequence[np.ndarray],
    *,
    output_path: str | Path,
    format: AnimationFormat | str,
    fps: float,
    interval_ms: int = 150,
    repeat: bool = True,
    dpi: int = 120,
    title: str | None = None,
    show_title: bool = False,
    close_figure: bool = True,
) -> str:
    """Write image arrays to a GIF or MP4 animation.

    Args:
        frames: Sequence of RGB or RGBA image arrays.
        output_path: Destination path.
        format: Animation output format.
        fps: Frames per second.
        interval_ms: Display interval in milliseconds.
        repeat: Whether GIF animations should loop.
        dpi: Output DPI.
        title: Optional animation title.
        show_title: Whether to display the title above the image.
        close_figure: Whether to close the temporary animation figure.

    Returns:
        Output path as a string.
    """

    normalized_format = normalize_animation_format(format)
    arrays = [normalize_image_array(frame) for frame in frames]
    validate_frame_shapes(arrays)

    if not arrays:
        raise AnimationWriterError("Cannot write animation with zero frames")

    path = Path(output_path)
    path.parent.mkdir(parents=True, exist_ok=True)

    figure, axes = plt.subplots(figsize=_figure_size_from_array(arrays[0], dpi=dpi), dpi=dpi)
    image_artist = axes.imshow(arrays[0])
    axes.axis("off")

    if show_title and title:
        axes.set_title(title)

    def update(frame_index: int) -> tuple[Any, ...]:
        image_artist.set_data(arrays[frame_index])
        return (image_artist,)

    animation = FuncAnimation(
        figure,
        update,
        frames=len(arrays),
        interval=int(interval_ms),
        blit=True,
        repeat=bool(repeat),
    )

    try:
        if normalized_format is AnimationFormat.GIF:
            writer = PillowWriter(fps=max(_EPSILON, float(fps)))
            animation.save(str(path), writer=writer, dpi=int(dpi))
        elif normalized_format is AnimationFormat.MP4:
            if not writers.is_available("ffmpeg"):
                raise AnimationWriterError(
                    "MP4 output requires ffmpeg, but Matplotlib could not find an ffmpeg writer. "
                    "Use GIF output or install ffmpeg in the Space environment."
                )
            writer = FFMpegWriter(fps=max(_EPSILON, float(fps)))
            animation.save(str(path), writer=writer, dpi=int(dpi))
        else:
            raise AnimationWriterError(f"Unsupported animation format: {normalized_format!r}")
    finally:
        if close_figure:
            plt.close(figure)

    return str(path)


def temporary_animation_path(
    format: AnimationFormat | str,
    *,
    output_dir: str | Path | None = None,
    prefix: str = "worldsmithai_animation_",
) -> str:
    """Return a temporary animation path with the requested extension."""

    normalized_format = normalize_animation_format(format)
    suffix = f".{normalized_format.value}"
    directory = None if output_dir is None else str(output_dir)

    if directory is not None:
        Path(directory).mkdir(parents=True, exist_ok=True)

    with tempfile.NamedTemporaryFile(
        suffix=suffix,
        prefix=prefix,
        dir=directory,
        delete=False,
    ) as handle:
        return handle.name


def normalize_animation_format(value: AnimationFormat | str | None) -> AnimationFormat:
    """Normalize an animation format value."""

    if value is None:
        return AnimationFormat.GIF

    if isinstance(value, AnimationFormat):
        return value

    normalized = str(value).lower().lstrip(".")
    return AnimationFormat(normalized)


def normalize_advance_mode(value: AnimationAdvanceMode | str) -> AnimationAdvanceMode:
    """Normalize an advance mode value."""

    if isinstance(value, AnimationAdvanceMode):
        return value

    return AnimationAdvanceMode(str(value))


def normalize_image_array(frame: np.ndarray) -> np.ndarray:
    """Normalize an image frame to uint8 RGB format."""

    array = np.asarray(frame)

    if array.ndim != 3:
        raise AnimationWriterError(
            f"Animation frame must have shape (height, width, channels), got {array.shape}"
        )

    if array.shape[2] not in {3, 4}:
        raise AnimationWriterError(
            f"Animation frame must have 3 or 4 channels, got shape {array.shape}"
        )

    if array.shape[2] == 4:
        array = array[:, :, :3]

    if array.dtype == np.uint8:
        return array.copy()

    numeric = array.astype(float)
    if numeric.max(initial=0.0) <= 1.0 and numeric.min(initial=0.0) >= 0.0:
        numeric = numeric * 255.0

    numeric = np.nan_to_num(numeric, nan=0.0, posinf=255.0, neginf=0.0)
    numeric = np.clip(numeric, 0.0, 255.0)

    return numeric.astype(np.uint8)


def validate_frame_shapes(frames: Sequence[np.ndarray]) -> None:
    """Validate that all animation frames share one shape."""

    if not frames:
        raise AnimationWriterError("Animation requires at least one frame")

    first_shape = tuple(frames[0].shape)

    for index, frame in enumerate(frames):
        shape = tuple(frame.shape)
        if shape != first_shape:
            raise AnimationWriterError(
                "All animation frames must have the same shape. "
                f"Frame 0 has {first_shape}, frame {index} has {shape}."
            )


def _figure_size_from_array(array: np.ndarray, *, dpi: int) -> tuple[float, float]:
    """Return Matplotlib figure size in inches for an image array."""

    height, width = array.shape[:2]
    safe_dpi = max(1, int(dpi))
    return max(1.0, width / safe_dpi), max(1.0, height / safe_dpi)


def _world_step(world: World) -> int | None:
    """Return current world step if available."""

    value = getattr(world, "step_count", None)
    if isinstance(value, Real) and not isinstance(value, bool):
        return int(value)
    return None


def _increment_world_step(world: World) -> None:
    """Increment ``world.step_count`` if possible."""

    current_step = _world_step(world)
    if current_step is None:
        try:
            setattr(world, "step_count", 1)
        except Exception:
            logger.debug("Could not initialize world.step_count", exc_info=True)
        return

    try:
        setattr(world, "step_count", current_step + 1)
    except Exception:
        logger.debug("Could not increment world.step_count", exc_info=True)


def _call_step_method(method: Callable[..., Any], world: World) -> Any:
    """Call a step-like method with zero or one world argument."""

    try:
        signature = inspect.signature(method)
    except (TypeError, ValueError):
        return method()

    required_parameters = [
        parameter
        for parameter in signature.parameters.values()
        if parameter.default is inspect.Parameter.empty
        and parameter.kind
        in {
            inspect.Parameter.POSITIONAL_ONLY,
            inspect.Parameter.POSITIONAL_OR_KEYWORD,
            inspect.Parameter.KEYWORD_ONLY,
        }
    ]

    if not required_parameters:
        return method()

    return method(world)


def _call_scheduler_phases(scheduler: Any, world: World) -> bool:
    """Call common scheduler phase methods when no direct step method exists."""

    phase_names = ("process_events", "execute_agents", "update_resources")
    called_any = False

    for name in phase_names:
        method = getattr(scheduler, name, None)
        if not callable(method):
            continue

        _call_step_method(method, world)
        called_any = True

    return called_any


ANIMATION_REGISTRY: Mapping[str, type[WorldAnimator]] = MappingProxyType(
    {
        WorldAnimator.name: WorldAnimator,
    }
)


__all__ = [
    "ANIMATION_REGISTRY",
    "AnimationAdvanceMode",
    "AnimationConfig",
    "AnimationError",
    "AnimationFormat",
    "AnimationFrame",
    "AnimationResult",
    "AnimationWriterError",
    "WorldAdvanceError",
    "WorldAnimator",
    "animate_world",
    "animate_world_to_gif_path",
    "animate_world_to_mp4_path",
    "capture_world_frame_arrays",
    "capture_world_frames",
    "normalize_animation_format",
    "normalize_image_array",
    "temporary_animation_path",
    "validate_frame_shapes",
    "write_animation_arrays",
]