"""Processing-interpretation tests across processor modes.

Pins the agreed semantics (not just current behaviour):
  * GPR object filtering is mode-agnostic — keep finite rows with score >= threshold
    (a normalized float for coherent GPR, a pair count for legacy GPR; same compare)
    inside the visible X/Z window, then apply draw limits. When more than
    max_detected_objects survive, ALL are hidden; legacy GPR passes draw_limits=None
    (count rules disabled) but is otherwise filtered identically — so heatmap markers
    match the objects-only view in both modes.
  * B-scan level/colormap/mean-subtraction/history transforms are deterministic.
  * ProcessingLiveConfig normalizes optional position lists to concrete int lists.
"""

from __future__ import annotations

import math
import unittest
from collections import deque

import numpy as np

from python_app.gui.controllers.app_window_plot.bscan_plot_mixin import (
    apply_mean_ascan_subtraction,
    bscan_levels,
    bscan_lookup_table,
    build_lut,
    rebuild_bscan_history_from_results,
)
from python_app.gui.controllers.app_window_plot.gpr_plot_mixin import AppWindowGprPlotMixin
from python_app.models.dataset_model import (
    ComboKey,
    ResultBlock,
    ResultCollection,
    ResultPayload,
)
from python_app.orchestration.gpr_locator import (
    apply_object_draw_limits,
    collection_has_gpr_payloads,
    collection_payload_by_name,
    collection_payloads_by_prefix,
    filter_object_rows,
    gpr_object_rows,
)
from python_app.orchestration.live_processing_config import ProcessingLiveConfig


def _table(name: str, rows) -> ResultPayload:
    return ResultPayload(processing_name=name, kind=4, table=np.asarray(rows, dtype=np.float32))


def _collection(payloads) -> ResultCollection:
    return ResultCollection(collection_id=1, monotonic_ns=1, collection_payloads=list(payloads))


# --------------------------------------------------------------------------- #
# Shared result-collection lookup helpers
# --------------------------------------------------------------------------- #
class CollectionLookupTest(unittest.TestCase):
    def test_payload_by_name_matches_name_and_optional_kind(self) -> None:
        col = _collection([_table("gpr_points", [[0, 0, 1]]), _table("gpr_region_centers", [[1, 1, 2, 9]])])
        self.assertIs(collection_payload_by_name(col, "gpr_points"), col.collection_payloads[0])
        self.assertIsNone(collection_payload_by_name(col, "missing"))
        self.assertIsNone(collection_payload_by_name(col, "gpr_points", kind=3))  # wrong kind

    def test_payloads_by_prefix_returns_all_in_order(self) -> None:
        col = _collection([
            ResultPayload(processing_name="gpr_region_mask_0", kind=3, image=np.zeros((1, 1), dtype=np.float32)),
            ResultPayload(processing_name="gpr_region_mask_1", kind=3, image=np.zeros((1, 1), dtype=np.float32)),
            _table("gpr_points", [[0, 0, 1]]),
        ])
        masks = collection_payloads_by_prefix(col, "gpr_region_mask_", kind=3)
        self.assertEqual([p.processing_name for p in masks], ["gpr_region_mask_0", "gpr_region_mask_1"])
        self.assertEqual(collection_payloads_by_prefix(col, "nope_"), [])

    def test_has_gpr_payloads(self) -> None:
        self.assertTrue(collection_has_gpr_payloads(_collection([_table("gpr_points", [[0, 0, 1]])])))
        self.assertFalse(collection_has_gpr_payloads(_collection([_table("pass_through", [[0, 0, 1]])])))
        self.assertFalse(collection_has_gpr_payloads(_collection([])))


