orisys-sdk/tests/test_force_stream_session.py

"""New force-side session (T06a): version/zero contract and old-session numerics.

The tests drive :class:`~orisys.force_stream_session.ForceStreamSession` with
T05-style measurements (T03 ``MeasurementFrame`` + same-frame raw feature) and
compare the scalar/uncalibrated numerical chain with the published old
``ForceSession`` on the same input sequence. They also pin the new contract:
reference/session changes invalidate the zero in every phase, only same-version
frames fill the baseline window, and duplicate/late/same-timestamp frames stay
diagnostics-only.
"""

from __future__ import annotations

import importlib.util
from pathlib import Path

import numpy as np
import pytest

from orisys.force_stream_session import ForceStreamOutput, ForceStreamSession
from orisys.image_pipeline import ImageMeasurement
from orisys.stream_contract import FrameStamp, MeasurementFrame

FRAME_NS = 33_333_333


def _load_helpers():
    path = Path(__file__).with_name("test_force_session.py")
    spec = importlib.util.spec_from_file_location("_force_session_test_helpers", path)
    module = importlib.util.module_from_spec(spec)
    assert spec.loader is not None
    spec.loader.exec_module(module)
    return module


helpers = _load_helpers()


def _measurement(
    frame_id: int,
    raw: tuple[float, float, float],
    *,
    version: int = 0,
    session_id: str = "s1",
    sensor_id: str = "TEST-1",
    capture_time_ns: int | None = None,
    source_frame_id: int | None = None,
    field: dict | None = None,
    valid_mask: np.ndarray | None = None,
) -> ImageMeasurement:
    stamp = FrameStamp(
        sensor_id=sensor_id,
        session_id=session_id,
        frame_id=frame_id,
        source_frame_id=frame_id if source_frame_id is None else source_frame_id,
        capture_time_ns=(
            frame_id * FRAME_NS if capture_time_ns is None else capture_time_ns
        ),
        clock_id="host",
        capture_time_basis="recorded_signal",
    )
    frame = MeasurementFrame(
        stamp=stamp,
        source_dt_s=None if frame_id == 1 else FRAME_NS / 1e9,
        reference_generation=version,
    )
    return ImageMeasurement(
        frame=frame,
        force_raw=raw,
        torque=0.0,
        flow=None,
        field=field,
        valid_mask=valid_mask,
    )


def _raw_for(frame_id: int) -> tuple[float, float, float]:
    return (
        15000.0 + 500.0 * frame_id,
        8000.0 + 200.0 * frame_id,
        7000.0 - 100.0 * frame_id,
    )


def _old_chain(calibration=None, settings=None):
    source = helpers.FakeMeasurementSource(_raw_for)
    session = helpers.make_session(source, calibration, settings=settings)
    session.request_zero(reference_id="old-ref", declared_unloaded=True)
    return session, source


def _new_chain(calibration=None, settings=None):
    return ForceStreamSession(
        sensor_family="box",
        calibration=calibration,
        settings=settings if settings is not None else helpers.legacy_settings(),
    )


def _compare_with_old(calibration=None, settings=None) -> None:
    old, source = _old_chain(calibration, settings)
    new = _new_chain(calibration, settings)
    new.request_zero(
        reference_id="new-ref", declared_unloaded=True, expected_reference_generation=0
    )
    old_raw, old_newton, old_raw_status, old_newton_status = [], [], [], []
    new_raw, new_newton, new_raw_status, new_newton_status = [], [], [], []
    new_zero_phases = []
    for frame_id in range(1, 41):
        old_output = helpers.process(old, source, frame_id)
        old_raw.append(np.asarray(old_output.raw))
        old_newton.append(np.asarray(old_output.newton))
        old_raw_status.append(list(old_output.measurement.sdk.raw_status))
        old_newton_status.append(list(old_output.measurement.newton_status))

        result = new.process(_measurement(frame_id, _raw_for(frame_id)))
        assert isinstance(result, ForceStreamOutput)
        new_raw.append(np.asarray(result.raw))
        new_newton.append(np.asarray(result.newton))
        new_raw_status.append(list(result.raw_status))
        new_newton_status.append(list(result.newton_status))
        new_zero_phases.append(result.zero.phase)

    assert np.array_equal(np.stack(old_raw), np.stack(new_raw), equal_nan=True)
    assert np.array_equal(np.stack(old_newton), np.stack(new_newton), equal_nan=True)
    assert old_raw_status == new_raw_status
    assert old_newton_status == new_newton_status
    assert new.zero.phase == "ready"
    assert new_zero_phases[:30] == ["collecting"] * 30
    assert new_zero_phases[30] == "ready"


