spatial_core.py

from __future__ import annotations

from dataclasses import dataclass
from enum import Enum, auto
import json
import math
from pathlib import Path
from typing import Iterable

import numpy as np


class GestureEventType(Enum):
    CLICK = auto()
    SCROLL = auto()
    SWIPE_LEFT = auto()
    SWIPE_RIGHT = auto()
    TOGGLE_CONTROL = auto()


@dataclass(frozen=True)
class GestureEvent:
    kind: GestureEventType
    amount: int = 0


@dataclass(frozen=True)
class AppConfig:
    grid_size: int = 4
    target_margin: float = 0.08
    calibration_samples: int = 12
    smooth_alpha: float = 0.3
    pinch_threshold_px: float = 40
    release_threshold_px: float = 52
    gesture_axis_threshold_px: float = 18
    scroll_step_px: float = 8
    scroll_sensitivity: float = 0.4
    double_pinch_window_s: float = 0.4
    swipe_threshold_px: float = 120
    cursor_safe_margin_px: int = 8
    gaze_model_confidence: float = 0.45
    invert_scroll: bool = True
    invert_desktop_swipe: bool = True
    show_face_dots: bool = True


def scroll_amount_for_action(amount: int, config: AppConfig) -> int:
    return -amount if config.invert_scroll else amount


def desktop_direction_for_action(event: GestureEventType, config: AppConfig) -> str:
    if event not in (GestureEventType.SWIPE_LEFT, GestureEventType.SWIPE_RIGHT):
        raise ValueError("Desktop direction requires a swipe event")
    direction = "left" if event == GestureEventType.SWIPE_LEFT else "right"
    if config.invert_desktop_swipe:
        direction = "right" if direction == "left" else "left"
    return direction


def design_matrix(x: np.ndarray, y: np.ndarray) -> np.ndarray:
    return np.column_stack((np.ones_like(x), x, y, x * y, x**2, y**2))


def fit_poly(
    src: Iterable[tuple[float, float]], dst: Iterable[tuple[float, float]]
) -> tuple[list[float], list[float]]:
    source = np.asarray(list(src), dtype=float)
    destination = np.asarray(list(dst), dtype=float)
    if source.shape[0] < 6 or source.shape != destination.shape:
        raise ValueError("Polynomial calibration requires at least six matching points")
    matrix = design_matrix(source[:, 0], source[:, 1])
    theta_x, *_ = np.linalg.lstsq(matrix, destination[:, 0], rcond=None)
    theta_y, *_ = np.linalg.lstsq(matrix, destination[:, 1], rcond=None)
    return theta_x.tolist(), theta_y.tolist()


def apply_poly(
    theta: tuple[list[float], list[float]] | list[list[float]], x: float, y: float
) -> tuple[float, float]:
    theta_x, theta_y = map(np.asarray, theta)
    row = design_matrix(np.array([x]), np.array([y]))[0]
    return float(row @ theta_x), float(row @ theta_y)


