metrics/diversity.py

"""
Generic diversity metrics for WorldSmithAI.

This module computes domain-agnostic diversity over agents, resources, or both.
It deliberately avoids hardcoded species assumptions. A "group" can be any
DSL-defined label, such as agent type, resource type, faction, strategy,
profession, role, technology, market sector, node class, or civilization tier.

The default grouping path is ``type``, which means:
    - agents are grouped by ``agent.type``
    - resources are grouped by ``resource.type``

Example:
    metric = DiversityMetric(collection="agents", group_by_path="type")
    result = metric.compute(world)

    print(result.richness)
    print(result.gini_simpson_index)

Weighted resource diversity:
    metric = DiversityMetric(
        collection="resources",
        group_by_path="type",
        weight_path="amount",
    )
    result = metric(world)

Future extensibility:
    - Add rolling-window diversity metrics.
    - Add spatial diversity over regions.
    - Add interaction-network diversity.
    - Add behavioral diversity based on executed actions.
    - Add diversity deltas for interestingness scoring.
    - Add "God Agent" hooks for judging world health and novelty.
"""

from __future__ import annotations

import copy
import logging
import math
from collections.abc import Iterable, Mapping, MutableSequence, Sequence
from dataclasses import dataclass, field
from enum import Enum
from numbers import Real
from types import MappingProxyType
from typing import TYPE_CHECKING, Any, ClassVar

import numpy as np

if TYPE_CHECKING:
    from core.agent import Agent
    from core.resource import Resource
    from core.world import World


logger = logging.getLogger(__name__)

_MISSING = object()
_EPSILON = 1.0e-12


class DiversityCollection(str, Enum):
    """Collections over which diversity may be computed."""

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


@dataclass(frozen=True)
class DiversityResult:
    """Result produced by a diversity metric computation.

    Attributes:
        metric_name: Stable metric name.
        collection: Collection analyzed, such as ``agents`` or ``resources``.
        group_by_path: Path used to derive group labels.
        weight_path: Optional path used to derive item weights.
        item_count: Number of world objects included after filters.
        total_weight: Sum of group weights. For unweighted diversity this is
            equal to ``item_count``.
        richness: Number of distinct groups.
        counts: Group weights keyed by group label.
        proportions: Group proportions keyed by group label.
        dominant_group: Group with the largest proportion, if any.
        dominance: Largest group proportion.
        simpson_concentration: Sum of squared proportions.
        gini_simpson_index: ``1 - simpson_concentration``.
        inverse_simpson_index: ``1 / simpson_concentration`` when defined.
        shannon_entropy: Shannon entropy using natural logarithms.
        pielou_evenness: Shannon entropy divided by maximum possible entropy.
        step: Optional world step.
        metadata: Additional JSON-compatible details.
    """

    metric_name: str
    collection: str
    group_by_path: str
    weight_path: str | None
    item_count: int
    total_weight: float
    richness: int
    counts: Mapping[str, float] = field(default_factory=dict)
    proportions: Mapping[str, float] = field(default_factory=dict)
    dominant_group: str | None = None
    dominance: float = 0.0
    simpson_concentration: float = 0.0
    gini_simpson_index: float = 0.0
    inverse_simpson_index: float = 0.0
    shannon_entropy: float = 0.0
    pielou_evenness: float = 0.0
    step: int | None = None
    metadata: Mapping[str, Any] = field(default_factory=dict)

    @property
    def group_count(self) -> int:
        """Return the number of distinct groups."""

        return self.richness

    @property
    def is_empty(self) -> bool:
        """Return whether no items contributed to the metric."""

        return self.item_count == 0 or self.total_weight <= 0.0

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

        return {
            "metric_name": self.metric_name,
            "collection": self.collection,
            "group_by_path": self.group_by_path,
            "weight_path": self.weight_path,
            "item_count": self.item_count,
            "total_weight": self.total_weight,
            "richness": self.richness,
            "group_count": self.group_count,
            "counts": copy.deepcopy(dict(self.counts)),
            "proportions": copy.deepcopy(dict(self.proportions)),
            "dominant_group": self.dominant_group,
            "dominance": self.dominance,
            "simpson_concentration": self.simpson_concentration,
            "gini_simpson_index": self.gini_simpson_index,
            "inverse_simpson_index": self.inverse_simpson_index,
            "shannon_entropy": self.shannon_entropy,
            "pielou_evenness": self.pielou_evenness,
            "step": self.step,
            "metadata": copy.deepcopy(dict(self.metadata)),
        }