# ------------------------------------------------------------ numeric equivalence
def test_scalar_curve_session_matches_old_force_session() -> None:
    _compare_with_old(helpers.box_cal())


def test_uncalibrated_session_matches_old_force_session() -> None:
    _compare_with_old(None)


def test_filtered_scalar_session_matches_old_force_session() -> None:
    settings = helpers.OutputSettings(
        sample_rate_hz=30.0, median_window=3, lowpass_hz=5.0
    )
    _compare_with_old(helpers.box_cal(), settings)


# ----------------------------------------------------- version / session contract
def _baseline(new: ForceStreamSession, frame_ids, *, version=0, session_id="s1"):
    return [
        new.process(
            _measurement(fid, _raw_for(fid), version=version, session_id=session_id)
        )
        for fid in frame_ids
    ]


def test_version_change_during_collecting_invalidates_the_zero() -> None:
    new = _new_chain(helpers.box_cal())
    new.request_zero(
        reference_id="ref", declared_unloaded=True, expected_reference_generation=0
    )
    collected = _baseline(new, range(1, 11))
    assert collected[-1].zero.phase == "collecting"

    changed = new.process(_measurement(11, _raw_for(11), version=2))
    assert changed.quality == "version_changed"
    assert changed.zero.phase == "failed"
    assert "reference generation changed" in changed.zero.failure_reason
    assert all(np.isnan(value) for value in changed.newton)

    # Only an explicit re-zero for the new version restores a valid N.
    new.request_zero(
        reference_id="ref-2", declared_unloaded=True, expected_reference_generation=2
    )
    ready = _baseline(new, range(12, 43), version=2)
    assert ready[-1].zero.phase == "ready"
    assert np.isfinite(ready[-1].newton[0])
    assert ready[-1].reference_generation == 2


def test_version_change_after_ready_invalidates_and_requires_explicit_reset() -> None:
    new = _new_chain()
    new.request_zero(
        reference_id="ref", declared_unloaded=True, expected_reference_generation=0
    )
    baseline = _baseline(new, range(1, 31))
    ready = _baseline(new, [31])
    assert ready[-1].zero.phase == "ready"
    assert np.isfinite(baseline[-1].raw[0])

    changed = new.process(_measurement(32, _raw_for(32), version=1))
    assert changed.quality == "version_changed"
    assert changed.zero.phase == "failed"
    assert all(np.isnan(value) for value in changed.newton)
    # A raw row is still in the result stream; the zero is not silently relabelled.
    assert all(np.isfinite(value) for value in changed.raw)

    still_failed = new.process(_measurement(33, _raw_for(33), version=1))
    assert still_failed.zero.phase == "failed"
    assert all(np.isnan(value) for value in still_failed.newton)


def test_session_change_invalidates_the_zero_and_emits_results() -> None:
    new = _new_chain()
    new.request_zero(
        reference_id="ref", declared_unloaded=True, expected_reference_generation=0
    )
    _baseline(new, range(1, 31))
    assert _baseline(new, [31])[-1].zero.phase == "ready"

    changed = new.process(_measurement(32, _raw_for(32), session_id="s2", version=0))
    assert changed.quality == "session_changed"
    assert changed.zero.phase == "failed"
    assert "source session changed" in changed.zero.failure_reason
    assert all(np.isfinite(value) for value in changed.raw)


def test_late_frame_cannot_fill_the_new_version_window() -> None:
    # Corrected in the Y2 rework: a frame older than the requested reset target is
    # "zero_pending", not "late" (it belongs to the still-active version and is
    # waiting for the reset confirmation). It still cannot fill the window.
    new = _new_chain()
    new.request_zero(
        reference_id="ref", declared_unloaded=True, expected_reference_generation=1
    )
    late = new.process(_measurement(1, _raw_for(1), version=0))
    assert late.quality == "zero_pending"
    assert late.late is False
    assert late.accepted is False
    assert late.zero.phase == "resetting", "pending frame must not advance the zero"

    accepted = _baseline(new, range(2, 12), version=1)
    assert accepted[-1].zero.phase == "collecting"


