#!/usr/bin/env python3
"""
nima_affordance_graph.py — The 3D Affordance Graph

This is the "graphing layer" — a graph data structure representing the
room's walkable topology + surface affordances. It drives:
  - Nima's movement AI (A* pathfinding through the room)
  - The VFX renderer (where to place + animate the avatar)
  - Wall enforcement (she literally cannot pass through walls)
  - Door passages (edges through doorways)
  - Furniture interaction (sit on couch, lie on bed, jump on cushions)

NEUROBIOLOGICAL MAPPING:
  This is the hippocampal cognitive map + parietal affordance competition:
    - Place cells (O'Keefe, 2014 Nobel) → graph nodes
    - Grid cells (Moser & Moser, 2014 Nobel) → spatial layout
    - Parietal affordance encoding → node action labels
    - Prefrontal path planning → A* search through the graph

  The brain doesn't represent space as pixels — it represents it as a
  graph of places with action possibilities. That's exactly what this is.

GRAPH STRUCTURE:
  Nodes = positions in the room (x, y, z) with affordance labels
  Edges = valid movements between positions
  Walls = no edge (impassable)
  Doors = edges through wall boundaries
  Furniture = nodes with special affordances (sit, lie, jump)

  Nima pathfinds via A* on this graph. When she wants to "sit on the
  couch," she finds: walk to couch → transition to sitting → z rises
  to couch height → posture animates to seated.
"""
from __future__ import annotations

import heapq
import json
import logging
import math
import time
from dataclasses import dataclass, field
from enum import Enum
from typing import Any, Dict, List, Optional, Set, Tuple

logger = logging.getLogger("NimaGraph")


class AffordanceType(Enum):
    """What actions Nima can perform at a location."""
    WALK      = "walk"
    STAND     = "stand"
    SIT       = "sit"
    LIE_DOWN  = "lie_down"
    JUMP_ON   = "jump_on"
    REST_HAND = "rest_hand"
    LEAN_ON   = "lean_on"
    DUCK      = "duck"        # under low surfaces
    AVOID     = "avoid"       # walls, obstacles


@dataclass
class GraphNode:
    """A node in the affordance graph — a position Nima can be at."""
    node_id: str
    position: Tuple[float, float, float]  # (x, y, z) meters
    surface_type: str = "floor"  # floor / furniture / wall / door
    height: float = 0.0          # surface height at this node
    material: str = "hard"       # hard / soft / rigid / fabric
    affordances: Set[AffordanceType] = field(default_factory=set)
    is_walkable: bool = True
    is_doorway: bool = False
    metadata: Dict[str, Any] = field(default_factory=dict)

    def to_dict(self) -> Dict[str, Any]:
        return {
            "node_id": self.node_id,
            "position": list(self.position),
            "surface_type": self.surface_type,
            "height": self.height,
            "material": self.material,
            "affordances": [a.value for a in self.affordances],
            "is_walkable": self.is_walkable,
            "is_doorway": self.is_doorway,
        }


@dataclass
class GraphEdge:
    """An edge between two nodes — a valid movement."""
    from_node: str
    to_node: str
    distance: float  # meters
    is_vertical: bool = False  # True if z changes (sit down, jump up)
    transition_type: str = "walk"  # walk / sit_down / stand_up / jump / duck


