import numpy as np
import scipy as sp
import sounddevice as sd
from threading import Thread
import time
import queue

import matplotlib.animation as anim
import matplotlib.pyplot as plt

from rtlsdr import RtlSdr
import asyncio
import dsp

class BfmDemod: 
    def __init__(self, sdr:RtlSdr, stream:sd.OutputStream, audio_sample_rate:int=48000, num_samples:int=8192):
        self.sdr = sdr
        self.stream = stream
        self.audio_sample_rate = audio_sample_rate
        self.buffer = queue.Queue(10)
        self.num_samples = num_samples

        self.fig, self.ax = plt.subplots()
        self.duration = self.num_samples/self.sdr.sample_rate
        #lower_extent = self.sdr.center_freq - self.sdr.sample_rate/2
        #upper_extent = self.sdr.center_freq + self.sdr.sample_rate/2

        #self.fig, self.ax = plt.subplots()
        #self.ax.xlabel = "Frequency [Hz]"
        #self.ax.ylabel = "Amplitude [dB]"
        #self.ax.set_xlim(left=lower_extent, right=upper_extent)
    
    # Plot the spectrum of the incoming samples
    def _spectrum_update(self, frame):
        self.ax.clear()
        if self.buffer is not None:
            spectrum = np.abs(sp.fft.fftshift(sp.fft.fft(self.buffer)))**2
            spectrum_db = dsp.db(spectrum)
            self.ax.plot(spectrum_db)
        return self.ax,

    # Begin reading and demodulation
    async def start(self):
        self.stream.start()
        plt.ion()
        self.scatter = self.ax.scatter([], [])
        plt.show()

        async for samples in self.sdr.stream(self.num_samples):
            message = dsp.quad_demod(samples)
            message = dsp.lowpass(message, 16E3, self.sdr.sample_rate)
            message /= np.max(np.abs(message))
            
            time = len(message)/self.sdr.sample_rate
            num_samples = int(time*self.audio_sample_rate)
            message = sp.signal.resample(message, num_samples)

            message = message.astype(np.float32)
            #self.stream.write(message)

            spectrum = np.abs(sp.fft.fftshift(sp.fft.fft(samples)))**2
            spectrum_db = dsp.db(spectrum)
            self.ax.clear()
            self.ax.plot(spectrum_db)

            self.ax.relim()
            self.ax.autoscale_view()

            self.fig.canvas.draw()
            self.fig.canvas.flush_events()

        await sdr.stop()
        

    def stop(self):
        self.sdr.cancel_read_async()
        self.reader.join()
        self.stream.stop()