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

"""RTL-SDR device integration for the RIA Toolkit."""

import time
import warnings
from typing import Optional

import numpy as np

try:
    from rtlsdr import RtlSdr
except ImportError as exc:  # pragma: no cover - dependency provided by end user
    raise ImportError("pyrtlsdr is required to use the RTLSDR class") from exc

from ria_toolkit_oss.data.recording import Recording
from ria_toolkit_oss.sdr.sdr import SDR, SDRParameterError


class RTLSDR(SDR):
    """SDR interface for RTL-SDR dongles using pyrtlsdr."""

    def __init__(self, identifier: Optional[str] = None):
        """
        Initialize a Pluto SDR device object and connect to the SDR hardware.

        This software supports the ADALM Pluto SDR created by Analog Devices.

        :param identifier: The value of the parameter that identifies the device.
        :type identifier: str = "192.168.3.1", "pluto.local", etc

        If no identifier is provided, it will select the first device found, with a warning.
        If more than one device is found with the identifier, it will select the first of those devices.
        """
        print(f"Initializing Pluto radio with identifier [{identifier}].")
        try:
            super().__init__()

            if identifier is None:
                self.radio = RtlSdr()
            else:
                self.radio = RtlSdr(identifier)

            self.rx_buffer_size = 256_000
            self.rx_channel = 0

            print(f"Initialized RTL-SDR with identifier [{identifier}].")

        except Exception as e:
            raise RuntimeError(f"RTL-SDR: Failed to find device with identifier '{identifier}'\nError: {e}")

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

        self.set_rx_sample_rate(sample_rate=sample_rate)
        self.set_rx_center_frequency(center_frequency=center_frequency)
        self.set_rx_gain(gain=gain, gain_mode=gain_mode)

        self.rx_channel = 0
        self.rx_buffer_size = self._calculate_optimal_buffer_size(sample_rate)
        print(f"RTL-SDR buffer: {self.rx_buffer_size} samples for {sample_rate/1e6:.1f} MS/s")

        if bias_t:
            self.set_bias_tee(True)
            time.sleep(1)

        self._rx_initialized = True
        self._tx_initialized = False

        return {"sample_rate": self.rx_sample_rate, "center_frequency": self.rx_center_frequency, "gain": self.rx_gain}

    def set_rx_sample_rate(self, sample_rate):
        """
        Set the sample rate of the receiver.
        Not callable during recording; RTL-SDR requires stream stop/restart to change sample rate.
        """
        if not ((sample_rate > 230e3 and sample_rate < 300e3) or (sample_rate > 900 and sample_rate < 3.2e6)):
            raise SDRParameterError(
                f"{self.__class__.__name__}: Sample rate {sample_rate/1e6:.3f} Msps "
                f"out of range: [{2:.3f} - {20:.3f} Msps]"
            )

        self.radio.sample_rate = float(sample_rate)
        self.rx_sample_rate = self.radio.sample_rate
        print(f"RTL RX Sample Rate = {self.radio.get_sample_rate()}")

    def set_rx_center_frequency(self, center_frequency):
        """
        Set the center frequency of the receiver.
        Not callable during recording; RTL-SDR requires stream stop/restart to change center frequency.
        """
        with self._param_lock:
            min_rate, max_rate = 25e6, 1.75e9
            if center_frequency < min_rate or center_frequency > max_rate:
                raise SDRParameterError(
                    f"{self.__class__.__name__}: Center frequency {center_frequency/1e9:.3f} GHz "
                    f"out of range: [{min_rate/1e9:.3f} - {max_rate/1e9:.3f} GHz]"
                )

            self.radio.center_freq = float(center_frequency)
            self.rx_center_frequency = self.radio.center_freq
            print(f"RTL RX Center Frequency = {self.radio.get_center_freq()}")

    def set_rx_gain(self, gain, gain_mode="absolute"):
        """
        Set the gain of the receiver. Callable during streaming.
        """
        with self._param_lock:
            available_gains = self.radio.get_gains()

            if gain is None:
                self.radio.gain = "auto"
                self.rx_gain = "auto"
            else:
                if not available_gains:
                    warnings.warn(
                        "No gain table reported by RTL-SDR; applying requested gain directly.",
                        RuntimeWarning,
                    )
                    target_gain = gain
                else:
                    min_gain = min(available_gains)
                    max_gain = max(available_gains)

                    if gain_mode == "relative":
                        if gain > 0:
                            raise SDRParameterError("When gain_mode = 'relative', gain must be < 0. This sets\
                                    the gain relative to the maximum possible gain.")
                        target_gain = max_gain + gain
                    else:
                        target_gain = gain

                    if target_gain < min_gain or target_gain > max_gain:
                        print(f"Requested gain {target_gain} dB out of range;\
                                clamping to valid span {min_gain}-{max_gain} dB.")
                        target_gain = min(max(target_gain, min_gain), max_gain)

                    target_gain = min(available_gains, key=lambda g: abs(g - target_gain))

                self.radio.set_gain(target_gain)
                self.rx_gain = self.radio.get_gain()

            print(f"RTL RX Gain = {self.radio.get_gain()}")
            print(f"Available RTL RX Gains: {available_gains}")

    def _calculate_optimal_buffer_size(self, sample_rate):
        """USB packet alignment for stability."""
        # RTL-SDR USB transfers in 16k chunks
        min_size = 16384
        max_size = 262144  # 256k

        # Target: 50ms of data per buffer
        target = int(sample_rate * 0.05)

        # Round up to 16k boundary
        size = ((target + 16383) // 16384) * 16384

        return max(min_size, min(size, max_size))

    def record(
        self,
        num_samples: Optional[int] = None,
        rx_time: Optional[int | float] = None,
    ) -> Recording:
        """
        Create a radio recording (iq samples and metadata) of a given length from the RTL-SDR.
        Either num_samples or rx_time must be provided.
        init_rx() must be called before record()

        :param num_samples: The number of samples to record.
        :type num_samples: int, optional
        :param rx_time: The time to record.
        :type rx_time: int or float, optional

        returns: Recording object (iq samples and metadata)
        """

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

        if num_samples is not None and rx_time is not None:
            raise SDRParameterError("Only input one of num_samples or rx_time")
        elif num_samples is not None:
            pass
        elif rx_time is not None:
            num_samples = int(rx_time * self.rx_sample_rate)
        else:
            raise SDRParameterError("Must provide input of one of num_samples or rx_time")

        # RTL-SDR has USB buffer limitations - use consistent 256k chunks
        # Always read full chunks to avoid USB overflow issues with partial reads
        max_samples_per_read = 262144  # 256k samples = stable chunk size
        num_full_reads = num_samples // max_samples_per_read
        remainder = num_samples % max_samples_per_read
        signal = np.array([], dtype=np.complex64)

        print("RTL-SDR Starting RX...")

        # Read full chunks
        for _ in range(num_full_reads):
            try:
                chunk = self.radio.read_samples(max_samples_per_read)
                signal = np.append(signal, chunk)
            except Exception as e:
                print(f"Error while reading samples: {e}")
                break

        # Read remainder if needed (round up to power of 2 for USB compatibility)
        if remainder > 0 and len(signal) == num_full_reads * max_samples_per_read:
            # Round up to next 16k boundary for USB stability
            padded_remainder = ((remainder + 16383) // 16384) * 16384
            try:
                chunk = self.radio.read_samples(padded_remainder)
                signal = np.append(signal, chunk[:remainder])  # Only keep what we need
            except Exception as e:
                print(f"Error while reading final chunk: {e}")

        print("RTL-SDR RX Completed.")

        metadata = {
            "source": self.__class__.__name__,
            "sample_rate": self.rx_sample_rate,
            "center_frequency": self.rx_center_frequency,
            "gain": self.rx_gain,
        }

        return Recording(data=signal, metadata=metadata)

    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().")

        print("RTL-SDR Starting RX...")
        self._enable_rx = True
        try:
            while self._enable_rx:
                samples = self.radio.read_samples(self.rx_buffer_size)
                callback(buffer=np.asarray(samples, dtype=np.complex64), metadata=None)
        finally:
            print("RTL-SDR RX Completed.")

    def _stream_tx(self, callback):  # pragma: no cover - RTL-SDR is RX only
        raise NotImplementedError("RTL-SDR does not support transmit operations")

    def init_tx(self, *args, **kwargs):  # pragma: no cover - RTL-SDR is RX only
        raise NotImplementedError("RTL-SDR does not support transmit operations")

    def tx_recording(
        self, recording: Recording | np.ndarray | list, num_samples=None, tx_time=None
    ):  # pragma: no cover - RTL-SDR is RX only
        raise NotImplementedError("RTL-SDR does not support transmit operations")

    def supports_bias_tee(self) -> bool:
        return True

    def set_bias_tee(self, enable: bool):
        self.radio.set_bias_tee(bool(enable))
        state = "enabled" if enable else "disabled"
        print(f"RTL-SDR bias tee {state}.")

    def set_clock_source(self, source):  # pragma: no cover - not applicable to RTL-SDR
        raise NotImplementedError("RTL-SDR does not support external clock configuration")

    def close(self):
        try:
            self.radio.close()
            del self.radio
        finally:
            self._enable_rx = False
            self._enable_tx = False

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