orisys-sdk/tests/test_force_session.py

"""SDK force session: fake-input chain, zero lifecycle, models and output chain.

The measurement block is injected (deterministic fake features / fields), so
these tests need no camera and no hardware. Real calibration documents are used
where the model semantics matter: the shipped box file for schema 2, minimal
valid schema 3 documents built from ``build_context`` for the gain models.
"""

from __future__ import annotations

import hashlib

import numpy as np
import pytest

from orisys import calibration as cal_mod
from orisys.calibration import CalibrationError
from orisys.calibration_context import build_context
from orisys.calibration_models import parse_model
from orisys.field_calibration import (
    BASELINE_POLICY,
    FIELD_IDENTITY,
    MAPPING_POLICY,
)
from orisys.force_readout import OutputSettings, process_output_sequence
from orisys.force_session import (
    WORKFLOW_VERSION,
    ForceOutput,
    ForceSession,
    FrameStamp,
)
from orisys.sensor import FrameMeasurement
from orisys.signal_processing import LowpassFilter

FRAME_S = 1.0 / 30.0


# ---------------------------------------------------------------------------
# Deterministic injected measurement source
# ---------------------------------------------------------------------------
class FakeMeasurementSource:
    """Injected measurement/feature source: no camera, no image processing."""

    def __init__(
        self,
        raw_for=None,
        *,
        field_for=None,
        centroid_for=None,
        output_mask=None,
        fail_frames=(),
    ):
        self._raw_for = raw_for or (lambda frame_id: (10.0 + frame_id, 20.0, 30.0))
        self._field_for = field_for
        self._centroid_for = centroid_for
        self.output_mask = output_mask
        self.fail_frames = set(fail_frames)
        self.events: list[tuple[str, int | None]] = []
        self.calls: list[tuple[int, bool, bool]] = []
        self.session_enabled = False
        self.reset_count = 0

    def enable_measurement_session(self) -> None:
        self.session_enabled = True

    def disable_measurement_session(self) -> None:
        self.session_enabled = False

    def reset_tracker(self) -> None:
        self.reset_count += 1
        self.events.append(("reset", None))

    def read_measurement_frame(
        self,
        source_frame,
        *,
        frame_id,
        timestamp_s=None,
        want_field=False,
        want_contact=False,
    ) -> FrameMeasurement:
        del source_frame, timestamp_s
        self.calls.append((frame_id, want_field, want_contact))
        if frame_id in self.fail_frames:
            self.events.append(("measure_failed", frame_id))
            raise RuntimeError(f"compute failed on frame {frame_id}")
        raw = tuple(float(value) for value in self._raw_for(frame_id))
        field = self._field_for(frame_id) if want_field else None
        centroid = self._centroid_for(frame_id) if want_contact else None
        self.events.append(("measure", frame_id))
        return FrameMeasurement(
            raw=raw,
            raw_valid=tuple(bool(np.isfinite(value)) for value in raw),
            field=field,
            output_valid_mask=self.output_mask,
            contact_found=(centroid is not None) if want_contact else None,
            single_contact=True if (want_contact and centroid is not None) else None,
            centroid=centroid,
            diagnostics=(("frame_id", frame_id),),
        )


def stamp(frame_id: int, *, session_id: str = "s1", t: float | None = None) -> FrameStamp:
    return FrameStamp(
        session_id=session_id,
        frame_id=frame_id,
        source_frame_id=frame_id,
        signal_time_s=FRAME_S * frame_id if t is None else t,
        time_basis="recorded_signal",
    )


def image() -> np.ndarray:
    return np.zeros((2, 2, 3), dtype=np.uint8)


_BOX_CALIBRATION = None


def box_cal():
    """The shipped schema 2 box document, loaded once (real file, no fake curves)."""
    global _BOX_CALIBRATION
    if _BOX_CALIBRATION is None:
        loaded = cal_mod.load(sensor_type="box", strict=True)
        assert loaded is not None
        _BOX_CALIBRATION = loaded
    return _BOX_CALIBRATION


def legacy_settings(**overrides) -> OutputSettings:
    base = {"sample_rate_hz": 30.0}
    base.update(overrides)
    return OutputSettings(**base)


def make_session(source, calibration=None, *, context=None, settings=None, family="box"):
    return ForceSession(
        source=source,
        sensor_family=family,
        context=context,
        calibration=calibration,
        settings=settings if settings is not None else legacy_settings(),
    )


def process(session, source, frame_id, *, session_id="s1", t=None) -> ForceOutput:
    return session.process_frame(image(), stamp(frame_id, session_id=session_id, t=t))


def collect_baseline(session, source, *, start=1, count=30, reference_id="ref-1"):
    session.request_zero(reference_id=reference_id, declared_unloaded=True)
    return [
        process(session, source, frame_id)
        for frame_id in range(start, start + count)
    ]


# ---------------------------------------------------------------------------
# T1: executable fake chain, one model call, cached reads, missing is not zero
# ---------------------------------------------------------------------------
def test_fake_chain_runs_and_same_output_is_cached() -> None:
    source = FakeMeasurementSource(lambda frame_id: (10000.0 + 100.0 * frame_id, 0.0, 0.0))
    session = make_session(source, box_cal())
    assert source.session_enabled, "constructing the session must enable the SDK seam"
    assert session.workflow_version == WORKFLOW_VERSION

    outputs = collect_baseline(session, source)
    assert len(outputs) == 30
    assert session.zero.phase == "collecting"
    assert session.zero.completed is False
    ready = process(session, source, 31)
    assert session.zero.phase == "ready"
    assert session.zero.completed is True
    assert session.zero.first_frame_id == 1
    assert session.zero.last_frame_id == 30
    assert session.zero.effective_from_frame_id == 31
    assert ready.workflow_version == WORKFLOW_VERSION

    # Same frame again: the SAME immutable output, no new measurement.
    repeat = process(session, source, 31)
    assert repeat is ready
    assert len(source.calls) == 31
    assert ready.raw == pytest.approx((13100.0, 0.0, 0.0))
    assert ready.newton[0] == pytest.approx(ready.measurement.newton[0])