def clamp_to_screen(
    x: float, y: float, width: int, height: int, margin: int = 0
) -> tuple[float, float]:
    margin = max(0, min(margin, (width - 1) // 2, (height - 1) // 2))
    return (
        min(max(x, margin), width - 1 - margin),
        min(max(y, margin), height - 1 - margin),
    )


def load_calibration(path: Path) -> tuple[list[float], list[float]] | None:
    try:
        data = json.loads(path.read_text())
        if (
            isinstance(data, list)
            and len(data) == 2
            and all(isinstance(axis, list) and len(axis) == 6 for axis in data)
            and all(math.isfinite(float(value)) for axis in data for value in axis)
        ):
            return data[0], data[1]
    except (OSError, ValueError, TypeError, json.JSONDecodeError):
        pass
    return None


class CalibrationSession:
    def __init__(self, screen_size: tuple[int, int], config: AppConfig):
        width, height = screen_size
        margin_x = width * config.target_margin
        margin_y = height * config.target_margin
        self.targets = [
            (
                margin_x + (width - 2 * margin_x) * column / (config.grid_size - 1),
                margin_y + (height - 2 * margin_y) * row / (config.grid_size - 1),
            )
            for row in range(config.grid_size)
            for column in range(config.grid_size)
        ]
        self.samples_per_target = config.calibration_samples
        self.sources: list[tuple[float, float]] = []
        self.destinations: list[tuple[float, float]] = []
        self.index = 0

    @property
    def complete(self) -> bool:
        return self.index >= len(self.targets)

    @property
    def current_target(self) -> tuple[float, float] | None:
        return None if self.complete else self.targets[self.index]

    def capture(self, gaze_samples: Iterable[tuple[float, float]]) -> bool:
        samples = list(gaze_samples)[-self.samples_per_target :]
        if self.complete or len(samples) < self.samples_per_target:
            return False
        x = float(np.median([sample[0] for sample in samples]))
        y = float(np.median([sample[1] for sample in samples]))
        self.sources.append((x, y))
        self.destinations.append(self.targets[self.index])
        self.index += 1
        return True

    def fit(self) -> tuple[tuple[list[float], list[float]], float]:
        if not self.complete:
            raise ValueError("Calibration is not complete")
        theta = fit_poly(self.sources, self.destinations)
        errors = [
            math.dist(apply_poly(theta, *source), destination)
            for source, destination in zip(self.sources, self.destinations)
        ]
        return theta, float(np.sqrt(np.mean(np.square(errors))))


class GestureRecognizer:
    """Turn pinch samples into mutually exclusive high-level gesture events."""

    def __init__(self, config: AppConfig):
        self.config = config
        self.pinching = False
        self.start: tuple[float, float] | None = None
        self.previous: tuple[float, float] | None = None
        self.axis: str | None = None
        self.moved = False
        self.pending_click_at: float | None = None
        self.toggle_on_release = False

    def reset(self) -> None:
        self.pinching = False
        self.start = None
        self.previous = None
        self.axis = None
        self.moved = False
        self.pending_click_at = None
        self.toggle_on_release = False

    def update(
        self, now: float, distance_px: float | None, point: tuple[float, float] | None
    ) -> list[GestureEvent]:
        events: list[GestureEvent] = []
        events.extend(self._flush_click(now))
        detected = (
            distance_px is not None
            and point is not None
            and distance_px
            < (
                self.config.release_threshold_px
                if self.pinching
                else self.config.pinch_threshold_px
            )
        )

        if detected and not self.pinching:
            self.pinching = True
            self.start = self.previous = point
            self.axis = None
            self.moved = False
            if self.pending_click_at is not None:
                self.pending_click_at = None
                self.toggle_on_release = True
            return events

        if detected and self.pinching and point is not None and self.start and self.previous:
            total_x = point[0] - self.start[0]
            total_y = point[1] - self.start[1]
            if self.axis is None and max(abs(total_x), abs(total_y)) >= self.config.gesture_axis_threshold_px:
                self.axis = "x" if abs(total_x) > abs(total_y) else "y"
                self.moved = True

            if self.axis == "x" and abs(total_x) >= self.config.swipe_threshold_px:
                kind = GestureEventType.SWIPE_RIGHT if total_x > 0 else GestureEventType.SWIPE_LEFT
                events.append(GestureEvent(kind))
                self.start = point
            elif self.axis == "y":
                delta = self.previous[1] - point[1]
                amount = int(delta * self.config.scroll_sensitivity)
                if abs(delta) >= self.config.scroll_step_px and amount:
                    events.append(GestureEvent(GestureEventType.SCROLL, amount))
                    self.previous = point
            else:
                self.previous = point
            return events

        if not detected and self.pinching:
            self.pinching = False
            self.start = self.previous = None
            self.axis = None
            if self.toggle_on_release:
                self.toggle_on_release = False
                if not self.moved:
                    events.append(GestureEvent(GestureEventType.TOGGLE_CONTROL))
            elif not self.moved:
                self.pending_click_at = now + self.config.double_pinch_window_s
            self.moved = False
        return events

    def _flush_click(self, now: float) -> list[GestureEvent]:
        if self.pending_click_at is not None and now >= self.pending_click_at:
            self.pending_click_at = None
            return [GestureEvent(GestureEventType.CLICK)]
        return []