# ------------------------------------------- duplicate / late / same timestamp
def test_duplicate_and_same_timestamp_frames_are_diagnostics_only() -> None:
    new = _new_chain()
    new.request_zero(
        reference_id="ref", declared_unloaded=True, expected_reference_generation=0
    )
    _baseline(new, range(1, 31))
    ready = _baseline(new, [31])[-1]
    assert ready.zero.phase == "ready"
    before = new.zero

    duplicate = new.process(_measurement(20, _raw_for(20)))
    assert duplicate.duplicate is True
    assert duplicate.quality == "duplicate"
    assert new.zero == before, "duplicate must not advance or invalidate the zero"

    # Same timestamp, different source frame: not a duplicate, but it is reported
    # as a late/order diagnostic and cannot advance the filter or invalidate.
    same_time = new.process(
        _measurement(32, _raw_for(32), capture_time_ns=31 * FRAME_NS)
    )
    assert same_time.duplicate is False
    assert same_time.quality == "late"
    assert same_time.late is True
    assert new.zero == before, "order diagnostic must not invalidate the zero"

    duplicate_again = new.process(_measurement(20, _raw_for(20)))
    assert duplicate_again.duplicate is True


# ------------------------------------------------------------ window limit gates
def test_max_gap_s_fails_the_baseline_window_with_reason() -> None:
    new = _new_chain()
    new.request_zero(
        reference_id="ref",
        declared_unloaded=True,
        expected_reference_generation=0,
        max_gap_s=0.05,
    )
    _baseline(new, range(1, 6))
    jumped = new.process(_measurement(6, _raw_for(6), capture_time_ns=6 * 500_000_000))
    assert jumped.quality == "window_failed"
    assert jumped.zero.phase == "failed"
    assert "max_gap_s" in jumped.zero.failure_reason


def test_max_window_s_fails_the_baseline_window_with_reason() -> None:
    new = _new_chain()
    new.request_zero(
        reference_id="ref",
        declared_unloaded=True,
        expected_reference_generation=0,
        max_window_s=0.5,
    )
    _baseline(new, range(1, 10))
    late_limit = new.process(
        _measurement(10, _raw_for(10), capture_time_ns=1_000_000_000)
    )
    assert late_limit.quality == "window_failed"
    assert late_limit.zero.phase == "failed"
    assert "max_window_s" in late_limit.zero.failure_reason


# -------------------------------------------------------------- stream contract
def test_every_frame_produces_a_header_carrying_result() -> None:
    new = _new_chain()
    new.request_zero(
        reference_id="ref", declared_unloaded=True, expected_reference_generation=0
    )
    rows = _baseline(new, range(1, 31)) + [
        new.process(_measurement(31, (float("nan"), 1.0, 2.0))),
    ]
    assert len(rows) == 31
    for row in rows:
        assert isinstance(row, ForceStreamOutput)
        assert isinstance(row.frame, MeasurementFrame)
        assert row.frame.stamp.session_id == "s1"
        assert isinstance(row.reference_generation, int)
    invalid = rows[-1]
    assert invalid.accepted is True
    assert invalid.zero.phase == "ready", (
        "an invalid frame must not invent a zero failure"
    )
    assert np.isnan(invalid.raw[0]), (
        "the failed axis stays NaN; other raw axes remain visible"
    )
    assert all(np.isnan(value) for value in invalid.newton)
    assert invalid.quality in ("invalid", "uncalibrated")


def test_new_session_does_not_expose_the_old_external_reset_api() -> None:
    new = _new_chain()
    for name in (
        "process_measurement",
        "request_zero_external_reset",
        "tracker_reset_seq",
        "process_frame",
    ):
        assert not hasattr(new, name), name