def test_one_model_call_per_frame_and_repeated_reads_do_not_advance() -> None:
    context = build_context(
        {
            "camera": {"resolution_stitch": [4, 4]},
            "tracker": {"dis_mode": 1, "hhd_downsample": 1},
        }
    )
    document = _global_gain_document(context, gains={"fn": 0.01})
    model = parse_model(document)
    calls: list[dict] = []
    original_predict = model.predict

    def counting_predict(**kwargs):
        calls.append(kwargs)
        return original_predict(**kwargs)

    model.predict = counting_predict
    source = FakeMeasurementSource(lambda frame_id: (100.0 + frame_id, 0.0, 0.0))
    session = ForceSession(
        source=source,
        sensor_family="box",
        context=context,
        calibration=model,
        settings=legacy_settings(),
    )
    session.request_zero(reference_id="ref", declared_unloaded=True)
    for frame_id in range(1, 31):
        process(session, source, frame_id)
    assert len(calls) == 30, "one inference per processed frame, baseline included"
    baseline_value = float(np.median([100.0 + frame_id for frame_id in range(1, 31)]))
    ready = process(session, source, 31)
    assert len(calls) == 31
    expected = (100.0 + 31 - baseline_value) * 0.01
    assert ready.newton[0] == pytest.approx(expected)
    assert ready.measurement.newton_status[0] == "ok"

    calls_before = len(calls)
    for _ in range(3):
        assert process(session, source, 31) is ready
    assert len(calls) == calls_before, "repeated reads must not infer again"

    reading = session.read_sdk_frame(image(), stamp(31))
    assert session.read_sdk_frame(image(), stamp(31)) is reading
    assert len(source.calls) == 31, "repeated reads must not measure again"


def test_missing_axes_stay_nan_and_a_valid_zero_is_zero() -> None:
    context = build_context(
        {
            "camera": {"resolution_stitch": [4, 4]},
            "tracker": {"dis_mode": 1, "hhd_downsample": 1},
        }
    )
    document = _global_gain_document(context, gains={"fn": 0.01})
    source = FakeMeasurementSource(lambda frame_id: (200.0, 7.0, -7.0))
    session = ForceSession(
        source=source,
        sensor_family="box",
        context=context,
        calibration=parse_model(document),
        settings=legacy_settings(),
    )
    # Before any unloaded reference: raw is valid, N is unavailable (not zero).
    first = process(session, source, 1)
    assert np.isfinite(first.raw).all()
    assert np.isnan(first.newton).all()
    assert first.measurement.newton_status[0] == "zero_unavailable"

    collect_from = 2
    session.request_zero(reference_id="ref", declared_unloaded=True)
    for frame_id in range(collect_from, collect_from + 30):
        process(session, source, frame_id)
    ready = process(session, source, collect_from + 30)
    assert ready.measurement.zero.completed is True
    # Unloaded reading: N is a measured zero, not a missing value.
    assert ready.newton[0] == pytest.approx(0.0)
    assert ready.measurement.newton_status[0] == "ok"
    # FN-only model: FX/FY can never be filled in.
    assert np.isnan(ready.newton[1]) and np.isnan(ready.newton[2])
    assert ready.measurement.newton_status[1:] == ("uncalibrated", "uncalibrated")


def test_failed_frame_never_reuses_the_previous_numbers() -> None:
    source = FakeMeasurementSource(lambda frame_id: (100.0, 0.0, 0.0), fail_frames={5})
    session = make_session(source, box_cal())
    session.request_zero(reference_id="ref", declared_unloaded=True)
    for frame_id in (1, 2, 3):
        process(session, source, frame_id)
    failed = process(session, source, 5)
    assert np.isnan(failed.raw).all()
    assert np.isnan(failed.newton).all()
    assert failed.measurement.sdk.raw_status == ("measurement_failed",) * 3
    assert session.zero.phase == "failed", "an invalid baseline frame fails the window"
    assert session.zero.completed is False


# ---------------------------------------------------------------------------
# T2: one unloaded-zero lifecycle for live and sequential replay
# ---------------------------------------------------------------------------
def test_twenty_nine_thirty_thirty_one_frames() -> None:
    source = FakeMeasurementSource(lambda frame_id: (100.0, 0.0, 0.0))
    session = make_session(source, box_cal())
    session.request_zero(reference_id="ref", declared_unloaded=True)
    for frame_id in range(1, 30):
        output = process(session, source, frame_id)
        assert np.isnan(output.newton).all()
    assert session.zero.collected_frames == 29
    assert session.zero.completed is False

    thirtieth = process(session, source, 30)
    assert session.zero.collected_frames == 30
    assert session.zero.phase == "collecting"
    assert session.zero.completed is False, "the 30th frame is still a baseline frame"
    assert np.isnan(thirtieth.newton).all()
    assert thirtieth.measurement.newton_status[0] == "zero_unavailable"

    thirty_first = process(session, source, 31)
    assert session.zero.completed is True
    assert session.zero.effective_from_frame_id == 31
    assert thirty_first.measurement.newton_status[0] == "ok"
    assert np.isfinite(thirty_first.newton[0])


def test_reset_is_applied_before_measuring_the_first_baseline_frame() -> None:
    source = FakeMeasurementSource(lambda frame_id: (100.0, 0.0, 0.0))
    session = make_session(source, box_cal())
    session.request_zero(reference_id="ref", declared_unloaded=True)
    first = process(session, source, 7)
    assert source.events[:2] == [("reset", None), ("measure", 7)]
    assert source.reset_count == 1
    assert session.zero.first_frame_id == 7
    assert np.isnan(first.newton).all()

    # A second request while collecting is a NEW atomic operation.
    session.request_zero(reference_id="ref-2", declared_unloaded=True)
    assert session.zero.phase == "resetting"
    assert session.zero.reference_id == "ref-2"
    process(session, source, 8)
    assert source.events[-2:] == [("reset", None), ("measure", 8)]
    assert source.reset_count == 2
    assert session.zero.first_frame_id == 8
    assert session.zero.collected_frames == 1