# --------------------------------------------------------------------------- #
# GPR object row extraction
# --------------------------------------------------------------------------- #
class GprObjectRowsTest(unittest.TestCase):
    def test_extracts_first_three_columns_of_points(self) -> None:
        rows = gpr_object_rows(_collection([_table("gpr_points", [[1, 2, 3], [4, 5, 6]])]))
        self.assertTrue(np.array_equal(rows, np.array([[1, 2, 3], [4, 5, 6]], dtype=np.float32)))

    def test_falls_back_to_region_centers_and_drops_extra_columns(self) -> None:
        # region_centers is [x, z, score, pixel_count]; only the first 3 cols are used.
        rows = gpr_object_rows(_collection([_table("gpr_region_centers", [[1, 2, 3, 99]])]))
        self.assertTrue(np.array_equal(rows, np.array([[1, 2, 3]], dtype=np.float32)))

    def test_empty_when_no_object_payloads(self) -> None:
        self.assertEqual(gpr_object_rows(_collection([_table("gpr_accumulator", [[1, 2, 3]])])).shape, (0, 3))

    def test_rejects_table_with_too_few_columns(self) -> None:
        self.assertEqual(gpr_object_rows(_collection([_table("gpr_points", [[1, 2]])])).shape, (0, 3))


# --------------------------------------------------------------------------- #
# Object draw limits (hide-all-when-over, then top-M)
# --------------------------------------------------------------------------- #
class ApplyObjectDrawLimitsTest(unittest.TestCase):
    @staticmethod
    def _rows(n: int) -> np.ndarray:
        return np.column_stack([np.arange(n), np.arange(n), np.arange(n)]).astype(np.float32)

    def test_none_limits_passes_through(self) -> None:
        rows = self._rows(5)
        self.assertTrue(np.array_equal(apply_object_draw_limits(rows, None), rows))

    def test_hides_all_when_over_max(self) -> None:
        self.assertEqual(apply_object_draw_limits(self._rows(6), (5, 3)).shape, (0, 3))

    def test_keeps_top_m_when_within_max(self) -> None:
        out = apply_object_draw_limits(self._rows(4), (5, 2))
        self.assertEqual(out.shape, (2, 3))
        self.assertTrue(np.array_equal(out, self._rows(4)[:2]))

    def test_max_zero_hides_any_objects(self) -> None:
        # max_detected_objects == 0 means "hide all" (any object exceeds it).
        self.assertEqual(apply_object_draw_limits(self._rows(1), (0, 5)).shape, (0, 3))

    def test_empty_input_passes_through(self) -> None:
        empty = np.zeros((0, 3), dtype=np.float32)
        self.assertEqual(apply_object_draw_limits(empty, (5, 3)).shape, (0, 3))


# --------------------------------------------------------------------------- #
# Mode-agnostic object filtering (the core both modes + the locator share)
# --------------------------------------------------------------------------- #
class FilterObjectRowsTest(unittest.TestCase):
    BOUNDS = {"x_bounds": (-2.0, 2.0), "z_bounds": (0.0, 10.0)}

    def _filter(self, rows, *, min_score, draw_limits=None):
        return filter_object_rows(np.asarray(rows, dtype=np.float32), min_score=min_score,
                                  draw_limits=draw_limits, **self.BOUNDS)

    def test_keeps_in_window_and_at_or_above_threshold(self) -> None:
        out = self._filter([[0.0, 5.0, 0.5], [1.0, 1.0, 0.9]], min_score=0.5)  # score==threshold kept (inclusive)
        self.assertEqual(out.shape[0], 2)

    def test_drops_below_threshold(self) -> None:
        out = self._filter([[0.0, 5.0, 0.4]], min_score=0.5)
        self.assertEqual(out.shape[0], 0)

    def test_window_bounds_are_inclusive(self) -> None:
        out = self._filter([[2.0, 10.0, 1.0], [-2.0, 0.0, 1.0]], min_score=0.0)  # exactly on each edge
        self.assertEqual(out.shape[0], 2)

    def test_drops_outside_window(self) -> None:
        out = self._filter([[2.001, 5.0, 1.0], [0.0, 10.001, 1.0]], min_score=0.0)
        self.assertEqual(out.shape[0], 0)

    def test_drops_non_finite_rows(self) -> None:
        out = self._filter([[np.nan, 5.0, 1.0], [0.0, np.inf, 1.0], [0.0, 5.0, 1.0]], min_score=0.0)
        self.assertEqual(out.shape[0], 1)

    def test_legacy_pair_count_threshold_uses_same_compare(self) -> None:
        # legacy GPR passes an integer pair-count threshold; the >= compare is identical.
        out = self._filter([[0.0, 5.0, 3.0], [0.0, 5.0, 2.0]], min_score=3, draw_limits=None)
        self.assertTrue(np.array_equal(out, np.array([[0.0, 5.0, 3.0]], dtype=np.float32)))

    def test_draw_limits_hide_all_when_over(self) -> None:
        rows = [[0.0, 5.0, 1.0]] * 4
        self.assertEqual(self._filter(rows, min_score=0.0, draw_limits=(3, 2)).shape[0], 0)

    def test_legacy_none_limits_skips_count_rule(self) -> None:
        rows = [[0.0, 5.0, 1.0]] * 4
        self.assertEqual(self._filter(rows, min_score=0.0, draw_limits=None).shape[0], 4)