def test_inputs_are_not_mutated_and_outputs_hold_no_internal_arrays() -> None:
    new = _new_chain()
    raw = _raw_for(1)
    flow = np.zeros((2, 2, 2), dtype=np.float32)
    measurement = _measurement(1, raw)
    sentinel = np.array(flow, copy=True)
    measurement.force_raw  # tuple attribute, no array to copy
    new.request_zero(
        reference_id="ref", declared_unloaded=True, expected_reference_generation=0
    )
    result = new.process(measurement)
    assert np.array_equal(flow, sentinel), "input arrays are read-only to the session"
    assert not isinstance(result.raw, np.ndarray)
    assert not isinstance(result.newton, np.ndarray)
    assert all(value == value for value in result.raw)  # finite tuple values


# --------------------------------------- normal reset order (Y1) and pending rows
def test_normal_reset_order_adopts_the_requested_version_and_completes_zero() -> None:
    """v0 frames in flight -> request_zero(v1) -> first v1 frame starts collecting.

    This is the order Studio uses. The request must not be invalidated by its own
    target version, and the still-in-flight v0 frames are pending, not late.
    """
    new = _new_chain(helpers.box_cal())
    seed = _baseline(new, range(1, 4), version=0)
    assert seed[-1].zero.phase != "ready"
    new.request_zero(
        reference_id="z", declared_unloaded=True, expected_reference_generation=1
    )

    in_flight = _baseline(new, range(4, 6), version=0)
    for row in in_flight:
        assert row.quality == "zero_pending", row.quality
        assert row.late is False, "in-flight current-version frames are not late"
        assert row.accepted is False
        assert row.zero.phase == "resetting", "pending frame must not advance the zero"

    version_one = _baseline(new, range(6, 37), version=1)
    assert version_one[0].zero.phase == "collecting"
    assert version_one[29].zero.phase == "collecting"
    assert version_one[30].zero.phase == "ready"
    assert np.isfinite(version_one[30].newton[0])
    assert version_one[0].reference_generation == 1


def test_request_arriving_after_the_first_new_version_frame_still_completes() -> None:
    new = _new_chain(helpers.box_cal())
    _baseline(new, range(1, 4), version=0)

    first_v1 = new.process(_measurement(4, _raw_for(4), version=1))
    assert first_v1.quality == "version_changed"
    assert first_v1.zero.phase == "failed"

    new.request_zero(
        reference_id="z-late", declared_unloaded=True, expected_reference_generation=1
    )
    frames = _baseline(new, range(5, 36), version=1)
    assert frames[0].zero.phase == "collecting"
    assert frames[30].zero.phase == "ready"
    assert np.isfinite(frames[30].newton[0])


def test_target_after_v1_waits_through_v1_inflight_then_adopts_v2() -> None:
    new = _new_chain(helpers.box_cal())
    new.request_zero(
        reference_id="v1", declared_unloaded=True, expected_reference_generation=1
    )
    _baseline(new, range(1, 32), version=1)
    assert new.zero.phase == "ready"
    new.request_zero(
        reference_id="v2", declared_unloaded=True, expected_reference_generation=2
    )
    pending = _baseline(new, range(32, 34), version=1)
    assert all(row.quality == "zero_pending" and not row.late for row in pending)
    assert all(row.zero.phase == "resetting" for row in pending)

    v2 = _baseline(new, range(34, 65), version=2)
    assert v2[0].zero.phase == "collecting"
    assert v2[30].zero.phase == "ready"


# ------------------------------------------------ branch frames register identity
def test_version_change_branch_frame_is_registered_for_duplicate_detection() -> None:
    new = _new_chain()
    _baseline(new, range(1, 3), version=0)
    changed = new.process(_measurement(3, _raw_for(3), version=1))
    assert changed.quality == "version_changed"
    after_change = new.zero
    repeat = new.process(_measurement(3, _raw_for(3), version=1))
    assert repeat.duplicate is True
    assert repeat.quality == "duplicate"
    assert new.zero == after_change, "duplicate must not be processed again"


# ------------------------------------------- T06b: field and FN-only calibration
def _field_payload_for(frame_id: int):
    return helpers._field_payload(0.5 + 0.05 * frame_id, 0.5)


def _field_models(tmp_path):
    document, context = helpers._field_document(tmp_path)
    old_model = helpers.parse_model(document, tmp_path / "field-calibration-old.json")
    new_model = helpers.parse_model(document, tmp_path / "field-calibration-new.json")
    return old_model, new_model, context