def test_invalid_frame_fails_the_window_and_a_new_request_restarts_it() -> None:
    source = FakeMeasurementSource(
        lambda frame_id: (100.0 if frame_id != 10 else float("nan"), 0.0, 0.0)
    )
    session = make_session(source, box_cal())
    session.request_zero(reference_id="ref", declared_unloaded=True)
    for frame_id in (1, 2, 3):
        process(session, source, frame_id)
    process(session, source, 10)
    assert session.zero.phase == "failed"
    assert session.zero.completed is False
    assert "invalid input frame" in session.zero.failure_reason
    for frame_id in (11, 12, 13):
        output = process(session, source, frame_id)
        assert np.isnan(output.newton).all()
    assert session.zero.completed is False

    # A fresh request re-runs the whole operation from a confirmed reset.
    source._raw_for = lambda frame_id: (100.0, 0.0, 0.0)
    session.request_zero(reference_id="ref-2", declared_unloaded=True)
    for frame_id in range(14, 44):
        process(session, source, frame_id)
    assert session.zero.collected_frames == 30
    ready = process(session, source, 44)
    assert session.zero.completed is True
    assert ready.measurement.newton_status[0] == "ok"


def test_cancel_and_late_results_cannot_restore_completion() -> None:
    source = FakeMeasurementSource(lambda frame_id: (100.0, 0.0, 0.0))
    session = make_session(source, box_cal())
    stale = session.read_sdk_frame(image(), stamp(1))
    session.request_zero(reference_id="ref", declared_unloaded=True)
    # A reading measured before the request can never fill the new window.
    progress = session.update_zero(stale)
    assert progress.collected_frames == 0
    assert progress.phase == "resetting"

    for frame_id in range(2, 32):
        process(session, source, frame_id)
    assert session.zero.collected_frames == 30
    assert session.zero.completed is False
    process(session, source, 32)
    assert session.zero.completed is True

    late = session.read_sdk_frame(image(), stamp(33))
    session.invalidate(reason="cancelled")
    assert session.zero.completed is False
    assert session.zero.phase == "failed"
    assert session.zero.failure_reason == "cancelled"
    progress = session.update_zero(late)
    assert progress.completed is False
    assert progress.phase == "failed"
    for frame_id in (34, 35, 36):
        output = process(session, source, frame_id)
        assert np.isnan(output.newton).all()
    assert session.zero.completed is False


def test_replay_window_is_enforced_and_live_and_replay_agree() -> None:
    live_source = FakeMeasurementSource(lambda frame_id: (100.0 + frame_id, 0.0, 0.0))
    live = make_session(live_source, box_cal())
    live.request_zero(reference_id="live", declared_unloaded=True)
    live_outputs = [process(live, live_source, frame_id) for frame_id in range(1, 32)]

    replay_source = FakeMeasurementSource(lambda frame_id: (100.0 + frame_id, 0.0, 0.0))
    replay = make_session(replay_source, box_cal())
    replay.request_zero(
        reference_id="replay", declared_unloaded=True, replay_window=(1, 30)
    )
    replay_outputs = [process(replay, replay_source, frame_id) for frame_id in range(1, 32)]

    assert [output.raw for output in live_outputs] == [
        output.raw for output in replay_outputs
    ]
    assert [output.newton for output in live_outputs] == [
        output.newton for output in replay_outputs
    ]
    assert live.zero.effective_from_frame_id == replay.zero.effective_from_frame_id

    short_source = FakeMeasurementSource(lambda frame_id: (100.0, 0.0, 0.0))
    short = make_session(short_source, box_cal())
    short.request_zero(
        reference_id="short", declared_unloaded=True, replay_window=(1, 20)
    )
    for frame_id in range(1, 21):
        process(short, short_source, frame_id)
    assert short.zero.phase == "collecting"
    process(short, short_source, 21)
    assert short.zero.phase == "failed"
    assert "replay window" in short.zero.failure_reason
    assert short.zero.completed is False


def test_zero_request_requires_the_unloaded_declaration_and_a_reference_id() -> None:
    source = FakeMeasurementSource()
    session = make_session(source, box_cal())
    with pytest.raises(ValueError, match="no load"):
        session.request_zero(reference_id="ref", declared_unloaded=False)
    with pytest.raises(ValueError, match="reference_id"):
        session.request_zero(reference_id="", declared_unloaded=True)
    with pytest.raises(ValueError, match="replay_window"):
        session.request_zero(
            reference_id="ref", declared_unloaded=True, replay_window=(5, 2)
        )
    assert session.zero.phase == "not_ready"


def test_measurement_failure_after_completion_keeps_the_baseline() -> None:
    source = FakeMeasurementSource(lambda frame_id: (100.0, 0.0, 0.0), fail_frames={33})
    session = make_session(source, box_cal())
    collect_baseline(session, source)
    process(session, source, 31)
    assert session.zero.completed is True
    failed = process(session, source, 33)
    assert np.isnan(failed.raw).all()
    assert np.isnan(failed.newton).all()
    assert session.zero.completed is True, "a failed frame is not a disconnect"
    resumed = process(session, source, 34)
    assert np.isfinite(resumed.raw).all()
    assert resumed.measurement.newton_status[0] == "ok"


def test_model_swap_clears_completion_and_a_rejected_swap_keeps_state() -> None:
    source = FakeMeasurementSource(lambda frame_id: (100.0, 0.0, 0.0))
    session = make_session(source, box_cal())
    collect_baseline(session, source)
    process(session, source, 31)
    assert session.zero.completed is True

    wrong = cal_mod.parse(
        {
            "schema_version": 2,
            "unit": "N",
            "channels": {
                "normal": {"method": "linear", "slope": 1.0, "offset": 0.0}
            },
            "meta": {"sensor_type": "mini", "sn": "OTHER"},
        }
    )
    with pytest.raises(CalibrationError):
        session.set_calibration(wrong)
    assert session.calibration is box_cal()
    assert session.zero.completed is True, "a rejected swap keeps model, baseline, state"

    session.set_calibration(box_cal())
    assert session.zero.completed is False
    assert session.zero.phase == "failed"
    assert session.zero.failure_reason == "model changed"
    output = process(session, source, 32)
    assert np.isnan(output.newton[0])


def test_new_source_session_invalidates_the_reference_automatically() -> None:
    source = FakeMeasurementSource(lambda frame_id: (100.0, 0.0, 0.0))
    session = make_session(source, box_cal())
    collect_baseline(session, source)
    process(session, source, 31)
    assert session.zero.completed is True
    output = process(session, source, 1, session_id="s2")
    assert session.zero.completed is False
    assert np.isfinite(output.raw).all()
    assert np.isnan(output.newton[0])


