"""Unit tests for switch-widened matrix capture.

Cover the targeted single-step acquisition on ``SwitchedMatrixRadarService`` and
verify the manual per-combo capture workflow uses it instead of sweeping the full
widened matrix (the regression that froze the GUI for the whole matrix per click).
"""

from __future__ import annotations

import time
import unittest
from unittest import mock

import numpy as np

from python_app.hardware_full.multi_device_service import MultiDeviceLibreVnaService
from python_app.hardware_full.switched_matrix_radar_service import SwitchedMatrixRadarService
from python_app.models.dataset_model import ComboKey, SweepCollection, TraceData
from python_app.models.run_config_model import RunConfigModel
from python_app.workflows.sequential_capture_workflow import SequentialCaptureSession

_INNER_INPUTS = 4
_INNER_OUTPUTS = 2
_POINTS = 8


class _FakeInnerMatrixRadar:
    """Matrix radar stub emitting the canonical 2x4 combo set per acquisition."""

    def __init__(self) -> None:
        self.acquire_count = 0

    def open(self) -> None:
        pass

    def close(self) -> None:
        pass

    def configure(self, sweep) -> None:
        pass

    def recover(self) -> None:
        pass

    def acquire_collection(self, collection_id: int = 1) -> SweepCollection:
        self.acquire_count += 1
        frequency_hz = np.linspace(1e6, 2e6, _POINTS, dtype=np.float32)
        traces = [
            TraceData(
                combo=ComboKey(input=input_pos, output=output_pos),
                frequency_hz=frequency_hz,
                s11=np.full(_POINTS, complex(self.acquire_count, 0), dtype=np.complex64),
                s21=np.full(_POINTS, complex(input_pos, output_pos), dtype=np.complex64),
            )
            for output_pos in range(_INNER_OUTPUTS)
            for input_pos in range(_INNER_INPUTS)
        ]
        return SweepCollection(
            collection_id=int(collection_id),
            monotonic_ns=time.monotonic_ns(),
            traces=traces,
        )


class _FakeSwitch:
    """Switch stub recording every position it is driven to."""

    def __init__(self, positions: int) -> None:
        self.positions = positions
        self.switched_to: list[int] = []

    def open(self) -> None:
        pass

    def close(self) -> None:
        pass

    def position_count(self) -> int:
        return self.positions

    def switch_to(self, position: int) -> None:
        self.switched_to.append(int(position))


def _switched_service(input_steps: int = 3) -> tuple[SwitchedMatrixRadarService, _FakeInnerMatrixRadar, _FakeSwitch]:
    inner = _FakeInnerMatrixRadar()
    input_switch = _FakeSwitch(input_steps)
    service = SwitchedMatrixRadarService(
        inner=inner,
        output_switch=None,
        input_switch=input_switch,
        inner_output_positions=_INNER_OUTPUTS,
        inner_input_positions=_INNER_INPUTS,
        settling_ms=0,
    )
    return service, inner, input_switch


class SwitchedMatrixComboAcquisitionTest(unittest.TestCase):
    """acquire_combo_collection must acquire exactly one physical switch step."""

    def test_acquires_only_the_step_containing_the_combo(self) -> None:
        service, inner, input_switch = _switched_service(input_steps=3)

        # Widened input 9 lives in physical step 9 // 4 = 2.
        collection = service.acquire_combo_collection(input_pos=9, output_pos=1)

        self.assertEqual(inner.acquire_count, 1)
        self.assertEqual(input_switch.switched_to, [2])
        self.assertEqual(len(collection.traces), _INNER_INPUTS * _INNER_OUTPUTS)
        combos = {(trace.combo.input, trace.combo.output) for trace in collection.traces}
        self.assertIn((9, 1), combos)
        # Every trace of the step is remapped into the widened axis of that step.
        self.assertEqual(
            combos,
            {(2 * _INNER_INPUTS + i, o) for i in range(_INNER_INPUTS) for o in range(_INNER_OUTPUTS)},
        )

    def test_rejects_out_of_range_combo(self) -> None:
        service, _inner, _input_switch = _switched_service(input_steps=3)
        with self.assertRaises(ValueError):
            service.acquire_combo_collection(input_pos=12, output_pos=0)
        with self.assertRaises(ValueError):
            service.acquire_combo_collection(input_pos=0, output_pos=2)

    def test_full_collection_still_covers_widened_matrix_in_canonical_order(self) -> None:
        service, inner, input_switch = _switched_service(input_steps=3)

        collection = service.acquire_collection(collection_id=7)

        self.assertEqual(inner.acquire_count, 3)
        self.assertEqual(input_switch.switched_to, [0, 1, 2])
        expected_combos = [
            (input_pos, output_pos)
            for output_pos in range(_INNER_OUTPUTS)
            for input_pos in range(3 * _INNER_INPUTS)
        ]
        self.assertEqual(
            [(trace.combo.input, trace.combo.output) for trace in collection.traces],
            expected_combos,
        )


class ManualComboCaptureUsesTargetedAcquisitionTest(unittest.TestCase):
    """The per-combo capture session must not sweep the full widened matrix."""

    @staticmethod
    def _switched_mock_config() -> RunConfigModel:
        config = RunConfigModel()
        config.radar.model = RunConfigModel.LIBREVNA_MULTI_MODEL
        config.radar.driver_mode = "mock"
        config.radar.multi_device.slave_serials = ["SLAVE_A", "SLAVE_B"]
        config.radar.multi_device.input_switch_positions = 3
        config.apply_device_model_constraints()
        return config

    def test_manual_capture_runs_one_inner_collection_per_median_sweep(self) -> None:
        config = self._switched_mock_config()
        session = SequentialCaptureSession(
            config=config,
            kind="s21_calibration",
            set_name="targeted_test",
            median_sweep_count=2,
        )
        with mock.patch.object(
            MultiDeviceLibreVnaService,
            "acquire_collection",
            autospec=True,
            side_effect=MultiDeviceLibreVnaService.acquire_collection,
        ) as inner_acquire:
            session.open()
            try:
                trace = session.capture_current_combo()
            finally:
                session.close()

        first_combo = config.combos[0]
        self.assertEqual(
            (trace.combo.input, trace.combo.output),
            (first_combo.input, first_combo.output),
        )
        # 2 median sweeps of ONE physical step — not 2 x 3 full-matrix steps.
        self.assertEqual(inner_acquire.call_count, 2)


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