visualization/charts.py

"""
Generic charting utilities for WorldSmithAI.

This module produces Matplotlib charts for grouped world time series such as
population curves, resource curves, metric curves, and arbitrary DSL-defined
state curves.

It is designed to work cleanly from a root-level Hugging Face Spaces ``app.py``.
It does not import Gradio and does not assume an ``app/`` package.

Gradio examples:
    from visualization.charts import plot_population_curves

    def simulate(prompt: str):
        world_history = run_world_from_prompt(prompt)
        return plot_population_curves(world_history)

    from visualization.charts import population_curves_to_png_path

    def simulate_file(prompt: str):
        world_history = run_world_from_prompt(prompt)
        return population_curves_to_png_path(world_history)

Input flexibility:
    Chart functions accept:
        - a single runtime world,
        - a sequence of runtime world snapshots,
        - a sequence of ChartSnapshot objects,
        - a sequence of dictionaries with ``step`` and ``values``,
        - a ChartData object.

Future extensibility:
    - Add confidence bands across multiple simulation runs.
    - Add stacked area charts.
    - Add event annotations.
    - Add dashboard composition for renderer + charts.
    - Add Plotly adapters while preserving Matplotlib defaults.
    - Add chart presets for diversity, entropy, stability, and interestingness.
"""

from __future__ import annotations

import copy
import logging
import math
import tempfile
from collections.abc import Iterable, 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.axes import Axes
from matplotlib.figure import Figure

if TYPE_CHECKING:
    from core.world import World


logger = logging.getLogger(__name__)

_MISSING = object()
_EPSILON = 1.0e-12


class ChartCollection(str, Enum):
    """World object collections that can be charted."""

    AGENTS = "agents"
    RESOURCES = "resources"
    BOTH = "both"


class ChartAggregation(str, Enum):
    """Aggregation modes used to convert world objects into chart values."""

    COUNT = "count"
    SUM = "sum"
    MEAN = "mean"


class ChartKind(str, Enum):
    """Common chart presets."""

    GENERIC = "generic"
    POPULATION = "population"
    RESOURCE = "resource"
    METRIC = "metric"


@dataclass(frozen=True)
class ChartSnapshot:
    """One time point in a grouped chart.

    Attributes:
        step: Optional simulation step. If absent, chart index is used.
        values: Mapping from series name to numeric value.
        metadata: Optional JSON-compatible metadata.
    """

    step: int | float | None
    values: Mapping[str, float]
    metadata: Mapping[str, Any] = field(default_factory=dict)

    def to_dict(self) -> dict[str, Any]:
        """Return a JSON-friendly representation of the snapshot."""

        return {
            "step": self.step,
            "values": copy.deepcopy(dict(self.values)),
            "metadata": copy.deepcopy(dict(self.metadata)),
        }


@dataclass(frozen=True)
class ChartData:
    """Grouped time-series data used by chart rendering.

    Attributes:
        snapshots: Ordered time snapshots.
        title: Optional chart title.
        metadata: Optional JSON-compatible metadata.
    """

    snapshots: tuple[ChartSnapshot, ...]
    title: str | None = None
    metadata: Mapping[str, Any] = field(default_factory=dict)

    @property
    def is_empty(self) -> bool:
        """Return whether the chart has no snapshots or no numeric values."""

        return not self.snapshots or not self.group_names

    @property
    def group_names(self) -> tuple[str, ...]:
        """Return all group or series names in deterministic order."""

        names: set[str] = set()
        for snapshot in self.snapshots:
            names.update(str(name) for name in snapshot.values.keys())
        return tuple(sorted(names))

    @property
    def step_values(self) -> tuple[float, ...]:
        """Return x-axis values, using snapshot index when step is absent."""

        values: list[float] = []

        for index, snapshot in enumerate(self.snapshots):
            if snapshot.step is None or not _is_number(snapshot.step):
                values.append(float(index))
            else:
                values.append(float(snapshot.step))

        return tuple(values)

    def series(self, *, fill_missing: float = 0.0) -> Mapping[str, tuple[tuple[float, float], ...]]:
        """Return chart series as mapping from name to ``(x, y)`` points."""

        x_values = self.step_values
        output: dict[str, tuple[tuple[float, float], ...]] = {}

        for group in self.group_names:
            points: list[tuple[float, float]] = []
            for x_value, snapshot in zip(x_values, self.snapshots):
                y_value = snapshot.values.get(group, fill_missing)
                points.append((float(x_value), float(y_value)))
            output[group] = tuple(points)

        return output

    def latest_values(self) -> Mapping[str, float]:
        """Return the latest snapshot values."""

        if not self.snapshots:
            return {}
        return copy.deepcopy(dict(self.snapshots[-1].values))

    def to_dict(self) -> dict[str, Any]:
        """Return a JSON-friendly representation of chart data."""

        return {
            "title": self.title,
            "snapshot_count": len(self.snapshots),
            "groups": list(self.group_names),
            "snapshots": [snapshot.to_dict() for snapshot in self.snapshots],
            "metadata": copy.deepcopy(dict(self.metadata)),
        }