def test_out_of_order_frame_is_explicit_and_clears_completion() -> None:
    source = FakeMeasurementSource(lambda frame_id: (100.0, 0.0, 0.0))
    session = make_session(source, box_cal())
    collect_baseline(session, source)
    process(session, source, 31)
    assert session.zero.completed is True
    late = process(session, source, 20)
    assert np.isnan(late.raw).all()
    assert late.measurement.sdk.raw_status == ("frame_order",) * 3
    assert session.zero.completed is False
    # And the same time going backwards is a discontinuity too.
    session.request_zero(reference_id="r2", declared_unloaded=True)
    process(session, source, 40)
    backwards = process(session, source, 41, t=0.5)
    assert backwards.measurement.sdk.raw_status == ("signal_time_order",) * 3


# ---------------------------------------------------------------------------
# T3: legacy scalar models subtract the session baseline exactly once
# ---------------------------------------------------------------------------
def _nonlinear_calibration():
    """A cross-breakpoint, deliberately nonlinear scalar document (schema 2)."""
    return cal_mod.parse(
        {
            "schema_version": 2,
            "unit": "N",
            "channels": {
                "normal": {
                    "method": "piecewise_linear",
                    "points": [[0.0, 0.0], [100.0, 1.0], [200.0, 3.0]],
                    "extrapolate": True,
                    "clamp_output_min": None,
                },
                "shear_x": {
                    "method": "piecewise_linear",
                    "signed": True,
                    "input": "components",
                    "points": [[0.0, 0.0], [100.0, 2.0]],
                    "extrapolate": True,
                    "clamp_output_min": None,
                },
                "shear_y": {
                    "method": "piecewise_linear",
                    "signed": True,
                    "input": "components",
                    "points": [[0.0, 0.0], [100.0, 4.0]],
                    "extrapolate": True,
                    "clamp_output_min": None,
                },
            },
            "meta": {"sensor_type": "box", "sn": "T3-1"},
        }
    )


def test_legacy_baseline_is_subtracted_once_across_a_breakpoint() -> None:
    calibration = _nonlinear_calibration()
    baseline_raw = (150.0, 10.0, 10.0)
    source = FakeMeasurementSource(lambda frame_id: baseline_raw)
    session = make_session(source, calibration)
    session.request_zero(reference_id="ref", declared_unloaded=True)
    for frame_id in range(1, 31):
        process(session, source, frame_id)
    process(session, source, 31)
    assert session.zero.completed is True

    loaded_raw = (250.0, 10.0, 10.0)
    source._raw_for = lambda frame_id: loaded_raw
    loaded = process(session, source, 32)
    # curve(250) - curve(150) = 4.0 - 2.0 = 2.0.  A raw-domain delta would give
    # curve(100) = 1.0; a double subtraction would give curve(250)-2*curve(150)=0.
    assert loaded.newton[0] == pytest.approx(2.0)
    # The historical entry point keeps its own semantics.
    assert calibration.convert(*loaded_raw)[0] == pytest.approx(4.0)


def test_legacy_zero_is_in_the_force_domain_and_negative_noise_survives() -> None:
    calibration = _nonlinear_calibration()
    source = FakeMeasurementSource(lambda frame_id: (150.0, 0.0, 0.0))
    session = make_session(source, calibration)
    collect_baseline(session, source)
    process(session, source, 31)
    assert session.zero.completed is True
    source._raw_for = lambda frame_id: (180.0, -50.0, 0.0)
    output = process(session, source, 32)
    # fn: curve(180) - curve(150) = 2.6 - 2.0 (one subtraction only).
    assert output.newton[0] == pytest.approx(0.6)
    # signed shear: curve(-50) - curve(0) = -1.0 (sign preserved, no clamp).
    assert output.newton[1] == pytest.approx(-1.0)
    assert output.newton[2] == pytest.approx(0.0)


def test_legacy_conversion_runs_before_output_filtering() -> None:
    calibration = _nonlinear_calibration()
    source = FakeMeasurementSource(lambda frame_id: (150.0, 0.0, 0.0))
    session = make_session(
        source,
        calibration,
        settings=legacy_settings(lowpass_hz=5.0, lowpass_order=2),
    )
    collect_baseline(session, source)
    ready = process(session, source, 31)
    assert ready.measurement.zero.completed is True
    assert ready.measurement.newton[0] == pytest.approx(0.0)

    baseline_force = calibration.convert(150.0, 0.0, 0.0)[0]
    raws = [150.0, 20.0, 400.0, 400.0]
    converted = [
        calibration.convert(raw, 0.0, 0.0)[0] - baseline_force for raw in raws
    ]
    reference_filter = LowpassFilter(
        cutoff_hz=5.0, sample_rate_hz=30.0, order=2, n_channels=1
    )
    expected = np.array(
        [float(reference_filter.process([value])[0]) for value in converted]
    )
    raw_strategy = LowpassFilter(
        cutoff_hz=5.0, sample_rate_hz=30.0, order=2, n_channels=1
    )
    filtered_raw = [
        float(raw_strategy.process([raw])[0]) for raw in raws
    ]
    wrong = np.array(
        [
            calibration.convert(raw, 0.0, 0.0)[0] - baseline_force
            for raw in filtered_raw
        ]
    )
    actual = [ready.newton[0]]
    for index, raw in enumerate(raws[1:], start=32):
        source._raw_for = lambda frame_id, raw=raw: (raw, 0.0, 0.0)
        actual.append(process(session, source, index).newton[0])
    assert np.allclose(actual, expected)
    assert not np.allclose(actual, wrong), "raw must be converted before it is filtered"


def test_legacy_median_follows_the_converted_triplet() -> None:
    calibration = _nonlinear_calibration()
    source = FakeMeasurementSource(lambda frame_id: (150.0, 0.0, 0.0))
    session = make_session(
        source, calibration, settings=legacy_settings(median_window=3)
    )
    collect_baseline(session, source)
    ready = process(session, source, 31)
    baseline_force = calibration.convert(150.0, 0.0, 0.0)[0]
    loaded = [(200.0, 0.0, 0.0), (300.0, 0.0, 0.0), (250.0, 0.0, 0.0)]
    converted = [
        calibration.convert(*raw)[0] - baseline_force for raw in loaded
    ]
    latest = []
    for index, raw in enumerate(loaded, start=32):
        source._raw_for = lambda frame_id, raw=raw: raw
        latest.append(process(session, source, index).newton[0])
    # A causal median over the converted newtons, warmed up from the zero frame.
    window = [ready.measurement.newton[0], *converted]
    expected = [float(np.median(window[max(0, k - 2) : k + 1])) for k in range(len(window))]
    assert latest == pytest.approx(expected[1:])


