ria-toolkit-oss/src/ria_toolkit_oss/sdr/thinkrf.py

"""ThinkRF integration for the RIA toolkit."""

from typing import Any, Dict, Optional

import numpy as np

try:
    from pyrf.devices.thinkrf import WSA
except ImportError as exc:  # pragma: no cover - optional dependency
    raise ImportError(
        "pyrf is required to use the ThinkRF integration. " "Install with: pip install ria-toolkit-oss[thinkrf]"
    ) from exc
except SyntaxError as exc:  # pragma: no cover - Python 2/3 compatibility issue
    import sys
    from pathlib import Path

    # pyrf ships with Python 2 syntax - try to auto-fix it
    print("\033[93mWARNING: pyrf has Python 2 syntax. Attempting automatic fix...\033[0m")
    try:
        from lib2to3.refactor import RefactoringTool, get_fixers_from_package

        import pyrf

        thinkrf_path = Path(pyrf.__file__).resolve().parent / "devices" / "thinkrf.py"
        print(f"Fixing: {thinkrf_path}")

        fixers = get_fixers_from_package("lib2to3.fixes")
        tool = RefactoringTool(fixers)
        tool.refactor_file(str(thinkrf_path), write=True)

        print("\033[92m✅ Fixed pyrf for Python 3. Please restart Python/reload the module.\033[0m")
        print("Or run: python -m ria_toolkit_oss.sdr.thinkrf_fix")
        sys.exit(1)  # Exit so user can reload
    except Exception as fix_exc:
        print(f"\033[91m❌ Auto-fix failed: {fix_exc}\033[0m")
        print("Manual fix: Run `python scripts/fix_pyrf_python3.py` from ria-toolkit-oss directory")
        raise exc

from ria_toolkit_oss.sdr.sdr import SDR, SDRParameterError