@dataclass(frozen=True)
class ChartResult:
    """Rendered chart result containing a figure and data."""

    figure: Figure
    data: ChartData
    path: str | None = None

    def to_dict(self) -> dict[str, Any]:
        """Return a JSON-friendly chart result summary.

        The Matplotlib figure itself is intentionally not serialized.
        """

        return {
            "path": self.path,
            "data": self.data.to_dict(),
        }


@dataclass
class ChartConfig:
    """Configuration for ``WorldChartRenderer``.

    Defaults are Gradio-friendly and produce standard Matplotlib figures that
    work with ``gr.Plot``.
    """

    kind: ChartKind | str = ChartKind.GENERIC

    collection: ChartCollection | str = ChartCollection.AGENTS
    group_by_path: str = "type"
    value_path: str | None = None
    weight_path: str | None = None
    aggregation: ChartAggregation | str = ChartAggregation.COUNT

    alive_only: bool = True
    include_missing: bool = False
    missing_label: str = "unknown"

    include_group_values: tuple[str, ...] = ()
    exclude_group_values: tuple[str, ...] = ()

    case_sensitive: bool = True
    strip_labels: bool = True
    include_collection_prefix: bool = False

    minimum_weight: float = 0.0

    title: str | None = None
    x_label: str = "Step"
    y_label: str = "Value"

    figure_size: tuple[float, float] = (7.0, 4.5)
    dpi: int = 120

    show_grid: bool = True
    show_legend: bool = True
    legend_location: str = "best"
    max_legend_items: int = 12

    show_markers: bool = False
    line_width: float = 1.8

    force_zero_y_min: bool = True
    y_padding_fraction: float = 0.08

    top_n_series: int | None = None
    other_series_label: str = "other"

    fill_missing_value: float = 0.0

    close_after_array: bool = True
    close_after_save: bool = True

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