# ---------------------------------------------------------------------------
# T5: shared output post-processing
# ---------------------------------------------------------------------------
def test_output_off_is_an_exact_passthrough_including_nan() -> None:
    source = FakeMeasurementSource(
        lambda frame_id: (100.0, 0.0, 0.0), fail_frames={33}
    )
    session = make_session(source, box_cal(), settings=legacy_settings())
    collect_baseline(session, source)
    ready = process(session, source, 31)
    assert ready.settings.median_window == 1
    assert ready.settings.lowpass_hz == 0.0
    assert ready.settings.raw_threshold == 0.0
    assert ready.settings.newton_threshold == 0.0
    assert ready.raw == ready.measurement.sdk.raw
    assert ready.newton == ready.measurement.newton
    failed = process(session, source, 33)
    assert np.isnan(failed.raw).all()
    assert np.isnan(failed.newton).all()


def test_nan_axis_does_not_enter_the_filter_history() -> None:
    chain = None
    from orisys.signal_processing import MultiAxisCausalFilter

    chain = MultiAxisCausalFilter(median_window=3, n_channels=3)
    first = chain.process((10.0, 20.0, 30.0))
    assert np.isfinite(first).all()
    skipped = chain.process((float("nan"), 20.0, 30.0), valid=(False, True, True))
    assert np.isnan(skipped[0])
    assert skipped[1] == pytest.approx(20.0)
    resumed = chain.process((10.0, 20.0, 30.0))
    # The NaN never entered the window: FN history is [10, 10], not [10, nan].
    assert resumed[0] == pytest.approx(10.0)


def test_thresholds_are_absolute_per_axis_and_do_not_touch_missing_axes() -> None:
    source = FakeMeasurementSource(lambda frame_id: (100.0, 50.0, -50.0))
    session = make_session(
        source,
        box_cal(),
        settings=legacy_settings(raw_threshold=51.0, newton_threshold=0.001),
    )
    collect_baseline(session, source)
    process(session, source, 31)
    assert session.zero.completed is True
    source._raw_for = lambda frame_id: (100.0, 50.0, -50.0)
    output = process(session, source, 32)
    # raw: |50| and |-50| are not above 51 -> suppressed to 0; FN stays.
    assert output.raw == pytest.approx((100.0, 0.0, 0.0))
    # newton: |N| <= threshold -> 0 on the calibrated axes, NaN stays NaN.
    assert output.newton[0] == pytest.approx(0.0)
    assert output.newton[1] == pytest.approx(0.0)
    assert np.isnan(output.newton[2]) or output.newton[2] == pytest.approx(0.0)


def test_effective_settings_report_a_cutoff_at_nyquist_as_off() -> None:
    source = FakeMeasurementSource(lambda frame_id: (100.0, 0.0, 0.0))
    session = make_session(
        source,
        box_cal(),
        settings=legacy_settings(sample_rate_hz=30.0, lowpass_hz=15.0),
    )
    output = process(session, source, 1)
    assert output.settings.lowpass_hz == 0.0
    assert output.settings.sample_rate_hz == pytest.approx(30.0)
    assert output.raw == output.measurement.sdk.raw


def test_filter_sample_rate_follows_signal_time_not_processing_speed() -> None:
    source = FakeMeasurementSource(lambda frame_id: (100.0, 0.0, 0.0))
    settings = legacy_settings(sample_rate_hz=30.0, lowpass_hz=5.0)
    session = make_session(source, box_cal(), settings=settings)
    for frame_id in range(1, 21):
        process(session, source, frame_id, t=frame_id / 60.0)
    assert session.output_settings.sample_rate_hz == pytest.approx(60.0, rel=0.05)


def test_offline_reprocessing_uses_the_same_chain_and_warm_up() -> None:
    source = FakeMeasurementSource(lambda frame_id: (150.0 + 10.0 * (frame_id % 7), 0.0, 0.0))
    settings = legacy_settings(median_window=3, lowpass_hz=4.0, lowpass_order=2)
    session = make_session(source, _nonlinear_calibration(), settings=settings)
    collect_baseline(session, source)
    ready_output = process(session, source, 31)
    assert ready_output.measurement.zero.completed is True

    frames = list(range(31, 46))
    outputs = [ready_output, *(process(session, source, frame_id) for frame_id in frames[1:])]
    unfiltered_raw = [output.measurement.sdk.raw for output in outputs]
    unfiltered_newton = [output.measurement.newton for output in outputs]
    signal_times = [frame * FRAME_S for frame in frames]
    reprocessed_raw, reprocessed_newton, effective = process_output_sequence(
        settings,
        unfiltered_raw,
        unfiltered_newton,
        signal_times_s=signal_times,
    )
    assert [output.raw for output in outputs] == reprocessed_raw
    assert [output.newton for output in outputs] == reprocessed_newton
    assert effective.sample_rate_hz == pytest.approx(30.0, rel=0.01)


def test_zero_completion_resets_the_output_filter_history() -> None:
    source = FakeMeasurementSource(lambda frame_id: (100.0, 0.0, 0.0))
    session = make_session(
        source, _nonlinear_calibration(), settings=legacy_settings(lowpass_hz=5.0)
    )
    collect_baseline(session, source)
    ready = process(session, source, 31)
    assert ready.measurement.zero.completed is True
    # A fresh IIR starts at its first sample: the first post-zero value is not
    # dragged down by the baseline frames' zero newtons.
    first_loaded = ready.newton[0]
    assert np.isfinite(first_loaded)
    clean = LowpassFilter(cutoff_hz=5.0, sample_rate_hz=30.0, order=2, n_channels=1)
    expected = float(
        clean.process([ready.measurement.newton[0]])[0]
    )
    assert first_loaded == pytest.approx(expected)


