orisys-sdk/src/orisys/finger/runtime.py

from __future__ import annotations

from dataclasses import dataclass
from pathlib import Path

import numpy as np

from ..capture import (
    CaptureConfig,
    CaptureSource,
    ReconnectPolicy,
    is_live_camera,
    normalize_source,
    open_capture,
)
from ..configs import load_config
from ..field_calibration import FIELD_IDENTITY, FIELD_PROFILE
from ..frame_processors import FrameUnwarper
from ..sensor import Sensor
from ..unroll import (
    build_or_load_domain,
    load_calib_json,
    resolve_uv_runtime_sizes,
    scale_camera_geometry,
)
from .exports import (
    DEFAULT_FINGER_EXPORT_DIR,
    DEFAULT_FINGER_MAP_CACHE,
    FingerExportPaths,
    load_finger_export_bundle,
)


def normalize_video_source(video):
    """Accept camera index, numeric string, or filesystem video path."""
    return normalize_source(video)


def open_source(video, config, *, auto_reconnect: bool = False, verbose: bool = False):
    """Open a capture source for finger pipelines (compat: returns ``.cap``-like handle).

    Historically returned a bare ``cv2.VideoCapture``. Callers that need the raw
    OpenCV handle can use ``source.cap``; prefer ``source.read()``.
    """
    capture_config = CaptureConfig.from_camera_dict(
        normalize_video_source(video),
        config.get("camera", {}),
        reconnect=ReconnectPolicy(enabled=auto_reconnect),
        verbose=verbose,
        background=False,
    )
    source = open_capture(capture_config, start_background=False)
    cap = source.cap
    cap._orisys_worker = source  # type: ignore[attr-defined]
    return cap


class TactileProcessor:
    """Narrow tactile-analysis wrapper around ``Sensor`` external-frame mode."""

    def __init__(self, sensor: Sensor, *, default_valid_mask=None):
        self._sensor = sensor
        self._default_valid_mask = default_valid_mask

    @property
    def raw_sensor(self) -> Sensor:
        """Escape hatch for legacy callers that still need the full ``Sensor``."""
        return self._sensor

    @property
    def info(self):
        return self._sensor.info

    @property
    def frame(self):
        return self._sensor.frame

    @property
    def img(self):
        return self._sensor.img

    @property
    def flow(self):
        return self._sensor.flow

    @property
    def vnormal(self):
        return self._sensor.vnormal

    @property
    def vshear(self):
        return self._sensor.vshear

    @property
    def depth_map(self):
        return self._sensor.depth_map

    @property
    def mask_stats(self):
        return self._sensor.mask_stats

    @property
    def valid_mask(self):
        return self._sensor.valid_mask

    @property
    def force_units(self) -> tuple[str, str, str]:
        return self._sensor.force_units

    @property
    def calibration_status(self):
        return self._sensor.calibration_status

    @property
    def calibration_result(self):
        return self._sensor.calibration_result

    def set_calibration_reference(self, raw, *, reference_id):
        return self._sensor.set_calibration_reference(raw, reference_id=reference_id)

    def push_frame(self, frame, *, already_gray=False, valid_mask=None):
        """Feed one processed tactile frame into the analysis pipeline."""
        mask = self._default_valid_mask if valid_mask is None else valid_mask
        return self._sensor.get_img(
            frame=frame,
            already_processed=True,
            already_gray=already_gray,
            valid_mask=mask,
        )

    def warmup(self):
        return self._sensor.warmup_deformation()

    def pipeline(self, produce, **kwargs):
        """Start a :class:`~orisys.pipeline.StreamPipeline` over this processor's own Sensor.

        ``produce`` is the frame source: a zero-argument callable returning the next unrolled
        frame, or ``None`` at end of input. :class:`FingerRuntime` supplies it through
        :meth:`FingerRuntime.pipeline`, so call this directly only when you are producing frames
        yourself.

        Any keyword argument is forwarded to ``StreamPipeline``, so ``contact=True`` and
        ``include=(...)`` work exactly as they do there.
        """
        from orisys.pipeline import StreamPipeline

        return StreamPipeline(self._sensor, produce, **kwargs)

    def compute_deformation(self, *args, **kwargs):
        return self._sensor.compute_deformation(*args, **kwargs)

    def compute_contact(self, *args, **kwargs):
        return self._sensor.compute_contact(*args, **kwargs)

    def read_info(self, *args):
        return self._sensor.read_info(*args)

    def reset(self):
        return self._sensor.reset()

    def disconnect(self):
        return self._sensor.disconnect()


