#!/usr/bin/env python3
"""
hackrf_fm_demod.py

A HackRF frequency translator + FM demodulator.

Usage:
    python3 hackrf_fm_demod.py <center_freq_Hz> <signal_offset_Hz> <output_wav>

Example:
    # Tune HackRF to 100 MHz, translate a 200 kHz-offset signal to baseband,
    # demodulate FM, and save audio to 'out.wav'
    python3 hackrf_fm_demod.py 100e6 200e3 out.wav
"""
import sys
import numpy as np
from SoapySDR import Device, SOAPY_SDR_RX, SOAPY_SDR_CF32
from scipy.signal import firwin, lfilter, resample_poly
import soundfile as sf

def main():
    if len(sys.argv) != 4:
        print("Usage: python3 {} <center_freq_Hz> <signal_offset_Hz> <output_wav>".format(sys.argv[0]))
        sys.exit(1)

    # Parse arguments
    center_freq = float(sys.argv[1])      # in Hz
    freq_offset = float(sys.argv[2])      # in Hz (positive shifts lower frequencies down)
    output_wav  = sys.argv[3]

    # SDR parameters
    sample_rate = 2_000_000.0             # 2 MS/s
    rx_gain     = 40.0                    # adjust per your environment

    # Audio parameters
    audio_rate  = 48_000                  # 48 kHz output
    record_secs = 10                      # seconds to record

    # Initialize HackRF via SoapySDR
    args = dict(driver="hackrf")
    sdr = Device(args)
    sdr.setSampleRate(SOAPY_SDR_RX, 0, sample_rate)
    sdr.setGain      (SOAPY_SDR_RX, 0, rx_gain)
    sdr.setFrequency (SOAPY_SDR_RX, 0, center_freq)

    # Set up the RX stream
    rx_stream = sdr.setupStream(SOAPY_SDR_RX, SOAPY_SDR_CF32, [0])
    sdr.activateStream(rx_stream)

    # Design a low-pass filter for FM baseband (cutoff ~100 kHz)
    num_taps = 101
    lpf_taps = firwin(num_taps, cutoff=100e3, fs=sample_rate)

    # Buffers for audio
    audio_chunks = []
    block_size    = 16_384
    total_blocks  = int(record_secs * sample_rate / block_size)
    sample_counter = 0

    for _ in range(total_blocks):
        sr = sdr.readStream(rx_stream, [block_size])
        buff = np.array(sr[0], dtype=np.complex64)

        # Frequency translation (mix down by freq_offset)
        t = (np.arange(len(buff)) + sample_counter) / sample_rate
        buff *= np.exp(-2j * np.pi * freq_offset * t)
        sample_counter += len(buff)

        # FM demodulation: phase difference
        phase_diff = np.angle(buff[1:] * np.conj(buff[:-1]))

        # Filter the demodulated signal
        filtered = lfilter(lpf_taps, 1.0, phase_diff)

        # Resample directly to audio_rate
        audio = resample_poly(filtered, audio_rate, int(sample_rate))
        audio_chunks.append(audio.astype(np.float32))

    # Clean up SDR
    sdr.deactivateStream(rx_stream)
    sdr.closeStream(rx_stream)

    # Concatenate and write to WAV
    audio_data = np.concatenate(audio_chunks)
    sf.write(output_wav, audio_data, audio_rate)
    print(f"✅ Demodulated audio saved to '{output_wav}'")

if __name__ == "__main__":
    main()