analysis/trajectory_extractor.py

"""
analysis.trajectory_extractor
=============================

Convert raw WorldSmithAI simulation history into structured trajectory data.

This module is intentionally placed above the simulation engine. It consumes
world history, metric logs, and event logs without mutating runtime agents,
resources, behaviors, policies, schedulers, or worlds.

The extractor is designed to be tolerant of several history formats because the
core engine may evolve independently. It accepts either:

    - a world-like object with ``history``
    - a raw sequence of snapshot mappings
    - a single snapshot mapping

The output is a ``TrajectoryData`` dataclass consumed by later analysis modules.

Example
-------

    extractor = TrajectoryExtractor()
    trajectory = extractor.extract(world)

    print(trajectory.step_count)
    print(trajectory.agent_type_counts)
    print(trajectory.population_timeseries)

Architecture
------------

    world.history
        ↓
    TrajectoryExtractor.extract(...)
        ↓
    TrajectoryData
        ↓
    simulation_analyzer / report_generator / scenario comparator
"""

from __future__ import annotations

import copy
import logging
import math
from collections import defaultdict
from collections.abc import Mapping, Sequence
from dataclasses import dataclass, field
from typing import Any, Protocol


logger = logging.getLogger(__name__)


class WorldHistoryProtocol(Protocol):
    """Structural protocol for world-like objects with history."""

    history: Any
    step_count: int


@dataclass(frozen=True)
class TrajectoryExtractionConfig:
    """Configuration for trajectory extraction.

    Attributes:
        count_alive_only:
            If true, population counts exclude agents with ``alive=False``.
        include_snapshot_events:
            If true, event-like records inside each snapshot are added to the
            normalized event log.
        include_snapshot_metrics:
            If true, metric-like records inside each snapshot are added to the
            normalized metric history.
        preserve_unknown_snapshot_fields:
            If true, selected unrecognized top-level snapshot keys are preserved
            in metadata for debugging.
        max_preserved_unknown_values:
            Maximum number of unknown snapshot values preserved in metadata.
    """

    count_alive_only: bool = True
    include_snapshot_events: bool = True
    include_snapshot_metrics: bool = True
    preserve_unknown_snapshot_fields: bool = False
    max_preserved_unknown_values: int = 20


@dataclass(frozen=True)
class TrajectoryData:
    """Structured simulation trajectory data.

    Attributes:
        population_timeseries:
            Mapping from step to counts by agent type.
        resource_timeseries:
            Mapping from step to total resource amount by resource type.
        event_log:
            Normalized event records. Each record should include a ``step`` key
            when the step is known.
        metric_history:
            Mapping from metric name to a list of per-step metric records.
        agent_type_counts:
            Final observed agent counts by type.
        step_count:
            Number of simulated steps represented by the trajectory. This is
            the highest observed step plus one when possible.
        first_step:
            First observed step, or ``None`` when there is no history.
        last_step:
            Last observed step, or ``None`` when there is no history.
        snapshot_count:
            Number of snapshots consumed.
        metadata:
            Additional extraction diagnostics.
    """

    population_timeseries: Mapping[int, Mapping[str, int]] = field(default_factory=dict)
    resource_timeseries: Mapping[int, Mapping[str, float]] = field(default_factory=dict)
    event_log: tuple[Mapping[str, Any], ...] = field(default_factory=tuple)
    metric_history: Mapping[str, tuple[Mapping[str, Any], ...]] = field(default_factory=dict)
    agent_type_counts: Mapping[str, int] = field(default_factory=dict)
    step_count: int = 0
    first_step: int | None = None
    last_step: int | None = None
    snapshot_count: int = 0
    metadata: Mapping[str, Any] = field(default_factory=dict)

    @property
    def is_empty(self) -> bool:
        """Return whether this trajectory contains no snapshots."""

        return self.snapshot_count == 0

    @property
    def final_population(self) -> Mapping[str, int]:
        """Return the final population count by agent type."""

        if self.last_step is None:
            return {}

        return dict(self.population_timeseries.get(self.last_step, {}))

    @property
    def final_resources(self) -> Mapping[str, float]:
        """Return the final resource totals by resource type."""

        if self.last_step is None:
            return {}

        return dict(self.resource_timeseries.get(self.last_step, {}))

    def population_delta(self) -> dict[str, int]:
        """Return final-minus-initial population delta by agent type."""

        if self.first_step is None or self.last_step is None:
            return {}

        initial = self.population_timeseries.get(self.first_step, {})
        final = self.population_timeseries.get(self.last_step, {})
        keys = set(initial.keys()) | set(final.keys())

        return {
            key: int(final.get(key, 0)) - int(initial.get(key, 0))
            for key in sorted(keys)
        }

    def resource_delta(self) -> dict[str, float]:
        """Return final-minus-initial resource amount delta by resource type."""

        if self.first_step is None or self.last_step is None:
            return {}

        initial = self.resource_timeseries.get(self.first_step, {})
        final = self.resource_timeseries.get(self.last_step, {})
        keys = set(initial.keys()) | set(final.keys())

        return {
            key: float(final.get(key, 0.0)) - float(initial.get(key, 0.0))
            for key in sorted(keys)
        }

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

        return {
            "population_timeseries": {
                str(step): dict(counts)
                for step, counts in sorted(self.population_timeseries.items())
            },
            "resource_timeseries": {
                str(step): dict(amounts)
                for step, amounts in sorted(self.resource_timeseries.items())
            },
            "event_log": [dict(event) for event in self.event_log],
            "metric_history": {
                metric_name: [dict(record) for record in records]
                for metric_name, records in self.metric_history.items()
            },
            "agent_type_counts": dict(self.agent_type_counts),
            "step_count": self.step_count,
            "first_step": self.first_step,
            "last_step": self.last_step,
            "snapshot_count": self.snapshot_count,
            "population_delta": self.population_delta(),
            "resource_delta": self.resource_delta(),
            "metadata": copy.deepcopy(dict(self.metadata)),
        }