@dataclass
class WorldChartRenderer:
    """Create generic charts from worlds, snapshots, or chart data.

    This class is intentionally independent from Gradio. Root-level ``app.py``
    can call ``render`` for ``gr.Plot``, ``render_to_array`` for ``gr.Image``,
    or ``save_temp_png`` for path-based outputs.
    """

    config: ChartConfig = field(default_factory=ChartConfig)

    name: ClassVar[str] = "world_chart_renderer"

    def snapshot_world(self, world: World) -> ChartSnapshot:
        """Create one chart snapshot from a runtime world."""

        collection = _normalize_collection(self.config.collection)
        aggregation = _normalize_aggregation(self.config.aggregation)
        items = _iter_world_items(world, collection)

        values, included_count, skipped_count = self._group_values(
            items,
            collection=collection,
            aggregation=aggregation,
        )

        return ChartSnapshot(
            step=_world_step(world),
            values=values,
            metadata={
                "collection": collection.value,
                "aggregation": aggregation.value,
                "included_count": included_count,
                "skipped_count": skipped_count,
            },
        )

    def build_data(self, source: Any) -> ChartData:
        """Convert flexible chart source input into ``ChartData``.

        Args:
            source: A ``ChartData``, ``ChartSnapshot``, runtime world, sequence
                of worlds, sequence of snapshots, or sequence of dictionaries.

        Returns:
            Normalized chart data.
        """

        if isinstance(source, ChartData):
            return self._prepare_data(source)

        if isinstance(source, ChartSnapshot):
            return self._prepare_data(
                ChartData(
                    snapshots=(source,),
                    title=self.config.title,
                    metadata=copy.deepcopy(dict(self.config.metadata)),
                )
            )

        if _looks_like_world(source):
            return self._prepare_data(
                ChartData(
                    snapshots=(self.snapshot_world(source),),
                    title=self.config.title,
                    metadata=copy.deepcopy(dict(self.config.metadata)),
                )
            )

        if isinstance(source, Mapping):
            return self._prepare_data(
                ChartData(
                    snapshots=(snapshot_from_mapping(source),),
                    title=self.config.title,
                    metadata=copy.deepcopy(dict(self.config.metadata)),
                )
            )

        if isinstance(source, Sequence) and not isinstance(source, (str, bytes)):
            snapshots = tuple(self._snapshot_from_item(item) for item in source)
            return self._prepare_data(
                ChartData(
                    snapshots=snapshots,
                    title=self.config.title,
                    metadata=copy.deepcopy(dict(self.config.metadata)),
                )
            )

        raise TypeError(
            "Chart source must be a World, ChartData, ChartSnapshot, mapping, "
            "or sequence of those objects"
        )

    def render(self, source: Any) -> Figure:
        """Render chart source into a Matplotlib figure.

        This return value can be used directly with ``gr.Plot``.
        """

        return self.render_result(source).figure

    def render_result(self, source: Any) -> ChartResult:
        """Render chart source and return both figure and normalized data."""

        data = self.build_data(source)
        figure, axes = plt.subplots(
            figsize=self.config.figure_size,
            dpi=int(self.config.dpi),
            constrained_layout=True,
        )

        self._draw_data(data, axes)

        return ChartResult(figure=figure, data=data)

    def render_to_array(self, source: Any) -> np.ndarray:
        """Render chart source into an RGB NumPy array.

        This return value can be used directly with ``gr.Image``.
        """

        result = self.render_result(source)
        array = figure_to_rgb_array(result.figure)

        if self.config.close_after_array:
            plt.close(result.figure)

        return array

    def save(self, source: Any, path: str | Path) -> str:
        """Render chart source to an image file path."""

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

        result = self.render_result(source)
        result.figure.savefig(output_path, dpi=int(self.config.dpi), bbox_inches="tight")

        if self.config.close_after_save:
            plt.close(result.figure)

        return str(output_path)

    def save_temp_png(
        self,
        source: Any,
        *,
        output_dir: str | Path | None = None,
        prefix: str = "worldsmithai_chart_",
    ) -> str:
        """Render chart source to a temporary PNG path."""

        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=".png",
            prefix=prefix,
            dir=directory,
            delete=False,
        ) as handle:
            path = handle.name

        return self.save(source, path)

    def _snapshot_from_item(self, item: Any) -> ChartSnapshot:
        """Convert one source item into a chart snapshot."""

        if isinstance(item, ChartSnapshot):
            return item

        if _looks_like_world(item):
            return self.snapshot_world(item)

        if isinstance(item, Mapping):
            return snapshot_from_mapping(item)

        raise TypeError(f"Unsupported chart source item: {item.__class__.__name__}")

    def _prepare_data(self, data: ChartData) -> ChartData:
        """Apply configured series limiting and metadata normalization."""

        prepared = limit_chart_series(
            data,
            top_n=self.config.top_n_series,
            other_label=self.config.other_series_label,
        )
        return ChartData(
            snapshots=prepared.snapshots,
            title=prepared.title or self.config.title,
            metadata={
                **copy.deepcopy(dict(prepared.metadata)),
                **copy.deepcopy(dict(self.config.metadata)),
                "kind": _normalize_kind(self.config.kind).value,
            },
        )

    def _draw_data(self, data: ChartData, axes: Axes) -> None:
        """Draw normalized chart data onto axes."""

        if data.is_empty:
            axes.text(
                0.5,
                0.5,
                "No chart data available",
                transform=axes.transAxes,
                ha="center",
                va="center",
            )
            self._style_axes(axes, data)
            return

        series = data.series(fill_missing=float(self.config.fill_missing_value))

        for index, (series_name, points) in enumerate(series.items()):
            x_values = [point[0] for point in points]
            y_values = [point[1] for point in points]

            label = (
                series_name
                if self.config.show_legend and index < int(self.config.max_legend_items)
                else "_nolegend_"
            )

            axes.plot(
                x_values,
                y_values,
                marker="o" if self.config.show_markers else None,
                linewidth=float(self.config.line_width),
                label=label,
            )

        self._style_axes(axes, data)

    def _style_axes(self, axes: Axes, data: ChartData) -> None:
        """Apply generic chart styling."""

        title = self.config.title or data.title
        if title:
            axes.set_title(title)

        axes.set_xlabel(self.config.x_label)
        axes.set_ylabel(self.config.y_label)

        if self.config.show_grid:
            axes.grid(True, linewidth=0.5, alpha=0.35)

        if self.config.force_zero_y_min:
            lower, upper = axes.get_ylim()
            if lower > 0.0:
                axes.set_ylim(bottom=0.0, top=upper)

        self._apply_y_padding(axes)

        if self.config.show_legend and data.group_names:
            axes.legend(loc=self.config.legend_location, fontsize="small")

    def _apply_y_padding(self, axes: Axes) -> None:
        """Apply lightweight y-axis padding."""

        lower, upper = axes.get_ylim()
        span = upper - lower

        if span <= _EPSILON:
            span = 1.0

        padding = span * max(0.0, float(self.config.y_padding_fraction))
        axes.set_ylim(lower - padding, upper + padding)

    def _group_values(
        self,
        items: Sequence[Any],
        *,
        collection: ChartCollection,
        aggregation: ChartAggregation,
    ) -> tuple[dict[str, float], int, int]:
        """Aggregate world objects into grouped chart values."""

        group_sums: dict[str, float] = {}
        group_denominators: dict[str, float] = {}
        included_count = 0
        skipped_count = 0

        include_values = {
            _normalize_label(value, case_sensitive=self.config.case_sensitive, strip=True)
            for value in self.config.include_group_values
        }
        exclude_values = {
            _normalize_label(value, case_sensitive=self.config.case_sensitive, strip=True)
            for value in self.config.exclude_group_values
        }

        for item in items:
            if self.config.alive_only and _is_agent_like(item) and not _is_alive(item):
                skipped_count += 1
                continue

            label = self._label_for(item, collection)
            if label is None:
                skipped_count += 1
                continue

            normalized_label = _normalize_label(
                label,
                case_sensitive=self.config.case_sensitive,
                strip=self.config.strip_labels,
            )

            if include_values and normalized_label not in include_values:
                skipped_count += 1
                continue

            if exclude_values and normalized_label in exclude_values:
                skipped_count += 1
                continue

            weight = self._weight_for(item)
            if weight <= max(0.0, float(self.config.minimum_weight)):
                skipped_count += 1
                continue

            value = self._value_for(item, aggregation)
            if value is None:
                skipped_count += 1
                continue

            if aggregation is ChartAggregation.COUNT:
                contribution = weight
                denominator = 1.0
            elif aggregation is ChartAggregation.SUM:
                contribution = value * weight
                denominator = 1.0
            else:
                contribution = value * weight
                denominator = weight

            group_sums[normalized_label] = group_sums.get(normalized_label, 0.0) + contribution
            group_denominators[normalized_label] = (
                group_denominators.get(normalized_label, 0.0) + denominator
            )
            included_count += 1

        if aggregation is ChartAggregation.MEAN:
            values = {
                label: group_sums[label] / max(group_denominators.get(label, 0.0), _EPSILON)
                for label in group_sums
            }
        else:
            values = group_sums

        return dict(sorted(values.items(), key=lambda pair: pair[0])), included_count, skipped_count

    def _label_for(self, item: Any, collection: ChartCollection) -> str | None:
        """Return group label for an item."""

        raw_label = _read_path(item, self.config.group_by_path, _MISSING)

        if raw_label is _MISSING or raw_label is None or str(raw_label) == "":
            if not self.config.include_missing:
                return None
            raw_label = self.config.missing_label

        label = _stable_label(raw_label)

        if self.config.include_collection_prefix and collection is ChartCollection.BOTH:
            return f"{_item_collection_name(item)}:{label}"

        return label

    def _weight_for(self, item: Any) -> float:
        """Return item weight."""

        if self.config.weight_path is None:
            return 1.0

        raw_weight = _read_path(item, self.config.weight_path, 0.0)
        if not _is_number(raw_weight):
            return 0.0

        return max(0.0, float(raw_weight))

    def _value_for(self, item: Any, aggregation: ChartAggregation) -> float | None:
        """Return item numeric value for aggregation."""

        if aggregation is ChartAggregation.COUNT:
            return 1.0

        if self.config.value_path is None:
            return 1.0

        raw_value = _read_path(item, self.config.value_path, _MISSING)
        if not _is_number(raw_value):
            return None

        return float(raw_value)


