agents/wayfinder/memory_graph.py

"""State memory graph — Module C.

A hash-deduplicated directed graph of observed game states. Nodes are
hashed frames; edges are actions taken. Tracks visit counts for novelty
detection and supports backtracking/loop avoidance.

Inspired by Blind Squirrel's graph-based exploration, but integrated
with the learned perception encoder for efficient similarity checks.
"""

from __future__ import annotations

import hashlib
import logging
from collections import defaultdict

import numpy as np

logger = logging.getLogger(__name__)


class MemoryGraph:
    """Directed graph of observed game states.

    Nodes are identified by frame hashes (MD5 of the raw 64×64 frame).
    Each node stores:
    - A latent vector (from the perception encoder) for similarity checks.
    - The score observed at that state.
    - A visit count (incremented each time the state is seen).

    Edges store the action that caused the transition.

    The graph resets per RESET (per level attempt), but the agent can
    optionally carry summary statistics across attempts.

    Attributes:
        nodes: Dict mapping frame_hash → node data.
        edges: Dict mapping (from_hash, to_hash) → edge data.
        adjacency: Dict mapping from_hash → list of (action, to_hash).
    """

    def __init__(self) -> None:
        """Initialize an empty memory graph."""
        self.nodes: dict[str, dict] = {}
        self.edges: dict[tuple[str, str], dict] = {}
        self.adjacency: dict[str, list[tuple[str, str]]] = defaultdict(list)

    @staticmethod
    def hash_frame(frame: np.ndarray) -> str:
        """Compute a stable hash for a 64×64 frame.

        Args:
            frame: 64×64 uint8 array of color indices.

        Returns:
            MD5 hex string of the frame data.
        """
        return hashlib.md5(frame.tobytes()).hexdigest()

    def add_node(
        self,
        frame_hash: str,
        latent: np.ndarray,
        score: float = 0.0,
    ) -> None:
        """Add or update a node in the graph.

        If the node already exists, increment its visit count.

        Args:
            frame_hash: Hash of the frame.
            latent: Latent vector from the perception encoder.
            score: Game score at this state.
        """
        if frame_hash in self.nodes:
            self.nodes[frame_hash]["visit_count"] += 1
            self.nodes[frame_hash]["score"] = score
        else:
            self.nodes[frame_hash] = {
                "latent": latent.copy(),
                "score": score,
                "visit_count": 1,
            }

    def add_edge(
        self,
        from_hash: str,
        to_hash: str,
        action: str,
        action_data: dict | None = None,
    ) -> None:
        """Add a directed edge to the graph.

        Args:
            from_hash: Source node hash.
            to_hash: Destination node hash.
            action: Action name that caused this transition.
            action_data: Optional action data (e.g. coordinates for ACTION6).
        """
        edge_key = (from_hash, to_hash)
        if edge_key in self.edges:
            self.edges[edge_key]["count"] += 1
        else:
            self.edges[edge_key] = {
                "action": action,
                "action_data": action_data,
                "count": 1,
            }
            self.adjacency[from_hash].append((action, to_hash))

    def novelty(self, frame_hash: str) -> float:
        """Compute novelty score for a frame.

        Novelty is high for unvisited or rarely-visited states, low for
        frequently-visited ones. Used by the intrinsic reward module.

        Formula: novelty = 1.0 / (1.0 + visit_count)

        Args:
            frame_hash: Hash of the frame to evaluate.

        Returns:
            Novelty score in (0, 1]. New states return 1.0.
        """
        if frame_hash not in self.nodes:
            return 1.0
        visit_count = self.nodes[frame_hash]["visit_count"]
        return 1.0 / (1.0 + visit_count)

    def shortest_path_to_unexplored(self, from_hash: str) -> list[str] | None:
        """Find shortest path from a node to any unexplored action.

        Uses BFS to find the nearest node that has untried actions
        (actions not yet taken from that node).

        Args:
            from_hash: Starting node hash.

        Returns:
            List of action names forming the path, or None if no
            unexplored action is reachable.
        """
        if from_hash not in self.nodes:
            return None

        # BFS
        from collections import deque

        queue: deque[tuple[str, list[str]]] = deque([(from_hash, [])])
        visited: set[str] = {from_hash}

        all_actions = {"RESET", "ACTION1", "ACTION2", "ACTION3",
                       "ACTION4", "ACTION5", "ACTION6"}

        while queue:
            current_hash, path = queue.popleft()

            # Check if current node has unexplored actions
            tried_actions = {
                edge_action
                for edge_action, _ in self.adjacency.get(current_hash, [])
            }
            unexplored = all_actions - tried_actions

            if unexplored and len(path) > 0:
                return path
            if unexplored:
                # We're at the start node with unexplored actions
                return []

            # Expand neighbors
            for edge_action, neighbor_hash in self.adjacency.get(current_hash, []):
                if neighbor_hash not in visited:
                    visited.add(neighbor_hash)
                    queue.append((neighbor_hash, path + [edge_action]))

        return None

    def get_visited_states(self) -> set[str]:
        """Return the set of all visited node hashes.

        Returns:
            Set of frame hashes that have been visited.
        """
        return set(self.nodes.keys())

    def get_transition_count(self) -> int:
        """Return total number of unique transitions (edges).

        Returns:
            Number of edges in the graph.
        """
        return len(self.edges)

    def stats(self) -> dict:
        """Return summary statistics about the graph.

        Returns:
            Dict with node_count, edge_count, avg_visits, max_visits.
        """
        visit_counts = [n["visit_count"] for n in self.nodes.values()]
        return {
            "node_count": len(self.nodes),
            "edge_count": len(self.edges),
            "avg_visits": np.mean(visit_counts) if visit_counts else 0.0,
            "max_visits": max(visit_counts) if visit_counts else 0,
        }

    def reset(self) -> None:
        """Clear all nodes and edges.

        Called when the agent RESETs a level or starts a new one.
        """
        self.nodes.clear()
        self.edges.clear()
        self.adjacency.clear()
        logger.debug("Memory graph reset")

    def serialize(self) -> dict:
        """Serialize the graph for checkpointing.

        Returns:
            Dict representation of the graph.
        """
        return {
            "nodes": {
                h: {"score": n["score"], "visit_count": n["visit_count"]}
                for h, n in self.nodes.items()
            },
            "edges": {
                f"{k[0]}->{k[1]}": v for k, v in self.edges.items()
            },
        }