#!/usr/bin/env python3

# read-iq.py ~ (c) Alex, R2AUK 2026
# https://eax.me/2026/2026-06-08-qo-100-gqrx-afc.html

import argparse
import numpy as np
import matplotlib.pyplot as plt

parser = argparse.ArgumentParser(
    description="Read raw IQ file and compute FFT magnitude spectrum."
)
parser.add_argument(
    "--input",
    required=True,
    help="Path to the IQ file")
parser.add_argument(
    "--sampling-rate",
    type=int,
    required=True,
    help="Sampling rate (samples per second)")
parser.add_argument(
    "--fft-size",
    type=int,
    default=65536,
    help="Number of FFT points (default: 65536)")
parser.add_argument(
    "--output",
    required=True,
    help="Output plot filename")
args = parser.parse_args()

def compute_magnitude(samples, fft_size):
    samples = samples - np.mean(samples) # remove DC offset
    windowed = samples * np.hanning(fft_size)
    spectrum = np.fft.fft(windowed, n=fft_size)
    return np.fft.fftshift(np.abs(spectrum))

def raw_to_complex(raw):
    iq = np.frombuffer(raw, dtype=np.float32)
    return iq[0::2] + 1j * iq[1::2]

def build_max_spectrum(filename, fft_size):
    half = fft_size // 2
    bytes_per_sample = 8  # I + Q, float32
    bytes_full = bytes_per_sample * fft_size
    bytes_half = bytes_per_sample * half

    with open(filename, "rb") as f:
        raw = f.read(bytes_full)
        if len(raw) < bytes_full:
            return None
        chunk = raw_to_complex(raw)
        max_spectrum = compute_magnitude(chunk, fft_size)
        while True:
            raw = f.read(bytes_half)
            if len(raw) < bytes_half:
                break
            chunk = np.concatenate([chunk[half:], raw_to_complex(raw)])
            spectrum = compute_magnitude(chunk, fft_size)
            max_spectrum = np.maximum(max_spectrum, spectrum)
    return max_spectrum

max_spectrum = build_max_spectrum(args.input, args.fft_size)

sample_spacing = 1.0 / args.sampling_rate
freqs = np.fft.fftshift(np.fft.fftfreq(args.fft_size, d=sample_spacing))

fig, ax = plt.subplots(figsize=(12, 5))
ax.plot(freqs, max_spectrum, color="steelblue", linewidth=0.5)
ax.set_xlabel("Frequency (Hz)")
ax.set_ylabel("Magnitude")
ax.set_title("FFT Spectrum")
ax.grid(True, linewidth=0.3)
fig.tight_layout()
fig.savefig(args.output, dpi=150)
print(f"Saved to {args.output}")