@dataclass
class ChartTracker:
    """Track chart snapshots over simulation time.

    This helper is useful in Gradio callbacks that run a world step-by-step and
    want to chart accumulated data afterward.
    """

    renderer: WorldChartRenderer = field(default_factory=WorldChartRenderer)
    max_history: int = 1000
    snapshots: list[ChartSnapshot] = field(default_factory=list)

    def update(self, world: World) -> ChartSnapshot:
        """Capture the current world as a chart snapshot and append it."""

        snapshot = self.renderer.snapshot_world(world)
        _append_bounded(self.snapshots, snapshot, self.max_history)
        return snapshot

    def data(self) -> ChartData:
        """Return tracked snapshots as ``ChartData``."""

        return ChartData(
            snapshots=tuple(self.snapshots),
            title=self.renderer.config.title,
            metadata=copy.deepcopy(dict(self.renderer.config.metadata)),
        )

    def render(self) -> Figure:
        """Render tracked snapshots into a Matplotlib figure."""

        return self.renderer.render(self.data())

    def render_to_array(self) -> np.ndarray:
        """Render tracked snapshots into an RGB NumPy array."""

        return self.renderer.render_to_array(self.data())

    def save_temp_png(self, *, output_dir: str | Path | None = None) -> str:
        """Render tracked snapshots to a temporary PNG path."""

        return self.renderer.save_temp_png(self.data(), output_dir=output_dir)

    def latest(self) -> ChartSnapshot | None:
        """Return latest tracked snapshot, if any."""

        if not self.snapshots:
            return None
        return self.snapshots[-1]

    def reset(self) -> None:
        """Clear tracked snapshots."""

        self.snapshots.clear()

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

        return {
            "max_history": self.max_history,
            "data": self.data().to_dict(),
        }