class AffordanceGraph:
    """
    The 3D spatial graph — Nima's cognitive map of the room.

    This is where Nima "thinks" about space. She doesn't see pixels;
    she sees a graph of places she can be, with actions she can take
    at each place. The graph is built from the SyntheticVisionComposite's
    spatial map, and is updated as Nima moves through the room.
    """

    def __init__(self, grid_resolution: float = 0.5) -> None:
        """
        Args:
            grid_resolution: meters between grid nodes (default 0.5m)
        """
        self.resolution = grid_resolution
        self.nodes: Dict[str, GraphNode] = {}
        self.edges: Dict[str, List[GraphEdge]] = {}  # node_id → edges
        self._room_bounds: Dict[str, float] = {}
        self._furniture: List[Dict[str, Any]] = []

    def build_from_spatial_map(self, spatial_map: Any) -> None:
        """
        Build the graph from a SyntheticVisionComposite spatial map.

        The spatial map provides:
          - room_bounds: the room's physical limits
          - surfaces: 3D point cloud of walls/floor/furniture
          - affordances: what actions are possible where
        """
        # Extract room bounds
        bounds = getattr(spatial_map, "room_bounds", {})
        if not bounds:
            bounds = {"x_min": 0, "x_max": 4, "y_min": 0, "y_max": 4, "z_min": 0, "z_max": 2.5}
        self._room_bounds = bounds

        # Clear existing graph
        self.nodes.clear()
        self.edges.clear()

        # Generate floor grid
        x_min, x_max = bounds.get("x_min", 0), bounds.get("x_max", 4)
        y_min, y_max = bounds.get("y_min", 0), bounds.get("y_max", 4)
        res = self.resolution

        for x in np.arange(x_min, x_max + res, res):
            for y in np.arange(y_min, y_max + res, res):
                node_id = self._node_id(x, y, 0.0)
                node = GraphNode(
                    node_id=node_id,
                    position=(round(float(x), 3), round(float(y), 3), 0.0),
                    surface_type="floor",
                    height=0.0,
                    material="hard",
                    affordances={AffordanceType.WALK, AffordanceType.STAND, AffordanceType.LIE_DOWN},
                )
                self.nodes[node_id] = node

        # Add walls (non-walkable boundary nodes)
        for x in np.arange(x_min, x_max + res, res):
            for y, wall_y in [(y_min, y_min), (y_max, y_max)]:
                node_id = self._node_id(x, wall_y, 0.0)
                if node_id in self.nodes:
                    self.nodes[node_id].surface_type = "wall"
                    self.nodes[node_id].is_walkable = False
                    self.nodes[node_id].affordances = {AffordanceType.AVOID}
        for y in np.arange(y_min, y_max + res, res):
            for x, wall_x in [(x_min, x_min), (x_max, x_max)]:
                node_id = self._node_id(wall_x, y, 0.0)
                if node_id in self.nodes:
                    self.nodes[node_id].surface_type = "wall"
                    self.nodes[node_id].is_walkable = False
                    self.nodes[node_id].affordances = {AffordanceType.AVOID}

        # Add furniture from spatial map surfaces
        surfaces = getattr(spatial_map, "surfaces", [])
        furniture_added = 0
        for surface in surfaces:
            if surface.surface_type == "furniture":
                self._add_furniture_node(surface)
                furniture_added += 1

        # Add doorways (edges through walls)
        self._add_doorways()

        # Build edges between adjacent walkable nodes
        self._build_edges()

        logger.info("[Graph] built: %d nodes, %d edges, %d furniture nodes",
                    len(self.nodes), sum(len(e) for e in self.edges.values()),
                    furniture_added)

    def _node_id(self, x: float, y: float, z: float) -> str:
        """Generate a node ID from coordinates."""
        return f"n_{round(x, 2)}_{round(y, 2)}_{round(z, 2)}"

    def _add_furniture_node(self, surface: Any) -> None:
        """Add a furniture surface as a graph node with affordances."""
        pos = surface.position
        node_id = self._node_id(pos[0], pos[1], pos[2])
        affordances: Set[AffordanceType] = set()
        height = surface.height

        if height < 0.3:
            # Very low furniture (rug, low step) → walkable
            affordances = {AffordanceType.WALK, AffordanceType.STAND}
        elif height < 0.6:
            if surface.material == "soft":
                # Couch, bed, cushion → sit, lie, jump
                affordances = {AffordanceType.SIT, AffordanceType.LIE_DOWN, AffordanceType.JUMP_ON}
            else:
                # Chair, low table → sit, rest hand
                affordances = {AffordanceType.SIT, AffordanceType.REST_HAND}
        elif height < 1.0:
            # Table height → rest hand, lean
            affordances = {AffordanceType.REST_HAND, AffordanceType.LEAN_ON}
        else:
            # Tall obstacle → avoid
            affordances = {AffordanceType.AVOID}

        node = GraphNode(
            node_id=node_id,
            position=(round(pos[0], 3), round(pos[1], 3), round(pos[2], 3)),
            surface_type="furniture",
            height=height,
            material=surface.material,
            affordances=affordances,
            is_walkable=height < 0.3,  # only walkable if very low
        )
        self.nodes[node_id] = node

    def _add_doorways(self) -> None:
        """Mark doorway nodes as walkable passages through walls."""
        # In a real implementation, doorways would be detected from the
        # spatial map (gaps in wall surfaces). For now, add a default
        # doorway in the center of one wall.
        bounds = self._room_bounds
        door_x = (bounds.get("x_min", 0) + bounds.get("x_max", 4)) / 2
        door_y = bounds.get("y_max", 4)
        door_node_id = self._node_id(door_x, door_y, 0.0)
        if door_node_id in self.nodes:
            self.nodes[door_node_id].is_walkable = True
            self.nodes[door_node_id].is_doorway = True
            self.nodes[door_node_id].surface_type = "door"
            self.nodes[door_node_id].affordances = {AffordanceType.WALK, AffordanceType.STAND}
            logger.debug("[Graph] doorway at %s", door_node_id)

    def _build_edges(self) -> None:
        """Build edges between adjacent walkable nodes."""
        res = self.resolution
        for node_id, node in self.nodes.items():
            if not node.is_walkable:
                continue
            x, y, z = node.position
            # Check 4 neighbors (N, S, E, W)
            for dx, dy in [(res, 0), (-res, 0), (0, res), (0, -res)]:
                neighbor_id = self._node_id(x + dx, y + dy, z)
                if neighbor_id in self.nodes:
                    neighbor = self.nodes[neighbor_id]
                    if neighbor.is_walkable:
                        edge = GraphEdge(
                            from_node=node_id,
                            to_node=neighbor_id,
                            distance=res,
                            transition_type="walk",
                        )
                        self.edges.setdefault(node_id, []).append(edge)

            # Check vertical edges (to furniture nodes for sit/lie/jump)
            for other_id, other in self.nodes.items():
                if other_id == node_id:
                    continue
                if other.surface_type != "furniture":
                    continue
                ox, oy, oz = other.position
                # If furniture is adjacent (within resolution) and higher
                horiz_dist = math.sqrt((x - ox)**2 + (y - oy)**2)
                if horiz_dist < res * 1.5 and oz > z:
                    if AffordanceType.SIT in other.affordances:
                        self.edges.setdefault(node_id, []).append(GraphEdge(
                            from_node=node_id,
                            to_node=other_id,
                            distance=math.sqrt(horiz_dist**2 + (oz - z)**2),
                            is_vertical=True,
                            transition_type="sit_down",
                        ))
                    elif AffordanceType.JUMP_ON in other.affordances:
                        self.edges.setdefault(node_id, []).append(GraphEdge(
                            from_node=node_id,
                            to_node=other_id,
                            distance=math.sqrt(horiz_dist**2 + (oz - z)**2),
                            is_vertical=True,
                            transition_type="jump",
                        ))

    def find_path(self,
                  start: Tuple[float, float, float],
                  goal: Tuple[float, float, float],
                  ) -> List[GraphNode]:
        """
        A* pathfinding from start to goal.

        Returns a list of GraphNodes representing the path, or empty
        list if no path exists.

        NEUROBIOLOGICAL ANALOGUE:
          This is prefrontal path planning — the brain's ability to
          plan a route through space before executing it. Place cells
          in the hippocampus fire in sequence during planning,
          "rehearsing" the route. A* does the same computation.
        """
        start_id = self._nearest_node(start)
        goal_id = self._nearest_node(goal)
        if start_id is None or goal_id is None:
            return []

        # A* search
        open_set: List[Tuple[float, str]] = [(0, start_id)]
        came_from: Dict[str, str] = {}
        g_score: Dict[str, float] = {start_id: 0}
        f_score: Dict[str, float] = {start_id: self._heuristic(start_id, goal_id)}

        while open_set:
            _, current_id = heapq.heappop(open_set)
            if current_id == goal_id:
                # Reconstruct path
                path = []
                cid = current_id
                while cid in came_from:
                    path.append(self.nodes[cid])
                    cid = came_from[cid]
                path.append(self.nodes[start_id])
                path.reverse()
                return path

            for edge in self.edges.get(current_id, []):
                neighbor_id = edge.to_node
                tentative_g = g_score[current_id] + edge.distance
                if neighbor_id not in g_score or tentative_g < g_score[neighbor_id]:
                    came_from[neighbor_id] = current_id
                    g_score[neighbor_id] = tentative_g
                    f_score[neighbor_id] = tentative_g + self._heuristic(neighbor_id, goal_id)
                    heapq.heappush(open_set, (f_score[neighbor_id], neighbor_id))

        return []  # no path found

    def _nearest_node(self, pos: Tuple[float, float, float]) -> Optional[str]:
        """Find the nearest walkable node to a position."""
        best_id = None
        best_dist = float("inf")
        for node_id, node in self.nodes.items():
            if not node.is_walkable:
                continue
            dist = math.sqrt(
                (node.position[0] - pos[0])**2 +
                (node.position[1] - pos[1])**2 +
                (node.position[2] - pos[2])**2
            )
            if dist < best_dist:
                best_dist = dist
                best_id = node_id
        return best_id

    def _heuristic(self, a_id: str, b_id: str) -> float:
        """Euclidean distance heuristic for A*."""
        a = self.nodes[a_id].position
        b = self.nodes[b_id].position
        return math.sqrt(
            (a[0] - b[0])**2 + (a[1] - b[1])**2 + (a[2] - b[2])**2
        )

    def find_affordance(self,
                        affordance: AffordanceType,
                        from_pos: Tuple[float, float, float],
                        ) -> Optional[GraphNode]:
        """
        Find the nearest node with a specific affordance.
        E.g., "find the nearest place I can sit."
        """
        best_node = None
        best_dist = float("inf")
        for node in self.nodes.values():
            if affordance in node.affordances:
                dist = math.sqrt(
                    (node.position[0] - from_pos[0])**2 +
                    (node.position[1] - from_pos[1])**2 +
                    (node.position[2] - from_pos[2])**2
                )
                if dist < best_dist:
                    best_dist = dist
                    best_node = node
        return best_node

    def get_node_at(self, x: float, y: float, z: float = 0.0) -> Optional[GraphNode]:
        """Get the node at a specific position."""
        node_id = self._node_id(x, y, z)
        return self.nodes.get(node_id)

    def get_stats(self) -> Dict[str, Any]:
        return {
            "total_nodes": len(self.nodes),
            "walkable_nodes": sum(1 for n in self.nodes.values() if n.is_walkable),
            "furniture_nodes": sum(1 for n in self.nodes.values() if n.surface_type == "furniture"),
            "doorway_nodes": sum(1 for n in self.nodes.values() if n.is_doorway),
            "total_edges": sum(len(e) for e in self.edges.values()),
            "room_bounds": self._room_bounds,
            "resolution": self.resolution,
        }

    def to_dict(self) -> Dict[str, Any]:
        """Serialize the graph for the renderer."""
        return {
            "nodes": {nid: n.to_dict() for nid, n in self.nodes.items()},
            "edges": {nid: [{"to": e.to_node, "dist": e.distance,
                             "type": e.transition_type}
                            for e in edges]
                      for nid, edges in self.edges.items()},
            "stats": self.get_stats(),
        }


