"""Tests for spatio-temporal dithering of the final 8-bit quantization.

The load-bearing property is *temporal convergence*: over many frames, the
dithered output time-averages back to the true sub-integer value (that's what
lets the eye see a higher effective bit depth instead of banding).
"""

from __future__ import annotations

import numpy as np

from ledgrab.core.capture.calibration import calibration_from_dict, calibration_to_dict
from ledgrab.core.capture.edge_interpolation import average_edge_to_leds
from ledgrab.utils.dither import ordered_dither_quantize


class TestOrderedDither:
    def test_single_frame_is_floor_or_ceil(self):
        vals = np.array([[100.4, 50.0, 200.7]], dtype=np.float32)
        out = ordered_dither_quantize(vals, frame_index=7)
        assert out.dtype == np.uint8
        assert out[0, 0] in (100, 101)
        assert out[0, 1] == 50  # exact integer never moves
        assert out[0, 2] in (200, 201)

    def test_integers_are_stable_across_frames(self):
        vals = np.array([[10.0, 20.0, 30.0]], dtype=np.float32)
        for f in range(20):
            out = ordered_dither_quantize(vals, frame_index=f)
            assert list(out[0]) == [10, 20, 30]

    def test_temporal_average_converges_to_true_value(self):
        vals = np.array([[100.4, 50.9, 200.25]], dtype=np.float32)
        acc = np.zeros(3, dtype=np.float64)
        n = 2000
        for f in range(n):
            acc += ordered_dither_quantize(vals, frame_index=f)[0]
        avg = acc / n
        assert abs(avg[0] - 100.4) < 0.3
        assert abs(avg[1] - 50.9) < 0.3
        assert abs(avg[2] - 200.25) < 0.3

    def test_same_threshold_across_channels_preserves_hue_steps(self):
        # All three channels share the per-LED threshold, so an equal-channel
        # grey never splits into a coloured pixel.
        vals = np.array([[123.5, 123.5, 123.5]], dtype=np.float32)
        for f in range(50):
            out = ordered_dither_quantize(vals, frame_index=f)
            assert out[0, 0] == out[0, 1] == out[0, 2]

    def test_clips_to_byte_range(self):
        vals = np.array([[-5.0, 255.9, 300.0]], dtype=np.float32)
        out = ordered_dither_quantize(vals, frame_index=3)
        assert out[0, 0] == 0
        assert out[0, 1] == 255
        assert out[0, 2] == 255


class TestEdgeReductionDither:
    def _half_edge(self):
        # 2 columns (black, white) → 1 LED whose mean is exactly 127.5, i.e. a
        # value the plain uint8 path can only represent as 127.
        e = np.zeros((2, 2, 3), dtype=np.uint8)
        e[:, 1, :] = 255
        return e

    def test_dithered_output_varies_by_frame(self):
        edge = self._half_edge()
        seen = {
            int(average_edge_to_leds(edge, "top", 1, {}, "k", dither=True, frame_index=f)[0, 0])
            for f in range(50)
        }
        # The 127.5 LED straddles two codes → dither flips it across frames.
        assert seen == {127, 128}

    def test_temporal_average_recovers_the_half_step(self):
        edge = self._half_edge()
        # Plain quantization truncates 127.5 → 127; dither recovers ~127.5.
        plain = int(average_edge_to_leds(edge, "top", 1, {}, "k")[0, 0])
        acc = np.zeros(3, dtype=np.float64)
        n = 1500
        for f in range(n):
            acc += average_edge_to_leds(edge, "top", 1, {}, "k", dither=True, frame_index=f)[0]
        avg = acc[0] / n
        assert plain == 127
        assert abs(avg - 127.5) < 0.2


class TestCalibrationRoundTrip:
    def test_dither_round_trips(self):
        cfg = calibration_from_dict(
            {
                "mode": "simple",
                "layout": "clockwise",
                "start_position": "bottom_left",
                "leds_top": 5,
                "dither": True,
            }
        )
        assert cfg.dither is True
        assert calibration_to_dict(cfg).get("dither") is True

    def test_dither_default_off_and_omitted(self):
        cfg = calibration_from_dict(
            {
                "mode": "simple",
                "layout": "clockwise",
                "start_position": "bottom_left",
                "leds_top": 5,
            }
        )
        assert cfg.dither is False
        assert "dither" not in calibration_to_dict(cfg)