@dataclass(frozen=True)
class FingerGeometry:
    """Resolved one-time geometry/bootstrap state for a finger UV runtime."""

    config: dict
    camera_config: dict
    export_paths: FingerExportPaths
    data: dict
    source_K: np.ndarray
    dist_coeffs: np.ndarray | None
    native_width: int
    native_height: int
    source_width: int
    source_height: int
    out_width: int
    out_height: int
    source_scale: float
    map_cache: str
    edge_expand: float


def resolve_finger_geometry(
    *,
    export_dir,
    config="finger",
    map_cache=DEFAULT_FINGER_MAP_CACHE,
    source_scale=None,
    calib=None,
) -> FingerGeometry:
    """Resolve config/export/camera geometry for a UV finger runtime."""
    config = load_config(config)
    camera_config = config.get("camera", {})
    out_width, out_height, source_scale = resolve_uv_runtime_sizes(
        camera_config,
        source_scale=source_scale,
    )
    bundle = load_finger_export_bundle(
        export_dir,
        map_cache=map_cache,
        source_scale=source_scale,
    )

    if calib:
        source_K, dist_coeffs, img_size = load_calib_json(calib)
        native_width, native_height = img_size
    else:
        source_K = bundle.data["K"]
        dist_coeffs = None
        native_width = int(bundle.data["width"])
        native_height = int(bundle.data["height"])

    source_K, source_width, source_height = scale_camera_geometry(
        source_K,
        native_width,
        native_height,
        source_scale,
    )

    return FingerGeometry(
        config=config,
        camera_config=camera_config,
        export_paths=bundle.export_paths,
        data=bundle.data,
        source_K=source_K,
        dist_coeffs=dist_coeffs,
        native_width=native_width,
        native_height=native_height,
        source_width=source_width,
        source_height=source_height,
        out_width=out_width,
        out_height=out_height,
        source_scale=float(source_scale),
        map_cache=bundle.map_cache,
        edge_expand=bundle.edge_expand,
    )