def _field_measurement(frame_id: int, *, include_field: bool = True, mask=None):
    return _measurement(
        frame_id,
        (0.0, 0.0, 0.0),
        field=_field_payload_for(frame_id) if include_field else None,
        valid_mask=np.ones((4, 4), dtype=bool) if mask is None else mask,
    )


def _field_old(tmp_path):
    old_model, new_model, context = _field_models(tmp_path)
    source = helpers.FakeMeasurementSource(
        lambda frame_id: (0.0, 0.0, 0.0),
        field_for=_field_payload_for,
        output_mask=np.ones((4, 4), dtype=bool),
    )
    session = helpers.ForceSession(
        source=source,
        sensor_family="finger",
        context=context,
        calibration=old_model,
        settings=helpers.legacy_settings(),
    )
    return session, source, new_model, context


def _compare_field_with_old(tmp_path) -> None:
    old, source, new_model, context = _field_old(tmp_path)
    new = ForceStreamSession(
        sensor_family="finger",
        context=context,
        calibration=new_model,
        settings=helpers.legacy_settings(),
    )
    assert new.measurement_requirements == (True, False)
    old.request_zero(reference_id="field-ref", declared_unloaded=True)
    new.request_zero(
        reference_id="field-ref",
        declared_unloaded=True,
        expected_reference_generation=0,
    )
    old_raw, old_newton, old_status = [], [], []
    new_raw, new_newton, new_status = [], [], []
    for frame_id in range(1, 41):
        old_out = helpers.process(old, source, frame_id)
        new_out = new.process(_field_measurement(frame_id))
        old_raw.append(np.asarray(old_out.raw))
        old_newton.append(np.asarray(old_out.newton))
        old_status.append(list(old_out.measurement.newton_status))
        new_raw.append(np.asarray(new_out.raw))
        new_newton.append(np.asarray(new_out.newton))
        new_status.append(list(new_out.newton_status))
    assert np.array_equal(np.stack(old_raw), np.stack(new_raw), equal_nan=True)
    assert np.array_equal(np.stack(old_newton), np.stack(new_newton), equal_nan=True)
    assert old_status == new_status
    assert new.zero.phase == "ready"
    # The field model outputs FN only; shear axes stay explicitly unavailable.
    assert np.isfinite(new_newton[-1][0])
    assert np.isnan(new_newton[-1][1]) and np.isnan(new_newton[-1][2])
    assert new_status[-1][1:] == ["uncalibrated", "uncalibrated"]


def _fn_only_models():
    context = helpers.build_context(
        {
            "camera": {"resolution_stitch": [4, 4]},
            "tracker": {"dis_mode": 1, "hhd_downsample": 1},
        }
    )
    document = helpers._global_gain_document(context, gains={"fn": 0.01})
    old_model = helpers.parse_model(document)
    new_model = helpers.parse_model(document)
    return old_model, new_model, context


def _fn_raw(frame_id: int) -> tuple[float, float, float]:
    return (100.0 + frame_id, 7.0, -7.0)


def _compare_fn_only_with_old() -> None:
    old_model, new_model, context = _fn_only_models()
    source = helpers.FakeMeasurementSource(_fn_raw)
    old = helpers.ForceSession(
        source=source,
        sensor_family="box",
        context=context,
        calibration=old_model,
        settings=helpers.legacy_settings(),
    )
    new = ForceStreamSession(
        sensor_family="box",
        context=context,
        calibration=new_model,
        settings=helpers.legacy_settings(),
    )
    assert new.measurement_requirements == (False, False)
    old.request_zero(reference_id="fn-ref", declared_unloaded=True)
    new.request_zero(
        reference_id="fn-ref",
        declared_unloaded=True,
        expected_reference_generation=0,
    )
    old_raw, old_newton, old_status = [], [], []
    new_raw, new_newton, new_status = [], [], []
    for frame_id in range(1, 41):
        old_out = helpers.process(old, source, frame_id)
        new_out = new.process(_measurement(frame_id, _fn_raw(frame_id)))
        old_raw.append(np.asarray(old_out.raw))
        old_newton.append(np.asarray(old_out.newton))
        old_status.append(list(old_out.measurement.newton_status))
        new_raw.append(np.asarray(new_out.raw))
        new_newton.append(np.asarray(new_out.newton))
        new_status.append(list(new_out.newton_status))
    assert np.array_equal(np.stack(old_raw), np.stack(new_raw), equal_nan=True)
    assert np.array_equal(np.stack(old_newton), np.stack(new_newton), equal_nan=True)
    assert old_status == new_status
    assert new.zero.phase == "ready"
    assert np.isfinite(new_newton[-1][0])
    assert np.isnan(new_newton[-1][1]) and np.isnan(new_newton[-1][2])
    assert new_status[-1][1:] == ["uncalibrated", "uncalibrated"]