def population_chart_config(**overrides: Any) -> ChartConfig:
    """Return a default config for population curves."""

    config = ChartConfig(
        kind=ChartKind.POPULATION,
        collection=ChartCollection.AGENTS,
        group_by_path="type",
        value_path=None,
        weight_path=None,
        aggregation=ChartAggregation.COUNT,
        alive_only=True,
        title="Population by Agent Type",
        x_label="Step",
        y_label="Population",
    )
    return _replace_config(config, overrides)


def resource_chart_config(*, weight_by_amount: bool = True, **overrides: Any) -> ChartConfig:
    """Return a default config for resource curves."""

    config = ChartConfig(
        kind=ChartKind.RESOURCE,
        collection=ChartCollection.RESOURCES,
        group_by_path="type",
        value_path="amount" if weight_by_amount else None,
        weight_path=None,
        aggregation=ChartAggregation.SUM if weight_by_amount else ChartAggregation.COUNT,
        alive_only=False,
        title="Resources by Type",
        x_label="Step",
        y_label="Resource Amount" if weight_by_amount else "Resource Count",
    )
    return _replace_config(config, overrides)


def metric_chart_config(**overrides: Any) -> ChartConfig:
    """Return a default config for precomputed metric curves."""

    config = ChartConfig(
        kind=ChartKind.METRIC,
        collection=ChartCollection.AGENTS,
        group_by_path="type",
        aggregation=ChartAggregation.COUNT,
        title="Metric Curve",
        x_label="Step",
        y_label="Metric Value",
    )
    return _replace_config(config, overrides)


def plot_population_curves(source: Any, *, config: ChartConfig | None = None) -> Figure:
    """Plot population curves from worlds, snapshots, mappings, or chart data."""

    renderer = WorldChartRenderer(config=config or population_chart_config())
    return renderer.render(source)


def plot_resource_curves(
    source: Any,
    *,
    config: ChartConfig | None = None,
    weight_by_amount: bool = True,
) -> Figure:
    """Plot resource curves from worlds, snapshots, mappings, or chart data."""

    renderer = WorldChartRenderer(config=config or resource_chart_config(weight_by_amount=weight_by_amount))
    return renderer.render(source)


def plot_metric_curves(source: Any, *, config: ChartConfig | None = None) -> Figure:
    """Plot generic metric curves from precomputed snapshots or mappings."""

    renderer = WorldChartRenderer(config=config or metric_chart_config())
    return renderer.render(source)


def population_curves_to_array(source: Any, *, config: ChartConfig | None = None) -> np.ndarray:
    """Render population curves to an RGB NumPy array for ``gr.Image``."""

    renderer = WorldChartRenderer(config=config or population_chart_config())
    return renderer.render_to_array(source)


def resource_curves_to_array(
    source: Any,
    *,
    config: ChartConfig | None = None,
    weight_by_amount: bool = True,
) -> np.ndarray:
    """Render resource curves to an RGB NumPy array for ``gr.Image``."""

    renderer = WorldChartRenderer(config=config or resource_chart_config(weight_by_amount=weight_by_amount))
    return renderer.render_to_array(source)


