camera.py

"""
Camera abstraction: real Pi Camera and simulation modes.
"""

import time
import numpy as np
from abc import ABC, abstractmethod
from typing import Optional, Tuple

# Pi Camera is imported lazily to allow simulation on non-Pi machines


class CameraBase(ABC):
    """Abstract camera interface."""

    @abstractmethod
    def capture(self) -> Optional[np.ndarray]:
        """Capture a grayscale frame. Returns None on failure."""
        ...

    @abstractmethod
    def release(self):
        """Release camera resources."""
        ...

    @property
    @abstractmethod
    def resolution(self) -> Tuple[int, int]:
        """(width, height) of captured frames."""
        ...


class PiCamera(CameraBase):
    """Real Pi Camera using picamera2 (Bookworm+)."""

    def __init__(self, resolution: Tuple[int, int] = (640, 480), fps: int = 10):
        self._resolution = resolution
        self._fps = fps
        self._camera = None

        try:
            from picamera2 import Picamera2
            self._camera = Picamera2()
            config = self._camera.create_preview_configuration(
                main={"size": resolution, "format": "YUV420"}
            )
            self._camera.configure(config)
            self._camera.start()
            time.sleep(0.5)  # Warmup
        except ImportError:
            raise RuntimeError(
                "picamera2 not installed. On Pi: pip install picamera2\n"
                "On desktop: use SimCamera for testing."
            )

    def capture(self) -> Optional[np.ndarray]:
        if self._camera is None:
            return None
        frame = self._camera.capture_array("main")
        # YUV420 -> grayscale
        y = frame[: self._resolution[1], : self._resolution[0]]
        return np.asarray(y, dtype=np.uint8)

    def release(self):
        if self._camera:
            self._camera.stop()

    @property
    def resolution(self) -> Tuple[int, int]:
        return self._resolution


class USBCamera(CameraBase):
    """USB webcam using OpenCV VideoCapture. Works on macOS, Linux, Windows."""

    def __init__(self,
                 device: int = 0,
                 resolution: Tuple[int, int] = (640, 480),
                 api_preference: int = None):
        """
        Args:
            device: Camera index (0 = built-in, 1 = first USB, etc.)
            resolution: Desired capture resolution
            api_preference: cv2.CAP_AVFOUNDATION (macOS), cv2.CAP_V4L2 (Linux), etc.
                            Auto-detected if None.
        """
        import cv2
        self._resolution = resolution

        if api_preference is None:
            # Auto-detect best backend
            import platform
            system = platform.system()
            if system == "Darwin":
                api_preference = cv2.CAP_AVFOUNDATION
            elif system == "Linux":
                api_preference = cv2.CAP_V4L2
            else:
                api_preference = cv2.CAP_DSHOW  # Windows

        self._cap = cv2.VideoCapture(device, api_preference)
        if not self._cap.isOpened():
            raise RuntimeError(
                f"Cannot open camera {device}. Try a different device number (0-3).\n"
                f"On macOS, check System Preferences → Privacy → Camera."
            )

        self._cap.set(cv2.CAP_PROP_FRAME_WIDTH, resolution[0])
        self._cap.set(cv2.CAP_PROP_FRAME_HEIGHT, resolution[1])
        self._cap.set(cv2.CAP_PROP_FPS, 30)

        # Warmup: discard first few frames (auto-exposure settling)
        for _ in range(5):
            self._cap.read()

        # Read actual resolution (camera may not support requested)
        actual_w = int(self._cap.get(cv2.CAP_PROP_FRAME_WIDTH))
        actual_h = int(self._cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
        if (actual_w, actual_h) != resolution:
            print(f"  Camera actual resolution: {actual_w}x{actual_h} (requested {resolution[0]}x{resolution[1]})")
            self._resolution = (actual_w, actual_h)

    def capture(self) -> Optional[np.ndarray]:
        if self._cap is None:
            return None
        ret, frame = self._cap.read()
        if not ret:
            return None
        import cv2
        gray = cv2.cvtColor(frame, cv2.COLOR_BGR2GRAY)
        return gray

    def release(self):
        if self._cap:
            self._cap.release()
            self._cap = None

    @property
    def resolution(self) -> Tuple[int, int]:
        return self._resolution


class SimCamera(CameraBase):
    """
    Simulated camera for testing on desktop.
    Flies a virtual drone over the map image.
    """

    def __init__(self,
                 map_path: str,
                 resolution: Tuple[int, int] = (640, 480),
                 altitude_m: float = 50.0,
                 gsd_cm_per_px: float = 3.0):
        import cv2
        self._map = cv2.imread(map_path)
        if self._map is None:
            raise FileNotFoundError(f"Cannot read map: {map_path}")

        self._map_gray = cv2.cvtColor(self._map, cv2.COLOR_BGR2GRAY)
        self._map_h, self._map_w = self._map_gray.shape
        self._resolution = resolution
        self._altitude = altitude_m
        self._gsd = gsd_cm_per_px

        # Drone position on map (center of image)
        self._drone_x = self._map_w / 2
        self._drone_y = self._map_h / 2
        self._heading = 0.0  # degrees, 0=up/north

        # Waypoints for figure-8 pattern
        self._t = 0.0

    def _simulated_position(self) -> Tuple[float, float, float]:
        """
        Generate a figure-8 flight path over the map.
        Returns (px, py, heading_deg).
        """
        self._t += 0.02

        # Figure-8: x = sin(t), y = sin(2t)/2
        amp_x = self._map_w * 0.35
        amp_y = self._map_h * 0.2

        offset_x = self._map_w / 2 + np.sin(self._t) * amp_x
        offset_y = self._map_h / 2 + np.sin(2 * self._t) * amp_y / 2

        # Heading follows the curve tangent
        dx = np.cos(self._t) * amp_x
        dy = np.cos(2 * self._t) * amp_y
        heading = np.degrees(np.arctan2(dx, dy)) % 360

        return offset_x, offset_y, heading

    def capture(self) -> Optional[np.ndarray]:
        self._drone_x, self._drone_y, self._heading = self._simulated_position()

        # Extract viewport from map
        viewport_m = self._altitude * self._gsd / 100.0  # viewport width in meters
        px_per_m = self._map_w / (self._map_w * self._gsd / 100.0)  # approximate
        viewport_px = int(viewport_m * px_per_m)
        viewport_px = min(viewport_px, self._map_w, self._map_h)
        half = viewport_px // 2

        x0 = int(max(0, self._drone_x - half))
        y0 = int(max(0, self._drone_y - half))
        x1 = int(min(self._map_w, self._drone_x + half))
        y1 = int(min(self._map_h, self._drone_y + half))

        if x1 <= x0 or y1 <= y0:
            return np.zeros((self._resolution[1], self._resolution[0]), dtype=np.uint8)

        crop = self._map_gray[y0:y1, x0:x1]

        # Add slight blur to simulate real camera
        import cv2
        crop = cv2.GaussianBlur(crop, (3, 3), 0.5)

        # Resize to camera resolution
        frame = cv2.resize(crop, self._resolution)

        # Add small noise
        noise = np.random.normal(0, 2, frame.shape).astype(np.uint8)
        frame = cv2.add(frame, noise)

        return frame

    @property
    def ground_truth(self) -> dict:
        """Return true position for evaluation."""
        return {
            "x_px": self._drone_x,
            "y_px": self._drone_y,
            "heading_deg": self._heading,
        }

    def release(self):
        pass

    @property
    def resolution(self) -> Tuple[int, int]:
        return self._resolution