class ThinkRF(SDR):
    """SDR adapter for ThinkRF analyzers using the PyRF API."""

    BASE_SAMPLE_RATE = 125_000_000
    SUPPORTED_DECIMATIONS = (1, 2, 4, 8, 16, 32, 64, 128, 256, 512, 1024)
    MAX_ONBOARD_SAMPLES = 33_500_000  # Confirmed: 512 packets @ dec 1 = 33.5M samples (268ms)
    DEFAULT_SPP = 65504  # VRT packet size (samples per packet)

    def __init__(self, identifier: Optional[str] = None):
        super().__init__()

        if identifier is None:
            raise SDRParameterError("ThinkRF requires an IP address or hostname identifier")

        self.identifier = identifier
        try:
            self.radio = WSA()
            self.radio.connect(identifier)
            self.radio.request_read_perm()
            print(f"Connected to ThinkRF at [{identifier}].")
        except Exception as exc:
            print(f"Failed to connect to ThinkRF at [{identifier}].")
            raise exc

        self.configure_frontend()
        self._last_context: Optional[Any] = None

    def configure_frontend(
        self,
        *,
        rfe_mode: str = "ZIF",
        attenuation: int = 0,
        gain_profile: str = "HIGH",
        trigger_config: Optional[Dict[str, Any]] = None,
        samples_per_packet: int = 65504,
        packets_per_block: int = 1,
        capture_mode: str = "block",
        stream_id: int = 1,
        min_stream_decimation: int = 16,
    ) -> None:
        """Persist settings applied during the next RX initialisation.

        ``capture_mode`` selects between buffered ``"block"`` captures that use
        the analyser's onboard RAM and ``"stream"`` captures that push data over
        GigE in real time. Streaming requires a sufficiently large decimation to
        keep within the link budget; ``min_stream_decimation`` forms the lower
        bound.
        """

        mode = capture_mode.lower()
        if mode not in {"block", "stream"}:
            raise SDRParameterError("capture_mode must be either 'block' or 'stream'")

        self._rfe_mode = rfe_mode
        self._attenuation = int(max(0, min(attenuation, 30)))
        self._gain_profile = gain_profile.upper()
        self._trigger_config = trigger_config
        self._samples_per_packet = int(samples_per_packet)
        self._packets_per_block = max(1, int(packets_per_block))
        self._capture_mode = mode
        self._stream_id = int(stream_id)
        self._min_stream_decimation = max(1, int(min_stream_decimation))
        self._streaming_active = False

    def init_rx(
        self,
        sample_rate: int | float,
        center_frequency: int | float,
        gain: int,
        channel: int,
        gain_mode: Optional[str] = "absolute",
        decimation: Optional[int] = None,
    ):
        if channel not in (0, None):
            raise SDRParameterError("ThinkRF supports only channel 0 for RX.")

        stream_mode = getattr(self, "_capture_mode", "block") == "stream"
        actual_decimation, _ = self.set_rx_sample_rate(
            sample_rate=sample_rate, decimation=decimation, stream_mode=stream_mode
        )

        self.radio.reset()
        self.radio.scpiset(":SYSTEM:FLUSH")
        try:
            self.radio.scpiset(":TRACE:STREAM:STOP")
        except Exception:
            pass

        self.radio.rfe_mode(self._rfe_mode)
        self.set_rx_center_frequency(center_frequency=center_frequency)
        self.set_rx_gain(gain=gain, gain_mode=gain_mode, actual_decimation=actual_decimation)

        self.radio.decimation(actual_decimation)
        if stream_mode:
            self.radio.scpiset(f":SENSE:DECIMATION {actual_decimation}")
        trigger = self._trigger_config or self._default_trigger(center_frequency)
        self.radio.trigger(trigger)

        self.radio.scpiset(f":TRACE:SPP {self._samples_per_packet}")
        if stream_mode:
            self._streaming_active = False
        else:
            print(
                f"ThinkRF: Configuring block capture - SPP={self._samples_per_packet}, PPB={self._packets_per_block}"
            )
            self.radio.scpiset(f":TRACE:BLOCK:PACKETS {self._packets_per_block}")
            self.radio.scpiset(":TRACE:BLOCK:DATA?")

        self.rx_buffer_size = self._samples_per_packet
        self.rx_channel = 0

        self._rx_initialized = True
        self._tx_initialized = False

    def set_rx_sample_rate(self, sample_rate, decimation, stream_mode):
        """
        Set the sample rate of the receiver.
        Not callable during recording; ThinkRF requires stream stop/restart to change sample rate.
        """
        # Enforce sample rate / decimation
        # Note: decimation parameter takes precedence if provided
        actual_decimation, actual_sample_rate = self.enforce_sample_rate(sample_rate, decimation)

        if stream_mode and actual_decimation < self._min_stream_decimation:
            enforced = self._min_stream_decimation
            print(
                "Requested ThinkRF sample rate exceeds typical GigE throughput; "
                f"enforcing decimation {enforced} for streaming."
            )
            actual_decimation = enforced
            actual_sample_rate = self.BASE_SAMPLE_RATE / actual_decimation

        self._decimation = actual_decimation
        self.rx_sample_rate = actual_sample_rate
        print(f"ThinkRF RX Sample Rate = {actual_sample_rate}")

        return actual_decimation, actual_sample_rate

    def set_rx_center_frequency(self, center_frequency):
        """
        Set the center frequency of the receiver. Callable during streaming.
        """
        with self._param_lock:
            self.radio.freq(int(center_frequency))
            self.rx_center_frequency = self.radio.freq
            print(f"ThinkRF RX Center Frequency = {self.radio.freq}")

    def set_rx_gain(self, gain, gain_mode, actual_decimation):
        attenuation = self._attenuation if gain is None else int(gain)  # gain
        attenuation = max(0, min(attenuation, 30))
        self.radio.attenuator(attenuation)

        gain_profile = self._gain_profile
        if gain_mode and isinstance(gain_mode, str) and gain_mode.upper() in {"LOW", "MEDIUM", "HIGH", "VLOW"}:
            gain_profile = gain_mode.upper()
        self.radio.gain(gain_profile.lower())  # WSA.gain() expects lowercase

        self.rx_gain = {
            "attenuation_dB": attenuation,
            "profile": gain_profile,
            "decimation": actual_decimation,
            "rfe_mode": self._rfe_mode,
            "spp": self._samples_per_packet,
            "ppb": self._packets_per_block,
        }

    def _stream_rx(self, callback):
        if not self._rx_initialized:
            raise RuntimeError("RX was not initialized. init_rx() must be called before _stream_rx() or record().")

        self._enable_rx = True
        packets_processed = 0
        stream_mode = getattr(self, "_capture_mode", "block") == "stream"

        if stream_mode and not self._streaming_active:
            try:
                self.radio.scpiset(f":TRACE:STREAM:START {self._stream_id}")
                self._streaming_active = True
            except Exception as exc:
                print(f"Failed to start ThinkRF stream: {exc}")
                return

        print("ThinkRF Starting RX...")
        while self._enable_rx:
            packet = self._safe_read(stream_mode, packets_processed)

            if packet is None:
                # No more packets available
                if not stream_mode and packets_processed >= self._packets_per_block:
                    # Finished reading block
                    break
                continue

            if packet.is_context_packet():
                self._last_context = packet
                continue

            if not packet.is_data_packet():
                # Unknown packet type - skip
                continue

            metadata = metadata = self._extract_metadata(packet)
            complex_buffer = self._extract_iq(packet)
            if complex_buffer is None:
                continue

            # Send packet data to callback (accumulation handled by parent)
            callback(buffer=complex_buffer, metadata=metadata)
            packets_processed += 1

            # In block mode, stop after receiving all packets in the block
            if not stream_mode and packets_processed >= self._packets_per_block:
                # Got all packets for this block
                break

        print("ThinkRF RX Completed.")
        if stream_mode and self._streaming_active:
            self._stop_stream()

        self.radio.scpiset(":SYSTEM:FLUSH")

    def _safe_read(self, stream_mode, packets_processed):
        packet = None
        try:
            packet = self.radio.read()
        except Exception as e:
            # In block mode, reaching end of block can cause exceptions
            if not stream_mode and packets_processed > 0:
                # We got some packets in block mode, so finish gracefully
                print(f"ThinkRF: Block read complete ({packets_processed} packets received)")
            else:
                print(f"ThinkRF read error: {e}")
        return packet

    def _extract_iq(self, packet):
        # packet.data is an iterable IQData object that yields (I, Q) tuples
        # Convert to numpy array: collect all [I, Q] pairs
        try:
            iq_pairs = list(packet.data)
            if not iq_pairs:
                return None
            iq_array = np.array(iq_pairs, dtype=np.float32)
            return (iq_array[:, 0] + 1j * iq_array[:, 1]).astype(np.complex64)
        except Exception as e:
            print(f"Error extracting IQ from packet.data: {e}")
            return None

    def _extract_metadata(self, packet):
        if not hasattr(packet, "fields"):
            return None
        metadata = packet.fields
        if metadata.get("sample_loss"):
            print("\033[93mWarning: ThinkRF sample overflow detected\033[0m")
        return metadata

    def _stop_stream(self):
        try:
            self.radio.scpiset(":TRACE:STREAM:STOP")
        except Exception:
            pass
        self._streaming_active = False

    def init_tx(
        self,
        sample_rate: int | float,
        center_frequency: int | float,
        gain: int,
        channel: int,
        gain_mode: Optional[str] = "absolute",
    ):
        raise NotImplementedError("ThinkRF devices do not support transmit operations")

    def _stream_tx(self, callback):
        raise NotImplementedError("ThinkRF devices do not support transmit operations")

    def set_clock_source(self, source):
        raise NotImplementedError("ThinkRF clock configuration is not implemented")

    def close(self):
        try:
            self.radio.scpiset(":TRACE:STREAM:STOP")
        except Exception:  # pragma: no cover - best effort cleanup
            pass
        try:
            self.radio.scpiset(":SYSTEM:FLUSH")
        except Exception:
            pass
        try:
            self.radio.disconnect()
        finally:
            self._enable_rx = False
            self._enable_tx = False
            print(f"Disconnected from ThinkRF at [{self.identifier}].")

    def supports_bias_tee(self) -> bool:
        return False

    def set_bias_tee(self, enable: bool):  # pragma: no cover - interface compliance
        raise NotImplementedError("ThinkRF radios do not expose a controllable bias-tee")

    def _derive_decimation(self, target_sample_rate: int | float) -> int:
        """
        Derive decimation from target sample rate.
        Always rounds DOWN decimation (UP sample rate) to meet or exceed user's requested rate.

        Example: 30 MS/s requested → dec 4 (31.25 MS/s), NOT dec 8 (15.625 MS/s)
        """
        if not target_sample_rate:
            return 1
        requested = float(target_sample_rate)
        if requested >= self.BASE_SAMPLE_RATE:
            return 1

        desired_decimation = self.BASE_SAMPLE_RATE / requested

        # Round DOWN decimation (UP sample rate) to meet or exceed requested rate
        # Find largest decimation that gives sample rate >= requested
        valid_decimations = [d for d in self.SUPPORTED_DECIMATIONS if d <= desired_decimation]

        if valid_decimations:
            # Use largest valid decimation (gives sample rate >= requested)
            best = max(valid_decimations)
        else:
            # Requested rate too low, use minimum decimation (max sample rate)
            best = self.SUPPORTED_DECIMATIONS[0]

        return int(best)

    def enforce_sample_rate(
        self, requested_sample_rate: int | float, decimation: Optional[int] = None
    ) -> tuple[int, float]:
        """
        Enforce valid sample rate and decimation.

        If decimation is provided, it takes precedence.
        Otherwise, derive decimation from requested sample rate.

        Returns:
            (decimation, actual_sample_rate)
        """
        if decimation is not None:
            # Decimation provided - validate and use it
            if decimation not in self.SUPPORTED_DECIMATIONS:
                # Round to nearest supported
                decimation = min(self.SUPPORTED_DECIMATIONS, key=lambda d: abs(d - decimation))
                print(f"ThinkRF: Requested decimation not supported. Using decimation={decimation}")
        else:
            # Derive from sample rate
            decimation = self._derive_decimation(requested_sample_rate)

        actual_sample_rate = self.BASE_SAMPLE_RATE / decimation

        if abs(actual_sample_rate - requested_sample_rate) > 1e3:  # More than 1 kHz difference
            print(
                f"ThinkRF: Requested {requested_sample_rate/1e6:.2f} MS/s → \
                    Using decimation={decimation} ({actual_sample_rate/1e6:.2f} MS/s)"
            )

        return decimation, actual_sample_rate

    def calculate_spp_ppb(self, num_samples: int, spp: Optional[int] = None) -> tuple[int, int]:
        """
        Calculate optimal SPP (samples per packet) and PPB (packets per block).

        Strategy:
        - Maximize SPP (use DEFAULT_SPP) unless num_samples < DEFAULT_SPP
        - Calculate PPB to get as close as possible to num_samples
        - Actual captured samples = SPP * PPB (may exceed num_samples slightly)

        Args:
            num_samples: Desired number of samples
            spp: Override SPP (for advanced users, not recommended)

        Returns:
            (spp, ppb)
        """
        if spp is not None:
            # User override - use as-is
            actual_spp = max(1, int(spp))
        else:
            # Maximize SPP unless samples requested is smaller
            if num_samples < self.DEFAULT_SPP:
                actual_spp = num_samples
            else:
                actual_spp = self.DEFAULT_SPP

        # Calculate PPB to get close to num_samples
        ppb = max(1, int(np.ceil(num_samples / actual_spp)))

        actual_samples = actual_spp * ppb
        if actual_samples != num_samples:
            print(
                f"ThinkRF: Requested {num_samples} samples → Capturing {actual_samples} (SPP={actual_spp}, PPB={ppb})"
            )

        return actual_spp, ppb

    def check_ram_limit(self, num_samples: int, decimation: int) -> None:
        """
        Check if requested capture exceeds onboard RAM limits.

        Raises warning if exceeds MAX_ONBOARD_SAMPLES at low decimations.
        For decimation 1 or 2, block captures are limited by onboard RAM.
        """
        if decimation <= 2 and num_samples > self.MAX_ONBOARD_SAMPLES:
            raise SDRParameterError(
                f"ThinkRF: Cannot capture {num_samples} samples at decimation {decimation}. "
                f"Onboard RAM limit is ~{self.MAX_ONBOARD_SAMPLES} samples for dec 1/2. "
                f"Either reduce num_samples or use stream mode (increase decimation to >=4)."
            )

    def _default_trigger(self, center_frequency: int | float) -> Dict[str, Any]:
        span = 40_000_000
        half = span // 2
        return {
            "type": "NONE",
            "fstart": int(center_frequency) - half,
            "fstop": int(center_frequency) + half,
            "amplitude": -100,
        }

    def supports_dynamic_updates(self) -> dict:
        return {"center_frequency": True, "sample_rate": False, "gain": False}