def population_curves_to_png_path(
    source: Any,
    *,
    path: str | Path | None = None,
    output_dir: str | Path | None = None,
    config: ChartConfig | None = None,
) -> str:
    """Render population curves to a PNG path."""

    renderer = WorldChartRenderer(config=config or population_chart_config())
    if path is not None:
        return renderer.save(source, path)
    return renderer.save_temp_png(source, output_dir=output_dir)


def resource_curves_to_png_path(
    source: Any,
    *,
    path: str | Path | None = None,
    output_dir: str | Path | None = None,
    config: ChartConfig | None = None,
    weight_by_amount: bool = True,
) -> str:
    """Render resource curves to a PNG path."""

    renderer = WorldChartRenderer(config=config or resource_chart_config(weight_by_amount=weight_by_amount))
    if path is not None:
        return renderer.save(source, path)
    return renderer.save_temp_png(source, output_dir=output_dir)


def chart_data_from_snapshots(snapshots: Sequence[ChartSnapshot | Mapping[str, Any]]) -> ChartData:
    """Build ``ChartData`` from chart snapshots or mapping snapshots."""

    normalized = tuple(
        snapshot if isinstance(snapshot, ChartSnapshot) else snapshot_from_mapping(snapshot)
        for snapshot in snapshots
    )
    return ChartData(snapshots=normalized)


def snapshot_from_mapping(data: Mapping[str, Any]) -> ChartSnapshot:
    """Convert a mapping into ``ChartSnapshot``.

    Supported shapes:
        ``{"step": 1, "values": {"a": 2}}``
        ``{"step": 1, "group_values": {"a": 2}}``
        ``{"step": 1, "counts": {"a": 2}}``
        ``{"step": 1, "a": 2, "b": 3}``
    """

    step = _mapping_step(data)
    values = _mapping_values(data)
    metadata = data.get("metadata", {})

    return ChartSnapshot(
        step=step,
        values=values,
        metadata=copy.deepcopy(dict(metadata)) if isinstance(metadata, Mapping) else {},
    )


def limit_chart_series(
    data: ChartData,
    *,
    top_n: int | None,
    other_label: str = "other",
) -> ChartData:
    """Limit chart to top-N series and combine remaining groups into ``other``."""

    if top_n is None or top_n <= 0:
        return data

    group_totals: dict[str, float] = {}
    for snapshot in data.snapshots:
        for group, value in snapshot.values.items():
            group_totals[group] = group_totals.get(group, 0.0) + abs(float(value))

    top_groups = {
        group
        for group, _ in sorted(
            group_totals.items(),
            key=lambda pair: (-pair[1], pair[0]),
        )[: int(top_n)]
    }

    new_snapshots: list[ChartSnapshot] = []
    for snapshot in data.snapshots:
        values: dict[str, float] = {}
        other_total = 0.0

        for group, value in snapshot.values.items():
            if group in top_groups:
                values[group] = float(value)
            else:
                other_total += float(value)

        if other_total != 0.0:
            values[other_label] = other_total

        new_snapshots.append(
            ChartSnapshot(
                step=snapshot.step,
                values=dict(sorted(values.items(), key=lambda pair: pair[0])),
                metadata=copy.deepcopy(dict(snapshot.metadata)),
            )
        )

    return ChartData(
        snapshots=tuple(new_snapshots),
        title=data.title,
        metadata=copy.deepcopy(dict(data.metadata)),
    )


def figure_to_rgb_array(figure: Figure) -> np.ndarray:
    """Convert a Matplotlib figure into an RGB NumPy array."""

    figure.canvas.draw()
    width, height = figure.canvas.get_width_height()

    try:
        buffer = figure.canvas.buffer_rgba()
        image = np.frombuffer(buffer, dtype=np.uint8).reshape(height, width, 4)
        return image[:, :, :3].copy()
    except AttributeError:
        raw = figure.canvas.tostring_rgb()
        image = np.frombuffer(raw, dtype=np.uint8).reshape(height, width, 3)
        return image.copy()


def close_figure(figure: Figure) -> None:
    """Close a Matplotlib figure to release memory."""

    plt.close(figure)


def _replace_config(config: ChartConfig, overrides: Mapping[str, Any]) -> ChartConfig:
    """Return a copy of a chart config with field overrides applied."""

    if not overrides:
        return config

    data = {
        field_name: copy.deepcopy(getattr(config, field_name))
        for field_name in config.__dataclass_fields__
    }
    data.update(dict(overrides))
    return ChartConfig(**data)