@dataclass
class DiversityMetric:
    """Compute generic group diversity over a world collection.

    By default, this metric computes diversity over alive agents grouped by
    their ``type`` attribute. The same class can compute resource diversity,
    strategy diversity, role diversity, faction diversity, or any other
    path-based grouping.

    Examples:
        Agent type diversity:
            metric = DiversityMetric()
            result = metric.compute(world)

        Resource amount diversity:
            metric = DiversityMetric(
                collection="resources",
                group_by_path="type",
                weight_path="amount",
            )
            result = metric.compute(world)

        Diversity by active goal:
            metric = DiversityMetric(
                collection="agents",
                group_by_path="state.current_goal",
                include_missing=True,
            )
            result = metric.compute(world)
    """

    name: ClassVar[str] = "diversity"

    collection: DiversityCollection | str = DiversityCollection.AGENTS
    group_by_path: str = "type"
    weight_path: str | None = None

    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
    metadata: Mapping[str, Any] = field(default_factory=dict)

    def compute(self, world: World) -> DiversityResult:
        """Compute diversity for the configured world collection.

        Args:
            world: Runtime world object.

        Returns:
            ``DiversityResult`` containing counts, proportions, and diversity
            indices.
        """

        collection = _normalize_collection(self.collection)
        items = self._items(world, collection)
        counts, included_count, skipped_count = self._counts(items, collection)

        result = diversity_from_counts(
            counts,
            metric_name=self.name,
            collection=collection.value,
            group_by_path=self.group_by_path,
            weight_path=self.weight_path,
            item_count=included_count,
            step=_world_step(world),
            metadata={
                **copy.deepcopy(dict(self.metadata)),
                "skipped_count": skipped_count,
                "alive_only": self.alive_only,
                "include_missing": self.include_missing,
                "minimum_weight": self.minimum_weight,
            },
        )

        logger.debug(
            "Computed %s diversity over %s: richness=%s total_weight=%.3f",
            self.name,
            collection.value,
            result.richness,
            result.total_weight,
        )

        return result

    def __call__(self, world: World) -> DiversityResult:
        """Compute diversity, allowing metric instances to be called directly."""

        return self.compute(world)

    def _items(self, world: World, collection: DiversityCollection) -> tuple[Any, ...]:
        """Return world items for the configured collection."""

        return _iter_world_items(world, collection)

    def _counts(
        self,
        items: Sequence[Any],
        collection: DiversityCollection,
    ) -> tuple[dict[str, float], int, int]:
        """Return group weights, included item count, and skipped item count."""

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

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

        for item in items:
            if self.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.case_sensitive,
                strip=self.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.minimum_weight)):
                skipped_count += 1
                continue

            counts[normalized_label] = counts.get(normalized_label, 0.0) + weight
            included_count += 1

        return counts, included_count, skipped_count

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

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

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

        label = _stable_label(raw_label)

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

        return label

    def _weight_for(self, item: Any) -> float:
        """Return the non-negative contribution weight for an item."""

        if self.weight_path is None:
            return 1.0

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

        return max(0.0, float(raw_weight))


@dataclass
class DiversityTracker:
    """Track diversity metric values over time.

    This class is useful for simulations that want a small in-memory time series
    without requiring a full metrics subsystem.
    """

    metric: DiversityMetric = field(default_factory=DiversityMetric)
    max_history: int = 1000
    history: list[DiversityResult] = field(default_factory=list)

    def update(self, world: World) -> DiversityResult:
        """Compute the metric for the current world and append it to history."""

        result = self.metric.compute(world)
        _append_bounded(self.history, result, self.max_history)
        return result

    def latest(self) -> DiversityResult | None:
        """Return the latest diversity result, if any."""

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

    def to_series(self, value_key: str = "gini_simpson_index") -> list[dict[str, Any]]:
        """Return a JSON-friendly time series for one result field.

        Args:
            value_key: Name of the ``DiversityResult`` attribute to extract.

        Returns:
            List of dictionaries with ``step`` and ``value`` keys.
        """

        series: list[dict[str, Any]] = []

        for index, result in enumerate(self.history):
            value = getattr(result, value_key, None)
            if not _is_number(value):
                continue

            series.append(
                {
                    "index": index,
                    "step": result.step,
                    "value": float(value),
                }
            )

        return series

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

        return {
            "metric": {
                "name": self.metric.name,
                "collection": _normalize_collection(self.metric.collection).value,
                "group_by_path": self.metric.group_by_path,
                "weight_path": self.metric.weight_path,
            },
            "max_history": self.max_history,
            "history": [result.to_dict() for result in self.history],
        }


def compute_diversity(
    world: World,
    *,
    collection: DiversityCollection | str = DiversityCollection.AGENTS,
    group_by_path: str = "type",
    weight_path: str | None = None,
    alive_only: bool = True,
    include_missing: bool = False,
) -> DiversityResult:
    """Convenience function for computing generic diversity."""

    metric = DiversityMetric(
        collection=collection,
        group_by_path=group_by_path,
        weight_path=weight_path,
        alive_only=alive_only,
        include_missing=include_missing,
    )
    return metric.compute(world)


