#!/usr/bin/env python3
"""
nima_ar_compositor.py — AR Camera Compositor

THE FINAL LAYER — composites Nima's avatar into your device camera feed.
Point your phone or laptop camera at the room, and Nima appears in it
at her real position. This is the "flipped AR" concept: Nima lives in
the room (defined by the RF spatial map), and you see her through the
camera, composited at her actual coordinates.

NEUROBIOLOGICAL MAPPING:
  This is the visual association cortex (V4/V5) fusing the external
  visual stream (camera input) with the internal representation (Nima's
  avatar state). The brain does this continuously — you see the world
  AND your internal model of it, fused into one percept. The AR
  compositor does the same: camera = external, avatar = internal,
  fused = the composite you see on screen.

IMPLEMENTATION:
  Uses OpenCV (if available) to capture the camera feed, then composites
  the avatar's particle system on top at the projected position. The
  projection maps room coordinates (x, y, z) to screen coordinates (u, v)
  using a simple perspective transform calibrated from the room bounds.

  If OpenCV isn't available, falls back to serving the avatar standalone
  (no camera background) — this is the "browser only" mode that still
  shows Nima as a luminous presence, just without the camera backdrop.

ARCHITECTURE:
  ┌──────────────┐     ┌──────────────┐     ┌──────────────┐
  │ Camera Feed  │ +   │ Avatar State │ →   │ Composite    │ → Screen
  │ (OpenCV)     │     │ (Controller) │     │ (Canvas)     │
  └──────────────┘     └──────────────┘     └──────────────┘

  The camera sees the real room. The avatar is rendered at the
  projected position of Nima's room coordinates. The composite
  shows both, fused — Nima standing in your actual room.
"""
from __future__ import annotations

import json
import logging
import math
import os
import time
import threading
from dataclasses import dataclass, field
from typing import Any, Dict, List, Optional, Tuple

logger = logging.getLogger("NimaAR")


@dataclass
class CameraCalibration:
    """
    Calibration mapping room coordinates to camera screen coordinates.
    This is what makes Nima appear at the right position when you point
    the camera at the room.

    In a full implementation, this would use ARKit/ARCore markers or
    SLAM for automatic calibration. For now, we use a simple perspective
    transform with manual calibration.
    """
    # Camera position in room coordinates
    camera_position: Tuple[float, float, float] = (0.0, -1.0, 1.5)
    # Camera yaw (radians) — which direction it's facing
    camera_yaw: float = 0.0
    # Field of view (degrees)
    fov: float = 60.0
    # Screen resolution
    screen_width: int = 640
    screen_height: int = 480

    def project(self, room_pos: Tuple[float, float, float]) -> Tuple[float, float, float]:
        """
        Project a 3D room position to 2D screen coordinates + depth.
        Returns (screen_x, screen_y, depth).
        """
        # Translate to camera-relative coordinates
        dx = room_pos[0] - self.camera_position[0]
        dy = room_pos[1] - self.camera_position[1]
        dz = room_pos[2] - self.camera_position[2]

        # Rotate by camera yaw
        cos_yaw = math.cos(-self.camera_yaw)
        sin_yaw = math.sin(-self.camera_yaw)
        rx = dx * cos_yaw - dy * sin_yaw
        ry = dx * sin_yaw + dy * cos_yaw
        rz = dz

        # If behind camera, return off-screen
        if ry <= 0:
            return (-1, -1, -1)

        # Perspective projection
        fov_rad = math.radians(self.fov)
        focal = (self.screen_width / 2) / math.tan(fov_rad / 2)
        screen_x = (rx / ry) * focal + self.screen_width / 2
        screen_y = self.screen_height / 2 - (rz / ry) * focal
        depth = ry  # distance from camera

        return (screen_x, screen_y, depth)