def test_field_model_matches_old_session_on_this_platform(tmp_path) -> None:
    _compare_field_with_old(tmp_path)


def test_fn_only_model_matches_old_session_on_this_platform() -> None:
    _compare_fn_only_with_old()


def test_three_calibration_types_match_the_old_session_on_this_platform(
    tmp_path,
) -> None:
    # Same interpreter/platform for both chains: no cross-OS golden comparison.
    _compare_with_old(helpers.box_cal())  # schema2 scalar curves
    _compare_field_with_old(tmp_path)  # schema3 field model (FN-only)
    _compare_fn_only_with_old()  # schema3 global gain (FN-only)


def test_field_model_requires_context_and_field_frame(tmp_path) -> None:
    from orisys.calibration import CalibrationError

    old_model, new_model, context = _field_models(tmp_path)
    with pytest.raises(CalibrationError, match="pipeline context"):
        ForceStreamSession(sensor_family="finger", calibration=old_model)

    session = ForceStreamSession(
        sensor_family="finger",
        context=context,
        calibration=new_model,
        settings=helpers.legacy_settings(),
    )
    session.request_zero(
        reference_id="field-ref",
        declared_unloaded=True,
        expected_reference_generation=0,
    )
    missing = session.process(_field_measurement(1, include_field=False))
    assert missing.zero.phase == "failed"
    assert "invalid_input" in missing.zero.failure_reason
    assert all(np.isnan(value) for value in missing.newton)
    assert missing.newton_status[0] == "invalid_input"


def test_field_model_uncalibrated_session_stays_explicit(tmp_path) -> None:
    _, _, context = _field_models(tmp_path)
    new = ForceStreamSession(
        sensor_family="finger",
        context=context,
        calibration=None,
        settings=helpers.legacy_settings(),
    )
    new.request_zero(
        reference_id="field-raw",
        declared_unloaded=True,
        expected_reference_generation=0,
    )
    rows = [_field_measurement(frame_id) for frame_id in range(1, 32)]
    results = [new.process(row) for row in rows]
    assert results[-1].zero.phase == "ready"
    assert all(np.isnan(value) for value in results[-1].newton)
    assert set(results[-1].newton_status) == {"uncalibrated"}
    assert results[-1].quality == "uncalibrated"


def test_fn_only_model_requires_context_and_rejects_invalid_frames() -> None:
    from orisys.calibration import CalibrationError

    old_model, new_model, context = _fn_only_models()
    with pytest.raises(CalibrationError, match="pipeline context"):
        ForceStreamSession(sensor_family="box", calibration=old_model)

    new = ForceStreamSession(
        sensor_family="box",
        context=context,
        calibration=new_model,
        settings=helpers.legacy_settings(),
    )
    new.request_zero(
        reference_id="fn-ref", declared_unloaded=True, expected_reference_generation=0
    )
    invalid = new.process(_measurement(1, (float("nan"), 7.0, -7.0)))
    assert invalid.zero.phase == "failed"
    assert "invalid input frame" in invalid.zero.failure_reason
    assert all(np.isnan(value) for value in invalid.newton)


def test_fn_only_model_uncalibrated_session_stays_explicit() -> None:
    _, _, context = _fn_only_models()
    new = ForceStreamSession(
        sensor_family="box",
        context=context,
        calibration=None,
        settings=helpers.legacy_settings(),
    )
    new.request_zero(
        reference_id="fn-raw", declared_unloaded=True, expected_reference_generation=0
    )
    results = [
        new.process(_measurement(frame_id, _fn_raw(frame_id)))
        for frame_id in range(1, 32)
    ]
    assert results[-1].zero.phase == "ready"
    assert all(np.isnan(value) for value in results[-1].newton)
    assert set(results[-1].newton_status) == {"uncalibrated"}
    assert results[-1].quality == "uncalibrated"