class FingerRuntime:
    """Reusable UV finger runtime composed from capture, UV remap, and tactile analysis."""

    def __init__(
        self,
        *,
        geometry: FingerGeometry,
        domain,
        tactile: TactileProcessor,
        frame_processor: FrameUnwarper,
        capture,
        capture_source: CaptureSource | None = None,
    ):
        self.geometry = geometry
        self.domain = domain
        self.tactile = tactile
        self.sensor = tactile
        self.frame_processor = frame_processor
        self.capture = capture
        self.capture_source = capture_source
        # Deprecated alias.
        self.capture_worker = capture_source

    @classmethod
    def open(
        cls,
        *,
        export_dir=DEFAULT_FINGER_EXPORT_DIR,
        config="finger",
        map_cache=DEFAULT_FINGER_MAP_CACHE,
        source_scale=None,
        calib=None,
        cal_path=None,
        calibration_dir=None,
        device_id=None,
        video=0,
        camera_id=None,
        flip_source_vertical=True,
        verbose=False,
        mirror=False,
        auto_reconnect: bool = False,
        capture_mode: str = "latest",
    ) -> FingerRuntime:
        """Open a live finger runtime from export bundle, config, and video source."""
        del capture_mode  # deprecated; latest-frame capture only
        if camera_id is not None:
            video = camera_id
        video = normalize_video_source(video)
        geometry = resolve_finger_geometry(
            export_dir=export_dir,
            config=config,
            map_cache=map_cache,
            source_scale=source_scale,
            calib=calib,
        )

        mesh_path, intrinsics_path, pose_path = geometry.export_paths.as_tuple()
        domain = build_or_load_domain(
            mesh_path,
            intrinsics_path,
            pose_path,
            geometry.map_cache,
            geometry.source_K,
            geometry.source_width,
            geometry.source_height,
            geometry.out_width,
            geometry.out_height,
            edge_expand=geometry.edge_expand,
            include_surface=True,
            verbose=verbose,
            source_scale=geometry.source_scale,
        )
        domain.remap.ensure_fast_maps()

        sensor = Sensor(
            "external", config_name=config, cal_path=cal_path,
            calibration_dir=calibration_dir, device_id=device_id,
            verbose=verbose, mirror=mirror,
        )
        if getattr(getattr(sensor, "calibration", None), "schema_version", None) == 3:
            from ..calibration_context import (
                build_context,
                domain_fingerprint,
                sha256_file,
            )

            model = sensor.calibration
            uv_map = getattr(model, "coordinate_map", None)
            if model.input_profile in ("fn_delta_autocal_uv/1", FIELD_PROFILE):
                if uv_map is None or uv_map.shape[:2] != (geometry.out_height, geometry.out_width):
                    raise ValueError("calibration coordinate map does not match finger output shape")
            field_identity = FIELD_IDENTITY if model.input_profile == FIELD_PROFILE else None
            sensor.bind_calibration_context(build_context(
                sensor._config, sensor.filters_state(),
                geometry_sha256=sha256_file(mesh_path), source_scale=geometry.source_scale,
                coordinates=model.pipeline.get("coordinates"),
                output_domain_sha256=domain_fingerprint(domain),
                processing_backend="uv_unroll",
                preprocessing={"flip_source_vertical": flip_source_vertical, "mirror": mirror},
                field=field_identity,
            ))

        capture_config = CaptureConfig.from_camera_dict(
            video,
            geometry.config.get("camera", {}),
            reconnect=ReconnectPolicy(enabled=auto_reconnect),
            verbose=verbose,
            background=False,
        )
        capture_source = open_capture(capture_config, start_background=False)
        if not capture_source.is_open():
            raise RuntimeError("Cannot open camera/video source")
        if is_live_camera(video):
            capture_source.start_background()
        capture = capture_source.cap

        if not is_live_camera(video):
            # File playback: wall-clock cadence is not signal time, so pin the
            # force-filter design rate to the configured camera fps.
            try:
                fps = float(geometry.config.get("camera", {}).get("fps", 0.0) or 0.0)
            except (TypeError, ValueError):
                fps = 0.0
            if fps > 0:
                sensor.configure(lowpass_sample_rate_hz=fps)
        frame_processor = FrameUnwarper(
            domain.remap,
            camera_matrix=geometry.source_K if geometry.dist_coeffs is not None else None,
            dist_coeffs=geometry.dist_coeffs,
            source_width=geometry.source_width,
            source_height=geometry.source_height,
            flip_source_vertical=flip_source_vertical,
        )
        tactile = TactileProcessor(sensor, default_valid_mask=frame_processor.valid_mask)
        return cls(
            geometry=geometry,
            domain=domain,
            tactile=tactile,
            frame_processor=frame_processor,
            capture=capture,
            capture_source=capture_source,
        )

    @property
    def valid_mask(self):
        return self.frame_processor.valid_mask

    def set_calibration_reference(self, raw, *, reference_id):
        """Set an explicitly measured unloaded reference for this finger stream."""
        self.tactile.raw_sensor.set_calibration_reference(raw, reference_id=reference_id)

    def set_calibration(self, calibration, *, device_id=None):
        """Replace the model at a caller-owned frame boundary, validating actual geometry."""
        from .. import calibration as cal_mod
        from ..calibration_context import build_context, domain_fingerprint, sha256_file

        sensor = self.tactile.raw_sensor
        selected = cal_mod.load(calibration, strict=True) if isinstance(calibration, (str, Path)) else calibration
        serial = device_id or sensor.device_id
        if selected is not None and serial:
            selected = cal_mod.validate_for_device(selected, "finger", serial)
        context = None
        if getattr(selected, "schema_version", None) == 3:
            if selected.input_profile in ("fn_delta_autocal_uv/1", FIELD_PROFILE):
                uv_map = selected.coordinate_map
                if uv_map is None or uv_map.shape[:2] != (self.geometry.out_height, self.geometry.out_width):
                    raise ValueError("calibration coordinate map does not match finger output shape")
            field_identity = FIELD_IDENTITY if selected.input_profile == FIELD_PROFILE else None
            context = build_context(
                self.geometry.config, sensor.filters_state(),
                geometry_sha256=sha256_file(self.geometry.export_paths.mesh),
                source_scale=self.geometry.source_scale, coordinates=selected.pipeline.get("coordinates"),
                output_domain_sha256=domain_fingerprint(self.domain), processing_backend="uv_unroll",
                preprocessing={"flip_source_vertical": self.frame_processor.flip_source_vertical,
                               "mirror": sensor.mirror},
                field=field_identity)
        sensor.set_calibration(selected, context=context)

    def declare_unloaded_reference(self, reference_id):
        """Declare the unloaded window for a field model; the SDK medians 30 weighted frames."""
        self.tactile.raw_sensor.declare_unloaded_reference(reference_id=reference_id)

    @property
    def calibration_status(self):
        """Copy the current per-axis calibration validity state."""
        return dict(self.tactile.raw_sensor.calibration_status)

    def read_source(self):
        if self.capture_source is not None and self.capture_source.running:
            return self.capture_source.read(copy=True, wait_s=1.0)
        if self.capture_source is not None:
            return self.capture_source.read_raw()
        return self.capture.read()

    def read_source_stamped(self):
        """Read the source image with its capture-thread identity and host clocks."""
        if self.capture_source is None:
            return None
        if self.capture_source.running:
            return self.capture_source.read_latest_stamped(copy=True, wait_s=1.0)
        return self.capture_source.read_raw_stamped()

    def process_frame(self, frame, *, color=False):
        """Prepare, unroll, and feed one source frame into the tactile processor."""
        if color:
            prepared = self.frame_processor.prepare_source(frame)
            unrolled = self.frame_processor.flatten_prepared(prepared)
            self.tactile.push_frame(unrolled)
            return unrolled

        prepared_gray = self.frame_processor.prepare_source_gray(frame)
        unrolled = self.frame_processor.flatten_prepared_gray(prepared_gray)
        self.tactile.push_frame(unrolled, already_gray=True)
        return unrolled

    # ------------------------------------------- session measurement seam
    def enable_measurement_session(self) -> None:
        """Hand force readout to a ``ForceSession`` (the explicit new entry)."""
        self.tactile.raw_sensor.enable_measurement_session()

    def disable_measurement_session(self) -> None:
        """Return to the historical automatic model/readout path."""
        self.tactile.raw_sensor.disable_measurement_session()

    def reset_tracker(self) -> None:
        """Reset the tactile tracker at a frame boundary (session zeroing half)."""
        self.tactile.raw_sensor.reset_tracker()

    def read_measurement_frame(
        self,
        source_frame,
        *,
        frame_id: int,
        timestamp_s: float | None = None,
        want_field: bool = False,
        want_contact: bool = False,
    ):
        """Unroll one source frame and measure it once, without calibration."""
        del timestamp_s  # signal time is carried by the caller's FrameStamp
        self.process_frame(source_frame)
        return self.tactile.raw_sensor.measure_current_frame(
            frame_id=frame_id,
            want_field=want_field,
            want_contact=want_contact,
        )

    def current_measurement(
        self, *, want_field: bool = False, want_contact: bool = False
    ):
        """Same-frame measurement features already produced by the pipeline (no recompute).

        Additive: forwards to the tactile sensor's session seam so a pipeline
        consumer can hand the frame to ``ForceSession`` without a second
        unroll/flow/HHD pass.
        """
        return self.tactile.raw_sensor.current_measurement(
            want_field=want_field, want_contact=want_contact
        )

    @property
    def calibration_context(self):
        """The pipeline context the open model was validated against, if any."""
        return self.tactile.raw_sensor.calibration_context

    @property
    def calibration(self):
        """The force calibration the open runtime holds (additive read-only accessor)."""
        return getattr(self.tactile.raw_sensor, "calibration", None)

    def grab_and_process(self, *, color=False):
        ok, frame = self.read_source()
        if not ok:
            return False, None
        return True, self.process_frame(frame, color=color)

    def get_img(self, *, color=False):
        """Sensor-style alias: capture, unroll, and return the tactile frame or None."""
        ok, unrolled = self.grab_and_process(color=color)
        return unrolled if ok else None

    def warmup(self, *, color=False):
        ok, unrolled = self.grab_and_process(color=color)
        if ok:
            self.tactile.warmup()
        return ok, unrolled

    def pipeline(self, **kwargs):
        """Return a :class:`~orisys.pipeline.StreamPipeline` for this runtime.

        The recommended way to run a finger sensor. Each finished frame arrives as a single
        :class:`~orisys.pipeline.FrameSnapshot`, so there are no stage calls to sequence and no
        tuple to unpack::

            with runtime.pipeline(contact=True, include=("image", "depth")) as pipe:
                while True:
                    snap = pipe.wait()
                    if snap is None:
                        break
                    print(snap.normal_force, snap.shear_x, snap.shear_y)

        Any keyword argument is forwarded to ``StreamPipeline``. Use ``contact=True`` to run
        contact analysis, and ``include=(...)`` to select which arrays each snapshot carries.
        """
        from orisys.pipeline import StreamPipeline

        return StreamPipeline.for_finger_runtime(self, **kwargs)

    def close(self):
        if self.capture_source is not None:
            self.capture_source.close()
            self.capture_source = None
            self.capture_worker = None
            self.capture = None
        elif self.capture is not None:
            self.capture.release()
            self.capture = None
        self.tactile.disconnect()