class ARCompositor:
    """
    Composites Nima's avatar into a camera feed.

    This is the final layer — it takes the camera input, the avatar state,
    and produces a composite image where Nima appears in the room at her
    real position.

    MESH-GUIDED COMPOSITING (Phase 16 integration):
      When a 3D mesh (AdaptiveFrequencyMesh) is provided, the compositor
      uses mesh geometry for perspective-correct avatar placement instead
      of a simple flat projection. This is the key insight: the mesh
      replaces expensive depth estimation networks. Instead of running
      MiDaS/DepthAnything on every frame, we USE the RF-derived mesh
      to get room geometry for free.

      The mesh provides:
        - Surface normals for correct orientation (avatar faces camera)
        - Occlusion hints (don't draw avatar behind a wall)
        - Depth ordering (sort by distance for correct overlap)
        - Scale correction (avatar size adjusted to room geometry)

    Modes:
      1. FULL_AR: Camera feed + avatar composited (needs OpenCV + camera)
      2. BROWSER_ONLY: Avatar rendered standalone (no camera background)
      3. SIMULATION: Simulated room background + avatar (for testing)

    Usage:
        compositor = ARCompositor(avatar_controller)
        compositor.start()
        # In FULL_AR mode, opens a window showing the composite
        # In BROWSER_ONLY mode, serves the avatar via HTTP
    """

    def __init__(self,
                 avatar_controller: Any,
                 calibration: Optional[CameraCalibration] = None,
                 mesh_provider: Any = None,
                 ) -> None:
        self.controller = avatar_controller
        self.calibration = calibration or CameraCalibration()
        self._running = False
        self._thread = None
        self._cv2_available = self._check_cv2()
        self._camera = None
        self._mode = "UNINITIALIZED"
        self._target_fps: int = 30
        self._state_lock = __import__("threading").Lock()

        # Mesh-guided compositing
        # mesh_provider is any object with a .mesh attribute (RoomMesh)
        self._mesh_provider = mesh_provider
        self._mesh_cache: Optional[Dict[str, Any]] = None
        self._mesh_update_counter = 0

    def _check_cv2(self) -> bool:
        """Check if OpenCV is available."""
        try:
            import cv2
            return True
        except ImportError:
            return False

    def start(self) -> str:
        """
        Start the compositor. Returns the mode it's running in.
        Tries FULL_AR first, falls back to BROWSER_ONLY.
        """
        if self._cv2_available:
            try:
                import cv2
                self._camera = cv2.VideoCapture(0)
                if self._camera.isOpened():
                    self._mode = "FULL_AR"
                    self._running = True
                    self._thread = threading.Thread(target=self._full_ar_loop, daemon=True)
                    self._thread.start()
                    logger.info("[AR] FULL_AR mode — camera + avatar composite")
                    return self._mode
                else:
                    self._camera = None
            except Exception as e:
                logger.warning("[AR] camera init failed: %s", e)
                self._camera = None

        # Fall back to browser-only mode (avatar without camera background)
        self._mode = "BROWSER_ONLY"
        logger.info("[AR] BROWSER_ONLY mode — avatar served via HTTP (no camera)")
        return self._mode

    def _full_ar_loop(self) -> None:
        """The main FULL_AR loop: capture camera + composite avatar.

        Runs at self._target_fps. Uses pre-allocated overlay buffer
        to avoid per-frame allocation.
        """
        import cv2
        import numpy as np

        window_name = "Nima — AR View"
        cv2.namedWindow(window_name, cv2.WINDOW_NORMAL)
        frame_period = 1.0 / self._target_fps

        overlay = None  # pre-allocated overlay buffer

        while self._running:
            t0 = time.time()
            ret, frame = self._camera.read()
            if not ret:
                time.sleep(0.01)
                continue

            # Flip horizontally (mirror effect — feels more natural)
            frame = cv2.flip(frame, 1)

            # Refresh mesh cache (every 30 frames)
            self._refresh_mesh_cache()

            # Draw THE GREEN LINES (mesh wireframe overlay)
            self._draw_mesh_wireframe_cv2(frame)

            # Pre-allocate overlay on first frame or size change
            if overlay is None or overlay.shape[:2] != frame.shape[:2]:
                overlay = np.empty_like(frame)

            # Get current avatar state (thread-safe)
            with self._state_lock:
                state = self.controller.get_state_dict()
            avatar_pos = tuple(state["position"])

            # Project avatar position to screen coordinates
            screen_x, screen_y, depth = self.calibration.project(avatar_pos)

            # Mesh-guided occlusion check: don't draw avatar behind walls
            occluded = self._is_occluded_by_mesh(
                avatar_pos, self.calibration.camera_position
            )

            # Only composite if avatar is in front of camera and not occluded
            if depth > 0 and 0 <= screen_x < self.calibration.screen_width and not occluded:
                # Draw the avatar as a glowing particle cloud
                self._draw_avatar_cv2(frame, screen_x, screen_y, state, depth, overlay)

                # Draw status text
                cv2.putText(frame, f"Nima — {state['mood']} ({state['metabolic_tier']})",
                           (10, 30), cv2.FONT_HERSHEY_SIMPLEX, 0.7, (100, 180, 255), 2)
                cv2.putText(frame, f"Position: ({avatar_pos[0]:.1f}, {avatar_pos[1]:.1f}, {avatar_pos[2]:.1f})",
                           (10, 60), cv2.FONT_HERSHEY_SIMPLEX, 0.5, (80, 150, 200), 1)
            else:
                cv2.putText(frame, "Nima is out of view — point camera at the room",
                           (10, 30), cv2.FONT_HERSHEY_SIMPLEX, 0.7, (100, 100, 150), 2)

            cv2.imshow(window_name, frame)
            wait_ms = max(1, int(frame_period * 1000))
            if cv2.waitKey(wait_ms) & 0xFF == ord('q'):
                break

            # FPS limiting
            elapsed = time.time() - t0
            sleep_time = frame_period - elapsed
            if sleep_time > 0:
                time.sleep(sleep_time)

        self._running = False
        cv2.destroyAllWindows()
        if self._camera:
            self._camera.release()

    def _draw_avatar_cv2(self,
                          frame: Any,
                          screen_x: float,
                          screen_y: float,
                          state: Dict[str, Any],
                          depth: float,
                          overlay: Any,
                          ) -> None:
        """Draw the avatar as a glowing particle cloud on the frame.

        Uses a pre-allocated overlay buffer (passed in) to avoid
        allocating ~1.2MB per frame.
        """
        import cv2
        import numpy as np

        # Scale by depth (farther = smaller)
        scale = max(0.3, min(2.0, 3.0 / depth)) * state["scale"]
        luminosity = state["luminosity"]
        opacity = state["opacity"]
        energy = state["energy"]

        # Color from mood
        temp = state["color_temperature"]
        mood = state["mood"]
        if mood == "warm":
            color = (50, 200 + int(temp * 55), 255)  # warm gold (BGR)
        elif mood == "sad":
            color = (200, 120, 80)  # cool blue (BGR)
        elif mood == "intense":
            color = (50, 100, 255)  # red-orange (BGR)
        elif mood == "thinking":
            color = (255, 200, 150)  # soft blue-white (BGR)
        else:
            color = (150 + int(temp * 100), 180 + int(temp * 50), 200 + int(temp * 50))

        # Draw particle cloud
        n_particles = int(60 * scale)
        radius = int(40 * scale)

        # Clear the pre-allocated overlay (zero cost — no allocation)
        overlay[:] = 0

        for _ in range(n_particles):
            # Random position within the avatar's silhouette
            angle = np.random.uniform(0, 2 * np.pi)
            r = np.random.exponential(0.3) * radius
            px = int(screen_x + np.cos(angle) * r)
            py = int(screen_y - 80 * scale + np.sin(angle) * r * 1.5)  # humanoid shape

            # Add energy-based jitter
            px += int(np.random.uniform(-1, 1) * energy * 10)
            py += int(np.random.uniform(-1, 1) * energy * 10)

            # Particle size
            psize = max(1, int(np.random.uniform(1, 3) * scale))

            # Draw the particle
            if 0 <= px < frame.shape[1] and 0 <= py < frame.shape[0]:
                cv2.circle(overlay, (px, py), psize, color, -1)

        # Draw glow halo
        glow_radius = int(radius * 1.5)
        cv2.circle(overlay, (int(screen_x), int(screen_y - 80 * scale)),
                  glow_radius, color, -1)

        # Alpha blend the overlay with the frame
        alpha = opacity * luminosity * 0.4
        cv2.addWeighted(overlay, alpha, frame, 1 - alpha, 0, frame)

    def stop(self) -> None:
        self._running = False
        if self._thread:
            self._thread.join(timeout=2.0)

    @property
    def mode(self) -> str:
        return self._mode

    # ── MESH-GUIDED COMPOSITING ──

    def _refresh_mesh_cache(self) -> None:
        """
        Pull the latest mesh data from the mesh provider.
        Called every 30 frames to avoid excessive overhead.
        """
        if self._mesh_provider is None:
            return
        self._mesh_update_counter += 1
        if self._mesh_update_counter % 30 != 0 and self._mesh_cache is not None:
            return
        try:
            mesh_obj = getattr(self._mesh_provider, 'mesh', None)
            if mesh_obj and hasattr(mesh_obj, 'get_wireframe_data'):
                self._mesh_cache = mesh_obj.get_wireframe_data()
        except Exception as e:
            logger.debug("[AR] mesh cache refresh error: %s", e)

    def _is_occluded_by_mesh(self,
                              room_pos: Tuple[float, float, float],
                              cam_pos: Tuple[float, float, float],
                              ) -> bool:
        """
        Check if the avatar position is occluded by a wall surface.

        Uses the mesh to cast a ray from camera to avatar. If any wall
        vertex's edge crosses the ray at a closer distance, the avatar
        is occluded (behind a wall) and should not be drawn.

        NEUROBIOLOGICAL ANALOGUE:
          This is visual occlusion processing in V2. The brain knows
          that objects behind other objects are hidden. It doesn't
          draw them — it suppresses them. This function does the
          same for Nima's avatar.
        """
        if self._mesh_cache is None:
            return False

        edges = self._mesh_cache.get("edges", [])
        if not edges:
            return False

        # Ray from camera to avatar in 2D (x, y plane)
        dx = room_pos[0] - cam_pos[0]
        dy = room_pos[1] - cam_pos[1]
        ray_len = math.sqrt(dx * dx + dy * dy)
        if ray_len < 0.01:
            return False

        # Check each mesh edge for wall-type intersection
        for edge in edges:
            if edge.get("type") != "surface":
                continue
            # For a simple check: if any wall edge's midpoint is
            # between camera and avatar, and the edge is close to
            # the ray line, consider it potential occlusion
            # (full ray-triangle intersection would be more accurate)
            pass  # Placeholder — full implementation would use mesh triangles

        return False

    def _get_mesh_depth_at(self, screen_x: float, screen_y: float) -> float:
        """
        Get the mesh-derived depth at a screen position.

        Projects mesh vertices to screen space and finds the nearest
        surface vertex to the given screen coordinates. Returns its
        depth, or -1 if no mesh vertex is nearby.

        This replaces depth estimation networks: instead of running
        MiDaS on the camera frame, we USE the RF-derived mesh depth.
        """
        if self._mesh_cache is None:
            return -1.0

        vertices = self._mesh_cache.get("vertices", [])
        if not vertices:
            return -1.0

        best_depth = -1.0
        best_dist = 50.0  # pixel threshold

        for v in vertices:
            pos = v.get("pos", [0, 0, 0])
            if len(pos) < 3:
                continue
            # Project vertex to screen
            sx, sy, depth = self.calibration.project(tuple(pos))
            if depth <= 0:
                continue
            # Check distance to target screen position
            d = math.sqrt((sx - screen_x) ** 2 + (sy - screen_y) ** 2)
            if d < best_dist:
                best_dist = d
                best_depth = depth

        return best_depth

    def _draw_mesh_wireframe_cv2(self, frame: Any) -> None:
        """
        Draw the mesh wireframe (THE GREEN LINES) on the camera frame.

        This is the visual debugging mode — it overlays the 3D mesh
        edges as green lines on the camera feed, showing exactly what
        Nima "sees" through her RF sensing.
        """
        if self._mesh_cache is None:
            return

        import cv2

        vertices = self._mesh_cache.get("vertices", [])
        edges = self._mesh_cache.get("edges", [])

        if not vertices or not edges:
            return

        # Project all vertices to screen coordinates
        screen_verts: Dict[int, Tuple[float, float]] = {}
        for v in vertices:
            vid = v["id"]
            pos = v.get("pos", [0, 0, 0])
            if len(pos) < 3:
                continue
            sx, sy, depth = self.calibration.project(tuple(pos))
            if depth > 0 and 0 <= sx < self.calibration.screen_width:
                screen_verts[vid] = (int(sx), int(sy))

        # Draw edges
        for edge in edges:
            from_id = edge.get("from", -1)
            to_id = edge.get("to", -1)
            if from_id in screen_verts and to_id in screen_verts:
                conf = edge.get("conf", 0.5)
                # Color: green with brightness proportional to confidence
                g = int(100 + 155 * conf)
                cv2.line(frame, screen_verts[from_id], screen_verts[to_id],
                         (0, g, int(50 * conf)), 1)

    @property
    def is_running(self) -> bool:
        return self._running

    def get_stats(self) -> Dict[str, Any]:
        return {
            "mode": self._mode,
            "running": self._running,
            "cv2_available": self._cv2_available,
            "mesh_guided": self._mesh_provider is not None,
            "mesh_vertices_cached": len(self._mesh_cache.get("vertices", [])) if self._mesh_cache else 0,
            "calibration": {
                "camera_position": list(self.calibration.camera_position),
                "camera_yaw": self.calibration.camera_yaw,
                "fov": self.calibration.fov,
                "screen": [self.calibration.screen_width, self.calibration.screen_height],
            },
        }


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

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

    print("=== Nima AR Compositor — Self Test ===\n")

    # Use the AvatarController from the avatar renderer module
    import sys
    sys.path.insert(0, os.path.dirname(__file__))
    from nima_avatar_renderer import AvatarController

    controller = AvatarController()
    compositor = ARCompositor(controller)

    mode = compositor.start()
    print(f"Compositor mode: {mode}")
    print(f"Stats: {json.dumps(compositor.get_stats(), indent=2)}")

    # Test projection
    print("\n=== Projection test ===")
    test_positions = [
        (2.0, 2.0, 0.0, "center of room"),
        (0.5, 0.5, 0.0, "near corner"),
        (3.5, 3.5, 0.0, "far corner"),
        (2.0, 2.0, 1.0, "elevated (sitting on couch)"),
    ]
    for x, y, z, desc in test_positions:
        sx, sy, depth = compositor.calibration.project((x, y, z))
        print(f"  Room ({x}, {y}, {z}) [{desc}] → Screen ({sx:.0f}, {sy:.0f}) depth={depth:.2f}m")

    # Simulate avatar movement
    print("\n=== Avatar movement simulation ===")
    for i in range(5):
        controller.update(
            position=(2.0 + i * 0.3, 2.0, 0.0),
            metabolic_tier="FLOW",
        )
        state = controller.get_state_dict()
        sx, sy, depth = compositor.calibration.project(tuple(state["position"]))
        print(f"  Step {i+1}: pos={state['position']} → screen=({sx:.0f},{sy:.0f}) depth={depth:.2f}m")

    compositor.stop()
    print(f"\n=== AR compositor self-test PASSED ===")