def compute_agent_type_diversity(
    world: World,
    *,
    alive_only: bool = True,
    include_missing: bool = False,
) -> DiversityResult:
    """Compute diversity over agent ``type`` labels."""

    return compute_diversity(
        world,
        collection=DiversityCollection.AGENTS,
        group_by_path="type",
        alive_only=alive_only,
        include_missing=include_missing,
    )


def compute_resource_type_diversity(
    world: World,
    *,
    weight_by_amount: bool = False,
    include_missing: bool = False,
) -> DiversityResult:
    """Compute diversity over resource ``type`` labels.

    Args:
        world: Runtime world object.
        weight_by_amount: If true, resource types are weighted by ``amount``.
            If false, each resource object contributes equally.
        include_missing: Whether resources missing ``type`` should be grouped
            under the missing label.
    """

    return compute_diversity(
        world,
        collection=DiversityCollection.RESOURCES,
        group_by_path="type",
        weight_path="amount" if weight_by_amount else None,
        alive_only=False,
        include_missing=include_missing,
    )


def diversity_from_counts(
    counts: Mapping[str, float],
    *,
    metric_name: str = "diversity",
    collection: str = "custom",
    group_by_path: str = "group",
    weight_path: str | None = None,
    item_count: int | None = None,
    step: int | None = None,
    metadata: Mapping[str, Any] | None = None,
) -> DiversityResult:
    """Compute diversity statistics from precomputed group counts.

    Args:
        counts: Mapping from group label to non-negative weight.
        metric_name: Name stored in the result.
        collection: Collection label stored in the result.
        group_by_path: Grouping path stored in the result.
        weight_path: Optional weight path stored in the result.
        item_count: Optional number of contributing items.
        step: Optional world step.
        metadata: Optional result metadata.

    Returns:
        ``DiversityResult`` with diversity statistics.
    """

    cleaned_counts = {
        str(label): max(0.0, float(weight))
        for label, weight in counts.items()
        if _is_number(weight) and float(weight) > 0.0
    }

    sorted_counts = dict(sorted(cleaned_counts.items(), key=lambda item: item[0]))
    total_weight = float(sum(sorted_counts.values()))
    richness = len(sorted_counts)

    if total_weight <= _EPSILON or richness == 0:
        return DiversityResult(
            metric_name=metric_name,
            collection=collection,
            group_by_path=group_by_path,
            weight_path=weight_path,
            item_count=0 if item_count is None else int(item_count),
            total_weight=0.0,
            richness=0,
            counts={},
            proportions={},
            step=step,
            metadata=metadata or {},
        )

    proportions = {
        label: weight / total_weight
        for label, weight in sorted_counts.items()
    }

    dominant_group, dominance = max(
        proportions.items(),
        key=lambda item: (item[1], item[0]),
    )

    probability_values = np.asarray(list(proportions.values()), dtype=float)
    simpson_concentration = float(np.sum(probability_values**2))
    gini_simpson_index = float(1.0 - simpson_concentration)
    inverse_simpson_index = (
        float(1.0 / simpson_concentration)
        if simpson_concentration > _EPSILON
        else 0.0
    )

    shannon_entropy = float(
        -np.sum(
            probability_values
            * np.log(np.clip(probability_values, _EPSILON, None))
        )
    )

    max_entropy = math.log(richness) if richness > 1 else 0.0
    pielou_evenness = (
        float(shannon_entropy / max_entropy)
        if max_entropy > _EPSILON
        else 1.0
    )

    return DiversityResult(
        metric_name=metric_name,
        collection=collection,
        group_by_path=group_by_path,
        weight_path=weight_path,
        item_count=len(sorted_counts) if item_count is None else int(item_count),
        total_weight=total_weight,
        richness=richness,
        counts=sorted_counts,
        proportions=proportions,
        dominant_group=dominant_group,
        dominance=float(dominance),
        simpson_concentration=simpson_concentration,
        gini_simpson_index=gini_simpson_index,
        inverse_simpson_index=inverse_simpson_index,
        shannon_entropy=shannon_entropy,
        pielou_evenness=pielou_evenness,
        step=step,
        metadata=metadata or {},
    )


def _normalize_collection(value: DiversityCollection | str) -> DiversityCollection:
    """Normalize a diversity collection value."""

    if isinstance(value, DiversityCollection):
        return value

    return DiversityCollection(str(value))


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 a runtime 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 DiversityCollection.AGENTS.value

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

    return "items"


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

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

    if collection is DiversityCollection.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 _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 path from an item.

    Supported special 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 a group value."""

    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 group label for deterministic comparisons."""

    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 an optional maximum history length."""

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


METRIC_REGISTRY: Mapping[str, type[DiversityMetric]] = MappingProxyType(
    {
        DiversityMetric.name: DiversityMetric,
    }
)


__all__ = [
    "DiversityCollection",
    "DiversityMetric",
    "DiversityResult",
    "DiversityTracker",
    "METRIC_REGISTRY",
    "compute_agent_type_diversity",
    "compute_diversity",
    "compute_resource_type_diversity",
    "diversity_from_counts",
]