@dataclass
class TrajectoryExtractor:
    """Extract structured trajectory data from completed simulations."""

    config: TrajectoryExtractionConfig = field(default_factory=TrajectoryExtractionConfig)

    def extract(
        self,
        world_or_history: Any,
        *,
        metric_history: Mapping[str, Any] | Sequence[Any] | None = None,
        event_log: Sequence[Any] | Mapping[str, Any] | None = None,
    ) -> TrajectoryData:
        """Extract trajectory data from a world-like object or raw history.

        Args:
            world_or_history:
                Either a world-like object with ``history`` or a raw history
                object.
            metric_history:
                Optional external metric history to merge with snapshot metrics.
            event_log:
                Optional external event log to merge with snapshot events.

        Returns:
            Normalized ``TrajectoryData``.
        """

        world = self._as_world_like(world_or_history)
        raw_history = self._history_from_input(world_or_history)
        snapshots = self._normalize_history(raw_history)

        if not snapshots and world is not None:
            current_snapshot = self._snapshot_from_world(world)
            if current_snapshot:
                snapshots = (current_snapshot,)

        population_timeseries: dict[int, dict[str, int]] = {}
        resource_timeseries: dict[int, dict[str, float]] = {}
        normalized_event_log: list[dict[str, Any]] = []
        normalized_metric_history: dict[str, list[dict[str, Any]]] = defaultdict(list)
        unknown_snapshot_metadata: list[dict[str, Any]] = []

        for index, raw_snapshot in enumerate(snapshots):
            snapshot = self._mapping_from_any(raw_snapshot)
            step = self._step_from_snapshot(snapshot, fallback=index)

            population_counts = self._population_counts_from_snapshot(snapshot)
            resource_amounts = self._resource_amounts_from_snapshot(snapshot)

            population_timeseries[step] = population_counts
            resource_timeseries[step] = resource_amounts

            if self.config.include_snapshot_events:
                self._append_snapshot_events(
                    snapshot=snapshot,
                    step=step,
                    output=normalized_event_log,
                )

            if self.config.include_snapshot_metrics:
                self._append_snapshot_metrics(
                    snapshot=snapshot,
                    step=step,
                    output=normalized_metric_history,
                )

            if self.config.preserve_unknown_snapshot_fields:
                preserved = self._preserve_unknown_snapshot_fields(snapshot, step=step)
                if preserved:
                    unknown_snapshot_metadata.append(preserved)

        self._append_external_events(
            event_log=event_log,
            world=world,
            output=normalized_event_log,
        )

        self._append_external_metrics(
            metric_history=metric_history,
            world=world,
            output=normalized_metric_history,
        )

        observed_steps = sorted(population_timeseries.keys() | resource_timeseries.keys())
        first_step = observed_steps[0] if observed_steps else None
        last_step = observed_steps[-1] if observed_steps else None

        if last_step is None:
            step_count = int(getattr(world, "step_count", 0)) if world is not None else 0
        else:
            step_count = max(last_step + 1, int(getattr(world, "step_count", 0)) if world is not None else 0)

        final_counts = (
            population_timeseries.get(last_step, {})
            if last_step is not None
            else {}
        )

        metadata: dict[str, Any] = {
            "source": "world" if world is not None else "history",
            "history_length": len(snapshots),
            "count_alive_only": self.config.count_alive_only,
        }

        if unknown_snapshot_metadata:
            metadata["unknown_snapshot_fields"] = unknown_snapshot_metadata[
                : self.config.max_preserved_unknown_values
            ]

        return TrajectoryData(
            population_timeseries={
                step: dict(counts)
                for step, counts in sorted(population_timeseries.items())
            },
            resource_timeseries={
                step: dict(amounts)
                for step, amounts in sorted(resource_timeseries.items())
            },
            event_log=tuple(normalized_event_log),
            metric_history={
                metric_name: tuple(records)
                for metric_name, records in sorted(normalized_metric_history.items())
            },
            agent_type_counts=dict(final_counts),
            step_count=step_count,
            first_step=first_step,
            last_step=last_step,
            snapshot_count=len(snapshots),
            metadata=metadata,
        )

    def extract_from_world(self, world: WorldHistoryProtocol) -> TrajectoryData:
        """Extract trajectory data directly from a world-like object."""

        return self.extract(world)

    def extract_from_history(self, history: Sequence[Any] | Mapping[str, Any]) -> TrajectoryData:
        """Extract trajectory data directly from a raw history object."""

        return self.extract(history)

    @staticmethod
    def _as_world_like(value: Any) -> Any | None:
        """Return value when it appears to be a world-like object."""

        if hasattr(value, "history"):
            return value

        return None

    @staticmethod
    def _history_from_input(value: Any) -> Any:
        """Return raw history from a world-like object or raw history input."""

        if hasattr(value, "history"):
            return getattr(value, "history")

        return value

    def _normalize_history(self, raw_history: Any) -> tuple[Any, ...]:
        """Normalize raw history into a tuple of snapshots."""

        if raw_history is None:
            return ()

        if isinstance(raw_history, Mapping):
            if "history" in raw_history:
                return self._normalize_history(raw_history["history"])

            if "snapshots" in raw_history:
                return self._normalize_history(raw_history["snapshots"])

            return (raw_history,)

        if isinstance(raw_history, Sequence) and not isinstance(raw_history, (str, bytes)):
            return tuple(raw_history)

        logger.warning(
            "TrajectoryExtractor received unsupported history type: %s",
            raw_history.__class__.__name__,
        )
        return ()

    def _snapshot_from_world(self, world: Any) -> dict[str, Any]:
        """Build a single best-effort snapshot from a world-like object."""

        snapshot_method = getattr(world, "snapshot", None)
        if callable(snapshot_method):
            try:
                snapshot = snapshot_method()
                if isinstance(snapshot, Mapping):
                    return dict(snapshot)
            except Exception:
                logger.debug("world.snapshot() failed during trajectory extraction", exc_info=True)

        to_dict = getattr(world, "to_dict", None)
        if callable(to_dict):
            try:
                snapshot = to_dict()
                if isinstance(snapshot, Mapping):
                    return dict(snapshot)
            except Exception:
                logger.debug("world.to_dict() failed during trajectory extraction", exc_info=True)

        snapshot: dict[str, Any] = {
            "step": getattr(world, "step_count", 0),
            "agents": getattr(world, "agents", {}),
            "resources": getattr(world, "resources", {}),
        }

        metrics = getattr(world, "metrics", None)
        if metrics is not None:
            snapshot["metrics"] = metrics

        return snapshot

    def _population_counts_from_snapshot(self, snapshot: Mapping[str, Any]) -> dict[str, int]:
        """Extract population counts by agent type from a snapshot."""

        if "agent_type_counts" in snapshot and isinstance(snapshot["agent_type_counts"], Mapping):
            return {
                str(agent_type): int(count)
                for agent_type, count in snapshot["agent_type_counts"].items()
                if self._is_numeric(count)
            }

        if "population" in snapshot and isinstance(snapshot["population"], Mapping):
            return {
                str(agent_type): int(count)
                for agent_type, count in snapshot["population"].items()
                if self._is_numeric(count)
            }

        agents = self._collection_values(
            self._first_present(
                snapshot,
                ("agents", "agent_states", "agent_snapshots"),
                default=(),
            )
        )

        counts: dict[str, int] = defaultdict(int)

        for agent in agents:
            agent_type = self._field(agent, "type", "agent_type", "kind", default="unknown")
            alive = self._field(agent, "alive", default=True)

            if self.config.count_alive_only and alive is False:
                continue

            counts[str(agent_type)] += 1

        return dict(sorted(counts.items()))

    def _resource_amounts_from_snapshot(self, snapshot: Mapping[str, Any]) -> dict[str, float]:
        """Extract resource total amounts by resource type from a snapshot."""

        for key in ("resource_timeseries", "resource_type_amounts", "resource_amounts"):
            value = snapshot.get(key)
            if isinstance(value, Mapping):
                return {
                    str(resource_type): float(amount)
                    for resource_type, amount in value.items()
                    if self._is_numeric(amount)
                }

        resources = self._collection_values(
            self._first_present(
                snapshot,
                ("resources", "resource_states", "resource_snapshots"),
                default=(),
            )
        )

        amounts: dict[str, float] = defaultdict(float)

        for resource in resources:
            resource_type = self._field(resource, "type", "resource_type", "kind", default="unknown")
            amount = self._field(resource, "amount", "quantity", "value", default=0.0)

            if self._is_numeric(amount):
                amounts[str(resource_type)] += float(amount)

        return dict(sorted(amounts.items()))

    def _append_snapshot_events(
        self,
        *,
        snapshot: Mapping[str, Any],
        step: int,
        output: list[dict[str, Any]],
    ) -> None:
        """Append normalized event records from a snapshot."""

        raw_events = self._first_present(
            snapshot,
            ("event_log", "events_processed", "events_fired", "events"),
            default=(),
        )

        for record in self._collection_values(raw_events):
            normalized = self._normalize_event_record(record, default_step=step)
            if normalized:
                output.append(normalized)

    def _append_snapshot_metrics(
        self,
        *,
        snapshot: Mapping[str, Any],
        step: int,
        output: dict[str, list[dict[str, Any]]],
    ) -> None:
        """Append metric records from a snapshot."""

        raw_metrics = self._first_present(
            snapshot,
            ("metrics", "metric_values", "metric_history"),
            default=None,
        )

        if raw_metrics is None:
            return

        self._append_metric_records(raw_metrics, default_step=step, output=output)

    def _append_external_events(
        self,
        *,
        event_log: Sequence[Any] | Mapping[str, Any] | None,
        world: Any | None,
        output: list[dict[str, Any]],
    ) -> None:
        """Append externally supplied or world-level event log records."""

        raw_event_log = event_log

        if raw_event_log is None and world is not None:
            raw_event_log = self._first_present_from_object(
                world,
                ("event_log", "events_processed", "events_fired"),
                default=None,
            )

        if raw_event_log is None:
            return

        for record in self._collection_values(raw_event_log):
            normalized = self._normalize_event_record(record, default_step=None)
            if normalized:
                output.append(normalized)

    def _append_external_metrics(
        self,
        *,
        metric_history: Mapping[str, Any] | Sequence[Any] | None,
        world: Any | None,
        output: dict[str, list[dict[str, Any]]],
    ) -> None:
        """Append externally supplied or world-level metric history."""

        raw_metric_history = metric_history

        if raw_metric_history is None and world is not None:
            raw_metric_history = self._first_present_from_object(
                world,
                ("metric_history", "metrics_history", "metrics"),
                default=None,
            )

        if raw_metric_history is None:
            return

        self._append_metric_records(raw_metric_history, default_step=None, output=output)

    def _append_metric_records(
        self,
        raw_metrics: Any,
        *,
        default_step: int | None,
        output: dict[str, list[dict[str, Any]]],
    ) -> None:
        """Normalize and append metric records."""

        if isinstance(raw_metrics, Mapping):
            for metric_name, metric_value in raw_metrics.items():
                if isinstance(metric_value, Sequence) and not isinstance(metric_value, (str, bytes, Mapping)):
                    for item in metric_value:
                        output[str(metric_name)].append(
                            self._normalize_metric_record(
                                item,
                                metric_name=str(metric_name),
                                default_step=default_step,
                            )
                        )
                else:
                    output[str(metric_name)].append(
                        self._normalize_metric_record(
                            metric_value,
                            metric_name=str(metric_name),
                            default_step=default_step,
                        )
                    )
            return

        if isinstance(raw_metrics, Sequence) and not isinstance(raw_metrics, (str, bytes)):
            for item in raw_metrics:
                item_mapping = self._mapping_from_any(item)

                metric_name = str(
                    self._first_present(
                        item_mapping,
                        ("name", "metric", "metric_name"),
                        default="unknown_metric",
                    )
                )
                output[metric_name].append(
                    self._normalize_metric_record(
                        item_mapping,
                        metric_name=metric_name,
                        default_step=default_step,
                    )
                )

    def _normalize_event_record(
        self,
        record: Any,
        *,
        default_step: int | None,
    ) -> dict[str, Any]:
        """Normalize one event record."""

        mapping = self._mapping_from_any(record)

        if not mapping:
            return {}

        if default_step is not None:
            mapping.setdefault("step", default_step)

        event_name = self._first_present(
            mapping,
            ("name", "event", "event_name", "type"),
            default=None,
        )
        if event_name is not None:
            mapping.setdefault("name", str(event_name))

        return self._json_safe_mapping(mapping)

    def _normalize_metric_record(
        self,
        record: Any,
        *,
        metric_name: str,
        default_step: int | None,
    ) -> dict[str, Any]:
        """Normalize one metric record."""

        if isinstance(record, Mapping):
            output = self._json_safe_mapping(record)
        else:
            output = {"value": self._json_safe(record)}

        output.setdefault("name", metric_name)

        if default_step is not None:
            output.setdefault("step", default_step)

        return output

    def _preserve_unknown_snapshot_fields(
        self,
        snapshot: Mapping[str, Any],
        *,
        step: int,
    ) -> dict[str, Any]:
        """Preserve selected unknown fields from a snapshot for diagnostics."""

        known_keys = {
            "step",
            "step_count",
            "time",
            "tick",
            "agents",
            "agent_states",
            "agent_snapshots",
            "agent_type_counts",
            "population",
            "resources",
            "resource_states",
            "resource_snapshots",
            "resource_timeseries",
            "resource_type_amounts",
            "resource_amounts",
            "events",
            "event_log",
            "events_processed",
            "events_fired",
            "metrics",
            "metric_values",
            "metric_history",
        }

        unknown: dict[str, Any] = {}

        for key, value in snapshot.items():
            if key in known_keys:
                continue
            unknown[str(key)] = self._json_safe(value)

        if not unknown:
            return {}

        return {
            "step": step,
            "fields": unknown,
        }

    @staticmethod
    def _step_from_snapshot(snapshot: Mapping[str, Any], *, fallback: int) -> int:
        """Return step value from snapshot with a fallback index."""

        for key in ("step", "step_count", "time", "tick"):
            if key in snapshot and TrajectoryExtractor._is_numeric(snapshot[key]):
                return int(snapshot[key])

        return int(fallback)

    @staticmethod
    def _collection_values(value: Any) -> tuple[Any, ...]:
        """Return collection values from mapping or sequence."""

        if value is None:
            return ()

        if isinstance(value, Mapping):
            return tuple(value.values())

        if isinstance(value, Sequence) and not isinstance(value, (str, bytes)):
            return tuple(value)

        return (value,)

    @staticmethod
    def _mapping_from_any(value: Any) -> dict[str, Any]:
        """Return a mapping representation of an arbitrary value."""

        if value is None:
            return {}

        if isinstance(value, Mapping):
            return dict(value)

        for method_name in ("to_dict", "snapshot", "model_dump"):
            method = getattr(value, method_name, None)
            if callable(method):
                try:
                    result = method()
                    if isinstance(result, Mapping):
                        return dict(result)
                except Exception:
                    logger.debug(
                        "Could not convert %s using %s()",
                        value.__class__.__name__,
                        method_name,
                        exc_info=True,
                    )

        output: dict[str, Any] = {}

        for attribute in ("id", "name", "type", "amount", "position", "alive", "state", "metadata"):
            if hasattr(value, attribute):
                output[attribute] = getattr(value, attribute)

        return output

    @staticmethod
    def _field(value: Any, *names: str, default: Any = None) -> Any:
        """Read the first present field from mapping or object."""

        if isinstance(value, Mapping):
            for name in names:
                if name in value:
                    return value[name]
            return default

        for name in names:
            if hasattr(value, name):
                return getattr(value, name)

        return default

    @staticmethod
    def _first_present(mapping: Mapping[str, Any], keys: Sequence[str], *, default: Any) -> Any:
        """Return first present mapping value."""

        for key in keys:
            if key in mapping and mapping[key] is not None:
                return mapping[key]

        return default

    @staticmethod
    def _first_present_from_object(value: Any, keys: Sequence[str], *, default: Any) -> Any:
        """Return first present object attribute value."""

        for key in keys:
            if hasattr(value, key):
                candidate = getattr(value, key)
                if candidate is not None:
                    return candidate

        return default

    @staticmethod
    def _is_numeric(value: Any) -> bool:
        """Return whether value is a finite int or float but not bool."""

        if isinstance(value, bool):
            return False

        if not isinstance(value, (int, float)):
            return False

        return math.isfinite(float(value))

    @classmethod
    def _json_safe_mapping(cls, mapping: Mapping[str, Any]) -> dict[str, Any]:
        """Return JSON-safe copy of a mapping."""

        return {
            str(key): cls._json_safe(value)
            for key, value in mapping.items()
        }

    @classmethod
    def _json_safe(cls, value: Any) -> Any:
        """Return a JSON-friendly representation."""

        if value is None or isinstance(value, (str, bool)):
            return value

        if isinstance(value, int) and not isinstance(value, bool):
            return value

        if isinstance(value, float):
            return value if math.isfinite(value) else None

        if isinstance(value, Mapping):
            return {
                str(key): cls._json_safe(item)
                for key, item in value.items()
            }

        if isinstance(value, Sequence) and not isinstance(value, (str, bytes)):
            return [cls._json_safe(item) for item in value]

        for method_name in ("to_dict", "snapshot", "model_dump"):
            method = getattr(value, method_name, None)
            if callable(method):
                try:
                    return cls._json_safe(method())
                except Exception:
                    logger.debug(
                        "Could not JSON-normalize %s using %s()",
                        value.__class__.__name__,
                        method_name,
                        exc_info=True,
                    )

        if hasattr(value, "tolist") and callable(value.tolist):
            try:
                return cls._json_safe(value.tolist())
            except Exception:
                logger.debug("Could not JSON-normalize tolist() value", exc_info=True)

        return repr(value)


def extract_trajectory(
    world_or_history: Any,
    *,
    metric_history: Mapping[str, Any] | Sequence[Any] | None = None,
    event_log: Sequence[Any] | Mapping[str, Any] | None = None,
    config: TrajectoryExtractionConfig | None = None,
) -> TrajectoryData:
    """Convenience function for extracting trajectory data."""

    extractor = TrajectoryExtractor(config=config or TrajectoryExtractionConfig())
    return extractor.extract(
        world_or_history,
        metric_history=metric_history,
        event_log=event_log,
    )


__all__ = [
    "TrajectoryData",
    "TrajectoryExtractionConfig",
    "TrajectoryExtractor",
    "WorldHistoryProtocol",
    "extract_trajectory",
]