"""Sequential capture workflow for preprocess asset dataset creation.""" from __future__ import annotations from contextlib import suppress from dataclasses import dataclass import time import numpy as np from python_app.hardware_full.multi_device_service import MultiDeviceLibreVnaService from python_app.hardware_full.single_radar_service import create_single_radar_service from python_app.hardware_full.switch_service import SwitchService from python_app.models.dataset_model import ComboKey, SweepCollection, TraceData from python_app.models.run_config_model import ComboModel, RunConfigModel from python_app.storage.npz_store import NpzStore, radar_key_from_config MULTI_DEVICE_MANUAL_CAPTURE_KINDS = frozenset({"s21_calibration", "s11_open", "s11_short", "s11_load"}) @dataclass(slots=True) class SequentialCaptureState: """Immutable view of sequential capture progress.""" kind: str set_name: str captured_count: int total_count: int current_combo: ComboModel | None can_undo: bool is_complete: bool variant_count: int = 1 supports_batch_capture: bool = True class SequentialCaptureSession: """Manage hardware and switch stepping for full combo capture sequence.""" def __init__(self, config: RunConfigModel, kind: str, set_name: str) -> None: """Create capture session for one preprocess asset set.""" if kind not in {"s21_calibration", "s21_reference", "s11_open", "s11_short", "s11_load", "s11_reference"}: raise RuntimeError(f"Unsupported capture kind: {kind}") if not set_name: raise RuntimeError("Set name is required") self._config = config self._kind = kind self._set_name = set_name self._is_multi_device = config.is_multi_device self._manual_multi_device_capture = self._is_multi_device and kind in MULTI_DEVICE_MANUAL_CAPTURE_KINDS self._combos = ( RunConfigModel.build_multi_device_virtual_combos() if self._is_multi_device else RunConfigModel.build_full_combos(config.input_switch.positions, config.output_switch.positions) ) if not self._combos: raise RuntimeError("No switch combinations available for capture") self._traces: list[TraceData] = [] self._next_index = 0 self._opened = False if self._is_multi_device: self._radar = MultiDeviceLibreVnaService( master_serial=config.radar.serial, slave_serials=list(config.radar.multi_device.slave_serials), force_external_reference=config.radar.multi_device.force_external_reference, recovery_attempts=config.radar.multi_device.recovery_attempts, backend_mode=config.radar.driver_mode, ) self._input_switch = None self._output_switch = None else: self._radar = create_single_radar_service(config) self._input_switch = SwitchService( name=config.input_switch.name, positions=config.input_switch.positions, mode=config.input_switch.driver_mode, driver=config.input_switch.driver, gpio_chip=config.input_switch.gpio_chip, pin_a=config.input_switch.pin_a, pin_b=config.input_switch.pin_b, invert_logic=config.input_switch.invert_logic, ) self._output_switch = SwitchService( name=config.output_switch.name, positions=config.output_switch.positions, mode=config.output_switch.driver_mode, driver=config.output_switch.driver, gpio_chip=config.output_switch.gpio_chip, pin_a=config.output_switch.pin_a, pin_b=config.output_switch.pin_b, invert_logic=config.output_switch.invert_logic, ) @property def kind(self) -> str: """Return canonical preprocess asset key for this capture session.""" return self._kind @property def set_name(self) -> str: """Return destination set name.""" return self._set_name def open(self) -> None: """Open radar and switch resources.""" if self._opened: return self._opened = True try: self._radar.open() self._radar.configure(self._config.radar.sweep) if self._input_switch is not None: self._input_switch.open() if self._output_switch is not None: self._output_switch.open() except Exception: self.close() raise def close(self) -> None: """Close all opened hardware resources.""" with suppress(Exception): if self._output_switch is not None: self._output_switch.close() with suppress(Exception): if self._input_switch is not None: self._input_switch.close() with suppress(Exception): self._radar.close() self._opened = False def state(self) -> SequentialCaptureState: """Return current progress snapshot.""" current_combo = self._current_combo() return SequentialCaptureState( kind=self._kind, set_name=self._set_name, captured_count=len(self._traces), total_count=len(self._combos), current_combo=current_combo, can_undo=bool(self._traces), is_complete=self.is_complete(), supports_batch_capture=not self._manual_multi_device_capture, ) def capture_current_combo(self) -> TraceData: """Capture one trace for current combo and advance sequence cursor.""" if not self._opened: raise RuntimeError("Capture session is not opened") combo = self._current_combo() if combo is None: raise RuntimeError("Capture session is already complete") if self._is_multi_device: collection = self._radar.acquire_collection(collection_id=1) if not collection.traces: raise RuntimeError("Multi-device capture returned no traces") if self._manual_multi_device_capture: trace = select_trace_for_combo(collection, combo) self._traces.append(trace) self._next_index += 1 return trace self._traces.extend(collection.traces) self._next_index = len(self._combos) return collection.traces[-1] assert self._input_switch is not None assert self._output_switch is not None self._output_switch.switch_to(combo.output) self._input_switch.switch_to(combo.input) if self._config.runtime.settling_ms > 0: time.sleep(self._config.runtime.settling_ms / 1000.0) sweep = self._radar.acquire() trace = TraceData( combo=ComboKey(input_pos=combo.input, output_pos=combo.output), frequency_hz=np.asarray(sweep.x, dtype=np.float32), s11=np.asarray(sweep.trace("s11"), dtype=np.complex64), s21=np.asarray(sweep.trace("s21"), dtype=np.complex64), ) self._traces.append(trace) self._next_index += 1 return trace def undo_last_capture(self) -> TraceData: """Remove the most recently captured trace and rewind cursor by one combo.""" if not self._opened: raise RuntimeError("Capture session is not opened") if not self._traces or self._next_index <= 0: raise RuntimeError("No captured combo is available to undo") if self._is_multi_device and not self._manual_multi_device_capture: if len(self._traces) != len(self._combos): raise RuntimeError("Capture session state is inconsistent; multi-device trace matrix is incomplete") removed_trace = self._traces[-1] self._traces.clear() self._next_index = 0 return removed_trace expected_combo = self._combos[self._next_index - 1] removed_trace = self._traces[-1] if ( int(removed_trace.combo.input_pos) != int(expected_combo.input) or int(removed_trace.combo.output_pos) != int(expected_combo.output) ): raise RuntimeError("Capture session state is inconsistent; last trace does not match rewind combo") self._next_index -= 1 self._traces.pop() return removed_trace def last_captured_trace(self) -> TraceData | None: """Return the most recently captured trace, if any.""" if not self._traces: return None return self._traces[-1] def captured_traces(self) -> list[TraceData]: """Return captured traces in acquisition order.""" return list(self._traces) def is_complete(self) -> bool: """Return `True` when all combos were captured.""" return self._next_index >= len(self._combos) def finalize(self, store: NpzStore) -> tuple[str, SweepCollection]: """Persist completed capture into store and return radar key + collection.""" if not self.is_complete(): raise RuntimeError("Capture session is not complete") collection = SweepCollection( collection_id=1, monotonic_ns=time.monotonic_ns(), traces=list(self._traces), ) radar_key = radar_key_from_config( model_name=self._config.radar.model, serial=self._config.radar.serial, sweep_start_hz=self._config.radar.sweep.start_hz, sweep_stop_hz=self._config.radar.sweep.stop_hz, sweep_points=self._config.radar.sweep.points, ifbw_hz=self._config.radar.sweep.if_bandwidth_hz, power_dbm=self._config.radar.sweep.power_dbm, extra_serials=self._config.radar_key_extra_parts() or None, ) store.save_set(self._kind, radar_key, self._set_name, collection) return radar_key, collection def _current_combo(self) -> ComboModel | None: """Return next combo to capture, or `None` if session is complete.""" if self._next_index >= len(self._combos): return None return self._combos[self._next_index] def select_trace_for_combo(collection: SweepCollection, combo: ComboModel) -> TraceData: """Return the trace matching a virtual combo from a full multi-device capture.""" for trace in collection.traces: if ( int(trace.combo.input_pos) == int(combo.input) and int(trace.combo.output_pos) == int(combo.output) ): return trace raise RuntimeError(f"Multi-device capture is missing trace for input={combo.input}, output={combo.output}")