# --------------------------------------------------------------------------- #
# B-scan transforms
# --------------------------------------------------------------------------- #
class BscanTransformTest(unittest.TestCase):
    def test_mean_ascan_subtraction(self) -> None:
        history = deque([np.array([1.0, 1.0], dtype=np.float32), np.array([3.0, 3.0], dtype=np.float32)])
        self.assertTrue(np.array_equal(apply_mean_ascan_subtraction(history, enabled=False),
                                       np.array([[1, 1], [3, 3]], dtype=np.float32)))
        self.assertTrue(np.array_equal(apply_mean_ascan_subtraction(history, enabled=True),
                                       np.array([[-1, -1], [1, 1]], dtype=np.float32)))

    def test_levels_abs_mode(self) -> None:
        self.assertEqual(bscan_levels(np.array([[1.0, 4.0], [2.0, 3.0]], dtype=np.float32), "abs"), (1.0, 4.0))

    def test_levels_abs_degenerate(self) -> None:
        low, high = bscan_levels(np.full((2, 2), 5.0, dtype=np.float32), "abs")
        self.assertEqual(low, 5.0)
        self.assertGreater(high, low)

    def test_levels_signed_mode_symmetric(self) -> None:
        self.assertEqual(bscan_levels(np.array([[-3.0, 1.0]], dtype=np.float32), "real"), (-3.0, 3.0))

    def test_build_lut_shape_and_endpoints(self) -> None:
        lut = build_lut(["#000000", "#ffffff"])
        self.assertEqual(lut.shape, (256, 3))
        self.assertEqual(lut.dtype, np.uint8)
        self.assertTrue(np.array_equal(lut[0], [0, 0, 0]))
        self.assertTrue(np.array_equal(lut[-1], [255, 255, 255]))

    def test_lookup_table_for_each_axis_mode(self) -> None:
        for mode in ("abs", "real", "phase"):
            self.assertEqual(bscan_lookup_table(mode).shape, (256, 3))


class RebuildBscanHistoryTest(unittest.TestCase):
    @staticmethod
    def _bscan_collection(cid: int, combo, depth, amps, *, name="bscan", kind=1) -> ResultCollection:
        trace = np.asarray(amps, dtype=np.float32).astype(np.complex64)
        payload = ResultPayload(processing_name=name, kind=kind,
                                frequency_hz=np.asarray(depth, dtype=np.float32), trace=trace)
        block = ResultBlock(combo=ComboKey(input=combo[0], output=combo[1]), payloads=[payload])
        return ResultCollection(collection_id=cid, monotonic_ns=cid, blocks=[block])

    def test_accumulates_sweeps_per_combo(self) -> None:
        history = [
            self._bscan_collection(1, (0, 0), [1.0, 2.0], [10.0, 20.0]),
            self._bscan_collection(2, (0, 0), [1.0, 2.0], [11.0, 21.0]),
        ]
        by_combo, axes = rebuild_bscan_history_from_results(history, history_limit=10, floor_collection_id=0)
        self.assertEqual(len(by_combo[(0, 0)]), 2)
        self.assertTrue(np.array_equal(axes[(0, 0)], np.array([1.0, 2.0], dtype=np.float32)))

    def test_skips_non_bscan_and_mismatched_payloads(self) -> None:
        history = [
            self._bscan_collection(1, (0, 0), [1.0, 2.0], [1.0, 2.0], name="other"),  # wrong name
            self._bscan_collection(2, (0, 0), [1.0, 2.0], [1.0, 2.0], kind=2),         # wrong kind
            self._bscan_collection(3, (0, 0), [1.0, 2.0, 3.0], [1.0, 2.0]),            # size mismatch
        ]
        by_combo, _ = rebuild_bscan_history_from_results(history, history_limit=10, floor_collection_id=0)
        self.assertEqual(by_combo, {})

    def test_depth_axis_change_resets_history(self) -> None:
        history = [
            self._bscan_collection(1, (0, 0), [1.0, 2.0], [10.0, 20.0]),
            self._bscan_collection(2, (0, 0), [1.0, 2.0, 3.0], [11.0, 21.0, 31.0]),  # new depth axis
        ]
        by_combo, axes = rebuild_bscan_history_from_results(history, history_limit=10, floor_collection_id=0)
        self.assertEqual(len(by_combo[(0, 0)]), 1)  # reset on axis change; only the latest sweep remains
        self.assertEqual(axes[(0, 0)].shape, (3,))

    def test_floor_collection_id_excludes_older(self) -> None:
        history = [
            self._bscan_collection(1, (0, 0), [1.0], [10.0]),
            self._bscan_collection(2, (0, 0), [1.0], [20.0]),
        ]
        by_combo, _ = rebuild_bscan_history_from_results(history, history_limit=10, floor_collection_id=1)
        self.assertEqual(len(by_combo[(0, 0)]), 1)  # only collection_id > 1