# ---------------------------------------------------------------------------
# Minimal legal schema 3 documents for the global-gain tests
# ---------------------------------------------------------------------------
def _global_gain_document(context, *, gains, outputs=None):
    axes = list(outputs or gains)
    return {
        "schema_version": 3,
        "unit": "N",
        "identity": {"sensor_type": "box", "sn": "TEST-1"},
        "pipeline": dict(context),
        "model": {
            "type": "global_gain/1",
            "input_profile": "force_delta/1",
            "output_frame": "sdk_force/1",
            "outputs": axes,
            "params": {
                "gains_N_per_raw": {axis: float(gains[axis]) for axis in axes},
                "valid_force_range_N": [-100.0, 100.0],
            },
        },
    }


# ---------------------------------------------------------------------------
# T4: schema 3 spatial gain and field models
# ---------------------------------------------------------------------------
_SQUARE_UV = [[0.0, 0.0], [3.0, 0.0], [0.0, 3.0], [3.0, 3.0]]
_SQUARE_TRIANGLES = [[0, 1, 2], [1, 3, 2]]


def _artifact(tmp_path, *, shape=(4, 4), uv=None, valid=None, name="map.npz"):
    height, width = shape
    if uv is None:
        rows, cols = np.mgrid[0:height, 0:width]
        uv = np.stack((cols, rows), axis=-1).astype(np.float64)
    if valid is None:
        valid = np.ones((height, width), dtype=bool)
    path = tmp_path / name
    np.savez(path, uv_by_output_px=uv, valid_mask=valid, origin_xy=np.array([0, 0]))
    return path, hashlib.sha256(path.read_bytes()).hexdigest()


def _uv_context(*, field_identity=None, coordinates):
    config = {
        "type": "compound_eye",
        "camera": {"resolution_raw": [4, 4], "resolution_stitch": [4, 4]},
        "tracker": {"dis_mode": 1, "hhd_downsample": 1},
        "filters": {"measurement": {"flow_ema_alpha": 1.0}},
        "contact": {},
    }
    return build_context(
        config,
        config["filters"],
        geometry_sha256="b" * 64,
        source_scale=0.5,
        coordinates=coordinates,
        processing_backend="uv_unroll",
        preprocessing={"flip_source_vertical": True, "mirror": False},
        field=field_identity,
    )


def _spatial_document(tmp_path, *, gains=None, coordinate_valid=None):
    path, digest = _artifact(tmp_path, valid=coordinate_valid)
    coordinates = {
        "frame": "autocal_flat",
        "unit": "px",
        "origin_xy": [0, 0],
        "map_sha256": digest,
    }
    context = _uv_context(coordinates=coordinates)
    document = {
        "schema_version": 3,
        "unit": "N",
        "identity": {"sensor_type": "finger", "sn": "T4-SPATIAL"},
        "pipeline": dict(context),
        "model": {
            "type": "spatial_gain/1",
            "input_profile": "fn_delta_autocal_uv/1",
            "output_frame": "sdk_force/1",
            "outputs": ["fn"],
            "params": {
                "points_uv": _SQUARE_UV,
                "gains_N_per_raw": gains or [0.001, 0.002, 0.003, 0.004],
                "triangles": _SQUARE_TRIANGLES,
                "valid_force_range_N": [-100.0, 100.0],
            },
            "artifacts": [
                {"role": "coordinate_map", "path": path.name, "sha256": digest}
            ],
        },
    }
    return document, context


def _field_document(tmp_path, *, gains=None, coordinate_valid=None):
    path, digest = _artifact(tmp_path, valid=coordinate_valid)
    coordinates = {
        "frame": "autocal_flat",
        "unit": "px",
        "origin_xy": [0, 0],
        "map_sha256": digest,
    }
    context = _uv_context(field_identity=FIELD_IDENTITY, coordinates=coordinates)
    document = {
        "schema_version": 3,
        "unit": "N",
        "identity": {"sensor_type": "finger", "sn": "T4-FIELD"},
        "pipeline": dict(context),
        "model": {
            "type": "field_gain/1",
            "input_profile": "hhd_normal_l1_field/1",
            "output_frame": "sdk_force/1",
            "outputs": ["fn"],
            "params": {
                "points_uv": _SQUARE_UV,
                "gains_N_per_raw": gains or [1.0, 1.0, 1.0, 1.0],
                "triangles": _SQUARE_TRIANGLES,
                "valid_force_range_N": [-100.0, 100.0],
                "mapping": dict(MAPPING_POLICY),
                "baseline": dict(BASELINE_POLICY),
            },
            "artifacts": [
                {"role": "coordinate_map", "path": path.name, "sha256": digest}
            ],
        },
    }
    return document, context


def _field_payload(dx=0.5, dy=0.5, *, shape=(2, 2), area_scale=4.0, mask=None):
    data = np.zeros((*shape, 2), dtype=np.float32)
    data[:, :, 0] = dx
    data[:, :, 1] = dy
    if mask is None:
        mask = np.ones(shape, dtype=bool)
    mask = np.asarray(mask, dtype=bool)
    data[~mask] = 0.0
    return {
        "d": data,
        "valid_mask": mask,
        "area_scale": float(area_scale),
        "grid_shape": tuple(shape),
        "reference_shape": (4, 4),
    }


def _field_session(tmp_path, *, gains=None, coordinate_valid=None, settings=None):
    document, context = _field_document(
        tmp_path, gains=gains, coordinate_valid=coordinate_valid
    )
    model = parse_model(document, tmp_path / "calibration.json")
    calls: list[dict] = []
    original_predict = model.predict

    def counting_predict(**kwargs):
        calls.append(kwargs)
        return original_predict(**kwargs)

    model.predict = counting_predict
    source = FakeMeasurementSource(
        lambda frame_id: (0.0, 0.0, 0.0),
        field_for=lambda frame_id: _field_payload(),
        output_mask=np.ones((4, 4), dtype=bool),
    )
    session = ForceSession(
        source=source,
        sensor_family="finger",
        context=context,
        calibration=model,
        settings=settings if settings is not None else legacy_settings(),
    )
    return session, source, model, calls