def _normalize_collection(value: ChartCollection | str) -> ChartCollection:
    """Normalize a chart collection value."""

    if isinstance(value, ChartCollection):
        return value
    return ChartCollection(str(value))


def _normalize_aggregation(value: ChartAggregation | str) -> ChartAggregation:
    """Normalize a chart aggregation value."""

    if isinstance(value, ChartAggregation):
        return value
    return ChartAggregation(str(value))


def _normalize_kind(value: ChartKind | str) -> ChartKind:
    """Normalize a chart kind value."""

    if isinstance(value, ChartKind):
        return value
    return ChartKind(str(value))


def _looks_like_world(value: Any) -> bool:
    """Return whether a value looks like a runtime world object."""

    return hasattr(value, "agents") or hasattr(value, "resources") or hasattr(value, "step_count")


def _world_step(world: World) -> int | None:
    """Return the 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 _is_number(value: Any) -> bool:
    """Return whether a value is a real numeric scalar, excluding booleans."""

    return isinstance(value, (Real, np.integer, np.floating)) and not isinstance(value, bool)


def _is_alive(item: Any) -> bool:
    """Return whether an item is alive when it exposes an ``alive`` field."""

    return bool(getattr(item, "alive", True))


def _is_agent_like(item: Any) -> bool:
    """Return whether an item looks like an agent."""

    return hasattr(item, "alive") or hasattr(item, "behaviors") or hasattr(item, "policy")


def _item_collection_name(item: Any) -> str:
    """Return a best-effort collection name for an item."""

    if _is_agent_like(item):
        return ChartCollection.AGENTS.value

    if hasattr(item, "amount"):
        return ChartCollection.RESOURCES.value

    return "items"


def _iter_world_items(world: World, collection: ChartCollection) -> tuple[Any, ...]:
    """Return world items for a configured chart collection."""

    if collection is ChartCollection.AGENTS:
        return _iter_collection(getattr(world, "agents", ()))

    if collection is ChartCollection.RESOURCES:
        return _iter_collection(getattr(world, "resources", ()))

    agents = _iter_collection(getattr(world, "agents", ()))
    resources = _iter_collection(getattr(world, "resources", ()))
    return agents + resources


def _iter_collection(raw_collection: Any) -> tuple[Any, ...]:
    """Return items from a mapping-backed or sequence-backed collection."""

    if raw_collection is None:
        return ()

    if isinstance(raw_collection, Mapping):
        values = raw_collection.values()
    elif isinstance(raw_collection, Iterable) and not isinstance(raw_collection, (str, bytes)):
        values = raw_collection
    else:
        values = (raw_collection,)

    return tuple(item for item in values if item is not None)


def _mapping_step(data: Mapping[str, Any]) -> int | float | None:
    """Extract a step value from a mapping snapshot."""

    for key in ("step", "world_step", "time", "t", "index"):
        value = data.get(key)
        if _is_number(value):
            numeric = float(value)
            return int(numeric) if numeric.is_integer() else numeric

    return None


def _mapping_values(data: Mapping[str, Any]) -> dict[str, float]:
    """Extract chart values from a mapping snapshot."""

    for key in ("values", "group_values", "counts", "probabilities", "component_scores"):
        raw_values = data.get(key)
        if isinstance(raw_values, Mapping):
            return {
                str(group): float(value)
                for group, value in sorted(raw_values.items(), key=lambda pair: str(pair[0]))
                if _is_number(value)
            }

    reserved = {
        "step",
        "world_step",
        "time",
        "t",
        "index",
        "metadata",
        "title",
        "name",
        "metric_name",
    }

    return {
        str(key): float(value)
        for key, value in sorted(data.items(), key=lambda pair: str(pair[0]))
        if key not in reserved and _is_number(value)
    }


def _split_path(path: str) -> tuple[str, ...]:
    """Split a dot-separated path into components."""

    return tuple(part for part in str(path).split(".") if part)


def _get_mapping_path(container: Mapping[str, Any], path: str, default: Any = _MISSING) -> Any:
    """Read a nested mapping value using dot notation."""

    parts = _split_path(path)
    if not parts:
        return default

    current: Any = container
    for part in parts:
        if not isinstance(current, Mapping) or part not in current:
            return default
        current = current[part]

    return current


def _get_object_path(root: Any, path: str, default: Any = _MISSING) -> Any:
    """Read nested values from mappings, sequences, or object attributes."""

    parts = _split_path(path)
    if not parts:
        return root

    current: Any = root
    for part in parts:
        if isinstance(current, Mapping):
            if part not in current:
                return default
            current = current[part]
            continue

        if isinstance(current, Sequence) and not isinstance(current, (str, bytes)) and part.isdigit():
            index = int(part)
            if index >= len(current):
                return default
            current = current[index]
            continue

        if not hasattr(current, part):
            return default

        current = getattr(current, part)

    return current


def _read_path(item: Any, path: str, default: Any = _MISSING) -> Any:
    """Read a generic path from an object or mapping.

    Supported prefixes:
        - ``state.foo``
        - ``state:foo``
        - ``memory.foo``
        - ``memory:foo``
        - ``metadata.foo``
        - ``metadata:foo``

    Without a prefix, object attributes and mapping keys are read directly.
    """

    normalized_path = str(path)

    if normalized_path.startswith("state."):
        state = getattr(item, "state", {})
        return _get_mapping_path(
            state if isinstance(state, Mapping) else {},
            normalized_path.removeprefix("state."),
            default,
        )

    if normalized_path.startswith("state:"):
        state = getattr(item, "state", {})
        return _get_mapping_path(
            state if isinstance(state, Mapping) else {},
            normalized_path.removeprefix("state:"),
            default,
        )

    if normalized_path.startswith("memory."):
        memory = getattr(item, "memory", {})
        return _get_mapping_path(
            memory if isinstance(memory, Mapping) else {},
            normalized_path.removeprefix("memory."),
            default,
        )

    if normalized_path.startswith("memory:"):
        memory = getattr(item, "memory", {})
        return _get_mapping_path(
            memory if isinstance(memory, Mapping) else {},
            normalized_path.removeprefix("memory:"),
            default,
        )

    if normalized_path.startswith("metadata."):
        metadata = getattr(item, "metadata", {})
        return _get_mapping_path(
            metadata if isinstance(metadata, Mapping) else {},
            normalized_path.removeprefix("metadata."),
            default,
        )

    if normalized_path.startswith("metadata:"):
        metadata = getattr(item, "metadata", {})
        return _get_mapping_path(
            metadata if isinstance(metadata, Mapping) else {},
            normalized_path.removeprefix("metadata:"),
            default,
        )

    return _get_object_path(item, normalized_path, default)


def _stable_label(value: Any) -> str:
    """Return a stable string label for arbitrary values."""

    if value is None:
        return "none"

    if isinstance(value, bool):
        return "true" if value else "false"

    if _is_number(value):
        numeric_value = float(value)
        if numeric_value.is_integer():
            return str(int(numeric_value))
        return str(numeric_value)

    if isinstance(value, Mapping):
        parts = [f"{key}={_stable_label(value[key])}" for key in sorted(value.keys(), key=str)]
        return "{" + ",".join(parts) + "}"

    if isinstance(value, Sequence) and not isinstance(value, (str, bytes)):
        return "[" + ",".join(_stable_label(item) for item in value) + "]"

    return str(value)


def _normalize_label(value: Any, *, case_sensitive: bool, strip: bool) -> str:
    """Normalize a label for deterministic grouping."""

    label = _stable_label(value)
    if strip:
        label = label.strip()
    if not case_sensitive:
        label = label.lower()
    return label


def _append_bounded(items: MutableSequence[Any], value: Any, max_items: int) -> None:
    """Append an item while enforcing a maximum history length."""

    items.append(value)
    if max_items > 0 and len(items) > max_items:
        del items[: len(items) - max_items]


CHART_REGISTRY: Mapping[str, type[WorldChartRenderer]] = MappingProxyType(
    {
        WorldChartRenderer.name: WorldChartRenderer,
    }
)


__all__ = [
    "CHART_REGISTRY",
    "ChartAggregation",
    "ChartCollection",
    "ChartConfig",
    "ChartData",
    "ChartKind",
    "ChartResult",
    "ChartSnapshot",
    "ChartTracker",
    "WorldChartRenderer",
    "chart_data_from_snapshots",
    "close_figure",
    "figure_to_rgb_array",
    "limit_chart_series",
    "metric_chart_config",
    "plot_metric_curves",
    "plot_population_curves",
    "plot_resource_curves",
    "population_chart_config",
    "population_curves_to_array",
    "population_curves_to_png_path",
    "resource_chart_config",
    "resource_curves_to_array",
    "resource_curves_to_png_path",
    "snapshot_from_mapping",
]