# --------------------------------------------------------------------------- #
# GPR display-range helpers (pure static math on the mixin)
# --------------------------------------------------------------------------- #
class GprDisplayHelperTest(unittest.TestCase):
    def test_normalized_range_orders_and_expands(self) -> None:
        self.assertEqual(AppWindowGprPlotMixin._normalized_display_range(2.0, 5.0), (2.0, 5.0))
        self.assertEqual(AppWindowGprPlotMixin._normalized_display_range(5.0, 2.0), (2.0, 5.0))  # reordered
        low, high = AppWindowGprPlotMixin._normalized_display_range(3.0, 3.0)  # zero span expands to 0.1
        self.assertAlmostEqual(high - low, 0.1)
        self.assertAlmostEqual(0.5 * (low + high), 3.0)

    def test_display_y_min_keeps_surface_margin(self) -> None:
        self.assertEqual(AppWindowGprPlotMixin._gpr_display_y_min(2.0, 5.0), 2.0)  # surface not visible
        self.assertAlmostEqual(AppWindowGprPlotMixin._gpr_display_y_min(0.0, 10.0), -0.3)  # 3% of span
        self.assertAlmostEqual(AppWindowGprPlotMixin._gpr_display_y_min(-1.0, 1.0), -1.06)  # min 0.06 margin

    def test_object_label_candidates_are_distinct_positions(self) -> None:
        candidates = AppWindowGprPlotMixin._gpr_object_label_candidates(0.0, 0.0, x_span=1.0, z_span=1.0)
        self.assertEqual(len(candidates), 10)
        self.assertTrue(all(math.isfinite(x) and math.isfinite(z) for x, z, _ in candidates))


# --------------------------------------------------------------------------- #
# ProcessingLiveConfig normalization
# --------------------------------------------------------------------------- #
class ProcessingLiveConfigTest(unittest.TestCase):
    def test_none_positions_become_empty_lists(self) -> None:
        cfg = ProcessingLiveConfig(gpr_input_positions=None, gpr_output_positions=None)
        self.assertEqual(cfg.gpr_input_positions, [])
        self.assertEqual(cfg.gpr_output_positions, [])

    def test_positions_coerced_to_ints(self) -> None:
        cfg = ProcessingLiveConfig(gpr_input_positions=[1.5, 2.9], gpr_output_positions=[0.0])
        self.assertEqual(cfg.gpr_input_positions, [1, 2])
        self.assertEqual(cfg.gpr_output_positions, [0])

    def test_to_dict_is_json_typed(self) -> None:
        data = ProcessingLiveConfig().to_dict()
        self.assertIsInstance(data["processor_mode"], str)
        self.assertIsInstance(data["gpr_input_positions"], list)


if __name__ == "__main__":
    unittest.main()