def _spatial_session(tmp_path):
    coordinate_valid = np.ones((4, 4), dtype=bool)
    coordinate_valid[0, 0] = False
    document, context = _spatial_document(
        tmp_path, coordinate_valid=coordinate_valid
    )
    model = parse_model(document, tmp_path / "calibration.json")
    calls: list[dict] = []
    original_predict = model.predict

    def counting_predict(**kwargs):
        calls.append(kwargs)
        return original_predict(**kwargs)

    model.predict = counting_predict
    loaded = {"raw": 100.0, "centroid": (1.0, 1.0)}
    source = FakeMeasurementSource(
        lambda frame_id: (loaded["raw"], 0.0, 0.0),
        centroid_for=lambda frame_id: loaded["centroid"],
    )
    session = ForceSession(
        source=source,
        sensor_family="finger",
        context=context,
        calibration=model,
        settings=legacy_settings(),
    )
    return session, source, model, calls, loaded


def test_spatial_model_uses_contact_mapping_and_one_inference_per_frame(tmp_path):
    session, source, model, calls, loaded = _spatial_session(tmp_path)
    session.request_zero(reference_id="ref", declared_unloaded=True)
    for frame_id in range(1, 31):
        process(session, source, frame_id)
    ready = process(session, source, 31)
    assert session.zero.completed is True
    assert len(calls) == 31, "exactly one inference per processed frame"
    assert ready.measurement.newton_status[0] == "ok"
    assert ready.newton[0] == pytest.approx(0.0)

    # Load: delta 100 raw at (1,1) -> barycentric gain (0.001+0.002+0.003)/3.
    loaded["raw"] = 200.0
    loaded_frame = process(session, source, 32)
    assert loaded_frame.newton[0] == pytest.approx(100.0 * 0.002)
    assert np.isnan(loaded_frame.newton[1])
    assert loaded_frame.measurement.newton_status[1:] == ("uncalibrated", "uncalibrated")
    assert len(calls) == 32

    # A contact on an unmapped coordinate cannot become a newton value.
    loaded["centroid"] = (0.0, 0.0)
    unmapped = process(session, source, 33)
    assert unmapped.measurement.newton_status[0] == "coordinate_unmapped"
    assert np.isnan(unmapped.newton[0])
    assert np.isfinite(unmapped.raw).all()


def test_field_model_baseline_is_weighted_and_delegated(tmp_path):
    session, source, model, calls = _field_session(tmp_path)
    session.request_zero(reference_id="field-ref", declared_unloaded=True)
    for frame_id in range(1, 31):
        output = process(session, source, frame_id)
        assert np.isnan(output.newton[0])
    assert model.field_baseline_N == pytest.approx(16.0)
    ready = process(session, source, 31)
    assert session.zero.completed is True
    assert len(calls) == 31
    diagnostics = dict(ready.measurement.diagnostics)
    assert diagnostics["baseline_N"] == pytest.approx(16.0)
    assert "baseline_raw_fn" not in diagnostics, "the weighted field baseline is not scalar"
    # The field model does not gate the unloaded window on contact location: the
    # measurement block is never even asked for a contact feature.
    assert all(not contact for _frame_id, _field, contact in source.calls)
    # Unloaded: the weighted sum equals the baseline through the same gain map.
    assert ready.newton[0] == pytest.approx(0.0)
    assert ready.measurement.newton_status[0] == "ok"

    # A loaded frame: q=3 per cell, P=48, FN = 48-16 = 32.
    source._field_for = lambda frame_id: _field_payload(1.5, 1.5)
    loaded = process(session, source, 32)
    assert loaded.newton[0] == pytest.approx(32.0)
    assert len(calls) == 32
    # FN-only field model never fills shear axes.
    assert np.isnan(loaded.newton[1]) and np.isnan(loaded.newton[2])
    assert loaded.measurement.newton_status[1:] == ("uncalibrated", "uncalibrated")


def test_field_model_mapping_error_is_explicit_and_does_not_fill_newtons(tmp_path):
    session, source, model, calls = _field_session(tmp_path)
    session.request_zero(reference_id="field-ref", declared_unloaded=True)
    for frame_id in range(1, 31):
        process(session, source, frame_id)
    ready = process(session, source, 31)
    assert ready.measurement.newton_status[0] == "ok"
    source.output_mask = np.ones((3, 3), dtype=bool)
    broken = process(session, source, 32)
    assert np.isfinite(broken.raw).all(), "raw stays available on a failed frame"
    assert np.isnan(broken.newton[0])
    assert broken.measurement.newton_status[0] == "invalid_input"
    assert len(calls) == 32


def test_scalar_and_field_baselines_are_separate(tmp_path):
    scalar_source = FakeMeasurementSource(lambda frame_id: (17.0, 3.0, 5.0))
    scalar = make_session(scalar_source, _nonlinear_calibration())
    collect_baseline(scalar, scalar_source)
    scalar_ready = process(scalar, scalar_source, 31)
    scalar_diagnostics = dict(scalar_ready.measurement.diagnostics)
    assert scalar_diagnostics["baseline_raw_fn"] == pytest.approx(17.0)
    assert "baseline_N" not in scalar_diagnostics

    field_session, field_source, _model, _calls = _field_session(tmp_path)
    field_session.request_zero(reference_id="field-ref", declared_unloaded=True)
    for frame_id in range(1, 31):
        process(field_session, field_source, frame_id)
    field_ready = process(field_session, field_source, 31)
    field_diagnostics = dict(field_ready.measurement.diagnostics)
    assert field_diagnostics["baseline_N"] == pytest.approx(16.0)
    assert "baseline_raw_fn" not in field_diagnostics
    # The two baselines are different quantities: a scalar raw median is not a
    # weighted field sum, and neither session may share the other's zero.
    assert scalar.zero.reference_id != field_session.zero.reference_id


