"""Display unit selection and the per-axis drawing rule.

Two defects are covered here:

* the live charts chose their unit from the device-level ``force_calibrated``
  flag, so a device with only its normal-force axis calibrated either showed
  every axis in raw counts or showed raw counts inside newton panels;
* the drawing rule asked for **all three** axes to be finite, so a valid FN
  disappeared from the chart whenever FX/FY had no calibrated value.
"""

from __future__ import annotations

import math

import numpy as np
import pytest

from orisys_app_shared.acquisition.sdk_worker import SdkPacket
from orisys_app_shared.force import select_display_forces
from orisys_app_shared.force.session import SessionSample, SessionStore
from orisys_app_shared.record.sample import to_record_sample

NAN = float("nan")


def _packet(**overrides) -> SdkPacket:
    frame = np.zeros((4, 4, 3), np.uint8)
    kwargs = dict(
        raw_frame=frame,
        processed_frame=frame,
        timestamp=1000.0,
        fps=30.0,
        process_ms=1.0,
        normal_force=1.5,
        shear_x=0.25,
        shear_y=-0.25,
        normal_force_raw=123.0,
        shear_x_raw=2.0,
        shear_y_raw=-2.0,
        force_calibrated=True,
        calibrated_axes=("fn", "fx", "fy"),
        frame_id=0,
    )
    kwargs.update(overrides)
    return SdkPacket(**kwargs)


def test_newton_mode_marks_uncalibrated_axes_unavailable():
    sample = to_record_sample(_packet(calibrated_axes=("fn",)))
    display = select_display_forces(units_pref="newton", sample=sample)
    assert display.unit == "N"
    assert display.values[0] == pytest.approx(1.5)
    assert math.isnan(display.values[1]) and math.isnan(display.values[2])
    assert display.available == (True, False, False)
    assert display.unavailable_axes() == ("fx", "fy")
    assert not display.using_raw_fallback


def test_raw_mode_shows_the_raw_counts_for_every_axis():
    sample = to_record_sample(_packet(calibrated_axes=("fn",)))
    display = select_display_forces(units_pref="raw", sample=sample)
    assert display.unit == "raw"
    assert display.values == (123.0, 2.0, -2.0)
    assert display.available == (True, True, True)
    assert not display.using_raw_fallback


def test_an_uncalibrated_device_falls_back_to_raw_and_says_so():
    sample = to_record_sample(
        _packet(force_calibrated=False, calibrated_axes=())
    )
    display = select_display_forces(units_pref="newton", sample=sample)
    assert display.unit == "raw"
    assert display.using_raw_fallback, "the fallback must be announceable"
    assert display.requested_unit == "N"
    assert display.values == (123.0, 2.0, -2.0)


def test_a_failed_readout_has_no_available_axis():
    sample = to_record_sample(
        _packet(valid=False, invalid_reason="compute failed", stale_forces=True)
    )
    for unit in ("N", "raw"):
        display = select_display_forces(units_pref=unit, sample=sample)
        assert display.available == (False, False, False)
        assert all(math.isnan(value) for value in display.values)


def _store_with_one_calibrated_axis() -> SessionStore:
    store = SessionStore()
    for index in range(5):
        store.append(
            SessionSample(
                index,
                index / 30.0,
                fn_raw=100.0 + index,
                fx_raw=2.0,
                fy_raw=-2.0,
                fn_n=1.0 + index * 0.1,
                fx_n=NAN,
                fy_n=NAN,
                valid=True,
                calibrated_axes=("fn",),
            )
        )
    return store


def test_a_valid_axis_is_drawn_even_when_the_others_have_no_value():
    """The all-axes rule hid a good FN because FX/FY were uncalibrated."""
    slice_ = _store_with_one_calibrated_axis().data()
    times, values, mask = slice_.plottable("N")
    assert mask.shape == (5, 3)
    assert mask[:, 0].all(), "the calibrated axis was hidden"
    assert not mask[:, 1].any()
    assert not mask[:, 2].any()

    fn_times, fn_values = slice_.plottable_axis("N", 0)
    assert fn_times.size == 5 == fn_values.size
    assert float(fn_values.max()) == pytest.approx(1.4)
    fx_times, _fx_values = slice_.plottable_axis("N", 1)
    assert fx_times.size == 0


def test_the_raw_triplet_is_still_complete_for_the_same_samples():
    slice_ = _store_with_one_calibrated_axis().data()
    for index in range(3):
        times, values = slice_.plottable_axis("raw", index)
        assert times.size == 5
        assert np.isfinite(values).all()


def test_an_invalid_sample_is_excluded_from_every_axis():
    store = _store_with_one_calibrated_axis()
    store.append(SessionSample(9, 9.0, fn_n=NAN, valid=False))
    slice_ = store.data()
    _times, _values, mask = slice_.plottable("N")
    assert not mask[-1].any(), "an invalid sample reached a drawing path"