# ── Need numpy for arange ──
import numpy as np


# ═══════════════════════════════════════════════════════════════════════════
# SELF-TEST
# ═══════════════════════════════════════════════════════════════════════════

if __name__ == "__main__":
    logging.basicConfig(level=logging.INFO,
                        format="%(asctime)s [%(levelname)s] %(message)s")

    print("=== Nima Affordance Graph — Self Test ===\n")

    # Build a graph from a simulated spatial map
    class FakeSurface:
        def __init__(self, pos, stype, height, material):
            self.position = pos
            self.surface_type = stype
            self.height = height
            self.material = material

    class FakeSpatialMap:
        def __init__(self):
            self.room_bounds = {"x_min": 0, "x_max": 4, "y_min": 0, "y_max": 4, "z_min": 0, "z_max": 2.5}
            # Add a couch at (1, 1, 0.4)
            self.surfaces = [
                FakeSurface((1.0, 1.0, 0.4), "furniture", 0.4, "soft"),
                FakeSurface((3.0, 3.0, 0.8), "furniture", 0.8, "rigid"),
            ]

    graph = AffordanceGraph(grid_resolution=0.5)
    graph.build_from_spatial_map(FakeSpatialMap())

    stats = graph.get_stats()
    print(f"Graph stats:")
    print(f"  Nodes: {stats['total_nodes']} ({stats['walkable_nodes']} walkable)")
    print(f"  Furniture: {stats['furniture_nodes']}")
    print(f"  Doorways: {stats['doorway_nodes']}")
    print(f"  Edges: {stats['total_edges']}")
    print()

    # Test pathfinding
    print("=== Pathfinding test ===")
    path = graph.find_path((0.5, 0.5, 0.0), (3.5, 3.5, 0.0))
    print(f"Path from (0.5, 0.5) to (3.5, 3.5): {len(path)} steps")
    for node in path[:5]:
        print(f"  {node.node_id} at {node.position} ({node.surface_type})")
    if len(path) > 5:
        print(f"  ... ({len(path) - 5} more steps)")
    print()

    # Test affordance search
    print("=== Affordance search ===")
    sit_spot = graph.find_affordance(AffordanceType.SIT, (0.5, 0.5, 0.0))
    if sit_spot:
        print(f"Nearest sit-able spot: {sit_spot.node_id} at {sit_spot.position}")
        print(f"  Material: {sit_spot.material}, Height: {sit_spot.height}m")

    jump_spot = graph.find_affordance(AffordanceType.JUMP_ON, (0.5, 0.5, 0.0))
    if jump_spot:
        print(f"Nearest jump-able spot: {jump_spot.node_id} at {jump_spot.position}")

    print(f"\n=== Graph self-test PASSED ===")