# ---------------------------------------------------------------------------
# T5 (continued): stage composition, independent raw/N states, reset triggers
# ---------------------------------------------------------------------------
def test_median_then_lowpass_then_threshold_composition() -> None:
    settings = legacy_settings(
        median_window=3, lowpass_hz=10.0, lowpass_order=2, raw_threshold=0.5
    )
    source = FakeMeasurementSource(lambda frame_id: (100.0, 0.0, 0.0))
    session = make_session(source, _nonlinear_calibration(), settings=settings)
    collect_baseline(session, source)
    ready = process(session, source, 31)
    assert ready.measurement.zero.completed is True
    samples = [0.0, 4.0, 0.0, 4.0, 4.0]
    raws = [(150.0 + sample, 0.0, 0.0) for sample in samples]
    outputs = []
    for index, raw in enumerate(raws, start=32):
        source._raw_for = lambda frame_id, raw=raw: raw
        outputs.append(process(session, source, index).raw[0])
    from orisys.signal_processing import CausalMedianFilter

    median = CausalMedianFilter(window=3, n_channels=1)
    lowpass = LowpassFilter(cutoff_hz=10.0, sample_rate_hz=30.0, order=2, n_channels=1)
    expected = []
    # Warm the reference chain from the ready frame's sample, as the session did.
    warmup = ready.measurement.sdk.raw[0]
    med = float(median.process([warmup])[0])
    assert float(lowpass.process([med])[0]) == pytest.approx(warmup)
    for raw in raws:
        med = float(median.process([raw[0]])[0])
        smooth = float(lowpass.process([med])[0])
        expected.append(0.0 if abs(smooth) <= 0.5 else smooth)
    assert outputs == pytest.approx(expected)


def test_raw_and_newton_states_are_independent() -> None:
    source = FakeMeasurementSource(lambda frame_id: (100.0, 0.0, 0.0))
    session = make_session(
        source,
        _nonlinear_calibration(),
        settings=legacy_settings(lowpass_hz=6.0),
    )
    collect_baseline(session, source)
    ready = process(session, source, 31)
    assert ready.measurement.zero.completed is True
    samples = [(160.0 + 5.0 * index, 0.0, 0.0) for index in range(6)]
    outputs = [ready]
    for index, raw in enumerate(samples, start=32):
        source._raw_for = lambda frame_id, raw=raw: raw
        outputs.append(process(session, source, index))
    raw_reference = LowpassFilter(cutoff_hz=6.0, sample_rate_hz=30.0, order=2, n_channels=1)
    newton_reference = LowpassFilter(
        cutoff_hz=6.0, sample_rate_hz=30.0, order=2, n_channels=1
    )
    expected_raw = [
        float(raw_reference.process([output.measurement.sdk.raw[0]])[0])
        for output in outputs
    ]
    expected_newton = [
        float(newton_reference.process([output.measurement.newton[0]])[0])
        for output in outputs
    ]
    assert [output.raw[0] for output in outputs] == pytest.approx(expected_raw)
    assert [output.newton[0] for output in outputs] == pytest.approx(expected_newton)
    # Both domains advance once per frame from their own valid samples only.
    assert all(np.isfinite(output.raw[0]) for output in outputs)
    assert all(np.isfinite(output.newton[0]) for output in outputs)


def test_invalid_newton_axis_stays_missing_under_active_filtering(tmp_path):
    session, source, _model, _calls = _field_session(
        tmp_path, settings=legacy_settings(lowpass_hz=5.0)
    )
    session.request_zero(reference_id="field-ref", declared_unloaded=True)
    for frame_id in range(1, 31):
        process(session, source, frame_id)
    ready = process(session, source, 31)
    assert ready.measurement.zero.completed is True
    loaded = process(session, source, 32)
    assert np.isfinite(loaded.raw).all(), "raw is still filtered normally"
    assert np.isnan(loaded.newton[1]), "an uncalibrated axis never becomes zero"
    assert np.isnan(loaded.newton[2])


def test_model_change_and_new_zero_reset_the_output_filter(tmp_path):
    def settings():
        return legacy_settings(lowpass_hz=5.0, median_window=3)

    first_source = FakeMeasurementSource(lambda frame_id: (150.0, 0.0, 0.0))
    first = make_session(first_source, _nonlinear_calibration(), settings=settings())
    collect_baseline(first, first_source)
    process(first, first_source, 31)
    assert first.zero.completed is True
    for index, raw in enumerate([(400.0, 0.0, 0.0)] * 5, start=32):
        first_source._raw_for = lambda frame_id, raw=raw: raw
        process(first, first_source, index)

    # A model swap clears completion and the filter history; a new zero then
    # starts both from scratch, exactly like a brand-new session.
    first.set_calibration(_nonlinear_calibration())
    assert first.zero.completed is False
    first.request_zero(reference_id="after-swap", declared_unloaded=True)
    first_source._raw_for = lambda frame_id: (150.0, 0.0, 0.0)
    for frame_id in range(40, 70):
        process(first, first_source, frame_id)
    swapped_ready = process(first, first_source, 70)
    assert first.zero.completed is True

    fresh_source = FakeMeasurementSource(lambda frame_id: (150.0, 0.0, 0.0))
    fresh = make_session(fresh_source, _nonlinear_calibration(), settings=settings())
    collect_baseline(fresh, fresh_source)
    fresh_ready = process(fresh, fresh_source, 31)
    assert swapped_ready.raw == fresh_ready.raw
    assert swapped_ready.newton == fresh_ready.newton

    # Same comparison one loaded frame later, where history would show.
    loaded = (400.0, 0.0, 0.0)
    first_source._raw_for = lambda frame_id: loaded
    swapped_loaded = process(first, first_source, 71)
    fresh_source._raw_for = lambda frame_id: loaded
    fresh_loaded = process(fresh, fresh_source, 32)
    assert swapped_loaded.raw == fresh_loaded.raw
    assert swapped_loaded.newton == fresh_loaded.newton


def test_old_readout_entry_keeps_filter_then_convert_order() -> None:
    """The historical entry point is not silently reordered by the new path."""
    from orisys.force_readout import ForceReadout

    calibration = _nonlinear_calibration()
    readout = ForceReadout({"force_median_window": 3}, config_fps=30.0)
    raws = [(200.0, 0.0, 0.0), (300.0, 0.0, 0.0), (250.0, 0.0, 0.0)]
    outputs = [readout.apply(raw, calibration.convert)[0] for raw in raws]
    # median of the raw FN values is 250 -> curve(250) = 4.0; the new session
    # ordering (convert first) would report curve(250)-curve(0) = 4.0 too, but
    # the point here is that the OLD call still filters before converting, which
    # is exactly the recorded historical behaviour.
    assert outputs[-1] == pytest.approx(4.0)
    assert readout.raw_unfiltered == (250.0, 0.0, 0.0)