# -*- coding: utf-8 -*-
"""
    Laborator PS cu RTL-SDR
    Cod initial de Matei Simtinica
    Unele modificari de Marios Choudary

    Nota: trebuie instalate cateva lucruri pentru ca acest script sa functioneze:

    1. Python3 și pip pentru Python3 (acest script se bazeaza pe folosirea Python3)

    2. Install librtlsdr (pentru a putea folosi dispozitivul RTL-SDR)

    Pe MAC OS:
    sudo port install rtl-sdr

    Sau pentru brew:
    https://macappstore.org/librtlsdr/

    Pe Linux
    sudo apt-get install  librtlsdr-dev

    3. Install matplotlib, scipy, sounddevice pentru Python3:
    pip3 install pyrtlsdr matplotlib scipy sounddevice

"""

from rtlsdr import RtlSdr # comentati asta daca nu folositi dispozitiv RTL
import matplotlib.pyplot as plt
import numpy as np
import scipy.signal as signal
#from IPython.display import Audio
from scipy.fftpack import fft
import sounddevice as sd
from scipy.io.wavfile import write
import math
import cmath

plt.rcParams['figure.dpi'] = 170
plt.rcParams['figure.figsize'] = (8, 4)

# Setati aici 1 daca folositi un RTL-SDR, 0 altfel
use_sdr = 0

F_station = int(88.5e6) # Trinitas FM
Fs = 2280000            # frecventa de esantionare (samplerate după downsampling intern)
N = 8192000             # numarul esantioanelor (t = N / Fs)

if use_sdr == 0:
    x1 = np.load('x1.npy')
else:
    # initializati device-ul RTL-SDR
    sdr = RtlSdr()

    # setati parametri device-ului
    sdr.sample_rate = Fs
    sdr.center_freq = F_station
    sdr.gain = 'auto'

    # cititi semnalul brut
    x1 = sdr.read_samples(N)
    np.save('x1.npy', x1)

    # eliberati resursele device-ului
    sdr.close()

# plotati spectrograma semnalului cu specgram
# Pentru detalii, vedeti aici: https://matplotlib.org/stable/api/_as_gen/matplotlib.pyplot.specgram.html
fig, ax = plt.subplots()
ax.specgram(x1, NFFT=2**10, Fs=Fs)
ax.set_title('Spectrograma semnalului (X1)')
ax.set_ylim(-Fs/2, Fs/2)
ax.ticklabel_format(style='plain')
ax.set_xlabel('Timp (s)')
ax.set_ylabel('Frecvență')
plt.show(block=False)
# sau salvati imaginea ca png, daca nu va merge plt.show
# plt.savefig('fig1.png')

print('Apasati o tastă pentru a continua...\n')
input()

# Acum să luăm eșantioane de la o frecvență puțin depărtată de centrul frecvenței anterioare
# Acest sceariu seamănă cu cazul în care facem întâi o eșantionare intermediară (IF=intermediate frequency)
# și apoi aplicăm downconvert ca să ajungem în baseband (ca să putem face filtrare și downsampling ușor)
F_station_bad = int(88.3e6) # La 200 kHz de Trinitas FM (sau la o distanță față o altă stație)

if use_sdr == 0:
    x2 = np.load('x2.npy')
else:
    # initializati device-ul RTL-SDR
    sdr = RtlSdr()

    # setati parametri device-ului
    sdr.sample_rate = Fs
    sdr.center_freq = F_station_bad
    sdr.gain = 'auto'

    # cititi semnalul brut
    x2 = sdr.read_samples(N)
    np.save('x2.npy', x2)

    # eliberati resursele device-ului
    sdr.close()


# plotati spectrograma semnalului
# Pentru detalii, vedeti aici: https://matplotlib.org/stable/api/_as_gen/matplotlib.pyplot.specgram.html
plt.figure()
plt.specgram(x2, NFFT=2**10, Fs=Fs)
plt.title('Spectrograma semnalului (X2)')
plt.ylim(-Fs/2, Fs/2)
plt.ticklabel_format(style='plain')
plt.show(block=False)
plt.xlabel('Timp')
plt.ylabel('Frecvență')
# sau salvati imaginea ca png, daca nu va merge plt.show
# plt.savefig('fig2.png')

# Q1: ce observați între cele două spectrograme (pt x1 și pt x2) ? Care este diferența ? De ce ?

print('Apasati o tastă pentru a continua...\n')
input()

# Acum să vedem spectrul celor două semnale cu FFT (varianta rapidă pentru DFT):
y1 = fft(x1)
nf1 = len(y1)
nf1_2 = int(nf1/2)
xf1 = np.linspace(0, Fs, nf1)
y2 = fft(x2)
nf2 = len(y2)
nf2_2 = int(nf2/2)
xf2 = np.linspace(0, Fs, nf2)

fig, (ax1, ax2) = plt.subplots(2)
fig.suptitle('Spectrul (FFT) pentru x1 (sus) și x2 (jos)')
ax1.plot(xf1[0:nf1_2], np.abs(y1[0:nf1_2]))
ax1.grid()
ax2.plot(xf2[0:nf2_2], np.abs(y2[0:nf2_2]))
ax2.set_xlabel('Frecvența')
ax2.grid()
plt.show(block=False)
# sau salvati imaginea ca png, daca nu va merge plt.show
# plt.savefig('fig6.png')

print('Apasati o tastă pentru a continua...\n')
input()

# Să folosim IQ downconvert pentru a aduce semnalul util din x2 în origine (baseband)
# Ar trebui să obținem în cele din urmă un semnal/spectru asemănător cu cel pentru x1
# Reminder: pentru a coborâ spectrul din f_c în 0 trebuie să înmulțim în timp cu o
# exponențială de frecvență f_c (în domeniul digital este raportul f_c/fs)
F_c = F_station - F_station_bad # sau puteți să puneți manual uitându-vă la spectru
lenx2 = len(x2)
x2n = np.zeros(lenx2, dtype=np.complex128)
fr = F_c / Fs
# TODO: înmulțiți aici fiecare element  x2[i] cu exponențiala e^(-j*2*pi*fr*i)
# pentru exponențială complexă folosiți cmath.exp
for i in range(lenx2):
    x2n[i] = 0

# Acum să vedem spectrul celor două semnale cu FFT (x2 și x2 după downconvert):
y2n = fft(x2n)
nf2n = len(y2n)
nf2n_2 = int(nf2n/2)
xx2n = np.linspace(0, Fs, nf2n)

fig, (ax1, ax2) = plt.subplots(2)
fig.suptitle('Spectrul (FFT) pentru x2 (sus) și x2 după downconvert (jos)')
ax1.plot(xf2[0:nf2_2], np.abs(y2[0:nf2_2]))
ax1.grid()
ax2.plot(xx2n[0:nf2n_2], np.abs(y2n[0:nf2n_2]))
ax2.set_xlabel('Frecvența')
ax2.grid()
plt.show(block=False)
# sau salvati imaginea ca png, daca nu va merge plt.show
# plt.savefig('fig6.png')

print('Apasati o tastă pentru a continua...\n')
input()

# TODO: Să facem și spectrograma pentru acest nou semnal (x2n) și să comparăm cu cea pentru x1
plt.figure()
plt.specgram(x2n, NFFT=2**10, Fs=Fs)
plt.title('Spectrograma semnalului (X2) după downconvert')
plt.ylim(-Fs/2, Fs/2)
plt.ticklabel_format(style='plain')
plt.show(block=False)
plt.xlabel('Timp')
plt.ylabel('Frecvență')
# sau salvati imaginea ca png, daca nu va merge plt.show
# plt.savefig('fig2.png')

print('Apasati o tastă pentru a continua...\n')
input()

# Folosiți acest nou semnal în loc de x1 (BONUS: încercați lab cu ambele variante: cu x1 original și cu acesta):
# TODO: scoați comentariul după ce ați implementat îmnulțirea de la IQ downconvert
#x1 = x2n


# semnalele FM sunt difuzate pe o latime de banda de 200kHz
f_bw = 200000

# filtrarea semnalului x1
# TODO 1: creati un filtru trece-jos digital folosind scipy.signal.firwin
# 1. Selectati un numar de coeficienti (taps) ai filtrului (incercati 8, 16, 32)
# 2. Calculati frecventa de taiere (cutoff), ca raport intre
#    frecventa maxima de interes (e.g. f_bw) si fs/2
# 3. alegeti o fereastra pentru atenuarea DFT leakage rezultat din filtru (e.g. 'hamming')
# 4. folositi functia firwin din scipy.signal pentru a obtine coeficientii unui filtru FIR
# 5. filtrati semnalul initial (x1) cu coeficientii obtinuti anterior, folosind functia lfilter din scipy.signal
ntaps = 8 # incercati mai multe variante: 8, 16, 32, vedeti apoi diferentele, prin spectrograma urmatoare
fcut = 2*f_bw/Fs  # calculati cat este frecventa de cutoff necesara, ca o valoare intre 0 si 1, unde 1 corespunde lui fs/2
b = 0 # folositi signal.firwin pentru a obtine coeficientii (b) ai unui filtru trece-jos:
      # https://docs.scipy.org/doc/scipy/reference/generated/scipy.signal.firwin.html#scipy.signal.firwin
x1f = signal.lfilter(b, 1, x1)

# Să vedem spectrul semnalului filtrat
y1f = fft(x1f)
nf1f = len(y1f)
nf1f_2 = int(nf1f/2)
xx1f = np.linspace(0, Fs, nf1f)

fig, ax = plt.subplots()
fig.suptitle('Spectrul (FFT) pentru x1 filtrat')
ax.plot(xx1f[0:nf1f_2], np.abs(y1f[0:nf1f_2]))
ax.set_xlabel('Frecvența')
ax.grid()
plt.show(block=False)


# plotati spectrograma semnalului dupa filtrare
# TODO 2: faceti spectrograma, ca mai sus, dar pentru semnalul x1f
# Observati diferentele intre spectrograma semnalului initial si cel filtrat
# pentru diferite valori ale numarului de coeficienti ai filtrului (ntaps)

fig, ax = plt.subplots()
ax.specgram(x1f, NFFT=2**10, Fs=Fs)
ax.set_title('Spectrograma semnalului filtrat (x1f)')
ax.set_ylim(-Fs/2, Fs/2)
ax.ticklabel_format(style='plain')
ax.set_xlabel('Timp (s)')
ax.set_ylabel('Frecvență')
plt.show(block=False)


print('Apasati o tastă pentru a continua...\n')
input()

# decimare
# TODO 3: decimati (faceti subsampling) semnalul, astfel incat sa aveti noua frecventa
# de esantionare f_bw   (i.e. fs' = f_bw)
dec_rate = 0 # calculati factorul de decimare (raportul intre frecvente) - important: trebuie să fie un întreg (int(...))
xd = signal.decimate(x1f, dec_rate)

# noua frecventa de esantionare
Fs_new = Fs / dec_rate

# plotati spectrograma semnalului decimat
# TODO 4: afisati spectrograma pentru semnalul obtinut dupa decimare

fig, ax = plt.subplots()
ax.specgram(xd, NFFT=2**10, Fs=Fs_new)
ax.set_title('Spectrograma semnalului filtrat decimat (xd)')
ax.set_ylim(-Fs_new/2, Fs_new/2)
ax.ticklabel_format(style='plain')
ax.set_xlabel('Timp (s)')
ax.set_ylabel('Frecvență')
plt.show(block=False)

print('Apasati o tastă pentru a continua...\n')
input()

# demodulati semnalul FM
# TODO 5: demodulati semnalul FM prin folosirea unui discriminator simplu de frecvente
# 1. Folosind semnalul decimat (xd), obtineti o copie a acestuia dar intarziat cu un element (xd_int = xd[1:])
# 2. Calculati conjugatul acestui semnal intarziat (puteti folosi numpy.conj)
# 3. Inmultiti semnalul decimat (xd) cu semnalul intarziat conjugat
# 4. calculati faza (unghiul) semnalului rezultat din inmultirea precedenta (puteti folosi numpy.angle)
# Aceasta faza va fi direct proportionala cu semnalul transmis, asa ca poate fi folosita direct.
xd_int = xd[1:]
xd_int_conj = 0 # calculati conjugatul cu numpy.conj
xx = xd[:-1] * xd_int_conj
x3 = 0 # calculati faza semnalului inmultit (xx) folosind numpy.angle

# vizualizati semnalul FM obtinut prin afisarea puterii spectrale (power spectral density)
# zonele colorate corespund componentelor relevante semnalului
# rosu -> audio mono
# portocaliu -> audio stereo
plt.figure()
plt.psd(x3, NFFT=2048, Fs=Fs_new, color='blue')
plt.title('Semnal FM decimat, apoi demodulat')
plt.axvspan(0,             15000,         color='red',    alpha=0.2)
plt.axvspan(19000-500,     19000+500,     color='green',  alpha=0.2)
plt.axvspan(19000*2-15000, 19000*2+15000, color='orange', alpha=0.2)
plt.axvspan(19000*3-1500,  19000*3+1500,  color='blue',   alpha=0.2)
plt.show(block=False)
# sau salvati imaginea ca png, daca nu va merge plt.show
# plt.savefig('fig4.png')

# Calculati si afisati FFT pentru semnalului demodulat
# TODO 6:
# 1. Calculati FFT din x3 prin functia fft din scipy.fftpack
# 2. afisati valoarea absoluta a spectrului (np.abs) intre 0 si fs/2
# Nota: puteti folosi np.linspace pentru a crea un vector de frecvente intre 0 si fs, util la plot
y3 = 0 # calculati aici spectrul lui x3 folosind scipy.fftpack.fft
y3 = fft(x3)
nf3 = len(y3)
nf3_2 = int(nf3/2)
xx3 = np.linspace(0, Fs, nf3)
fig, ax = plt.subplots()
fig.suptitle('Spectrul (FFT) pentru semnalul demodulat')
ax.plot(xx3[0:nf3_2], np.abs(y3[0:nf3_2]))
ax.set_xlabel('Frecvența')
ax.grid()
plt.show(block=False)
# sau salvati imaginea ca png, daca nu va merge plt.show
# plt.savefig('fig5.png')

print('Apasati o tastă pentru a continua...\n')
input()

# extragerea semnalului audio (mono) prin filtrare si decimare
# TODO 7:
# 1. filtrati semnalul precedent (x3) pentru a retine doar componentele intre 0 si 15 kHz (corespunde semnalului mono)
# 2. Calculati si afisati FFT pentru semnalul filtrat si comparati cu FFT-ul semnalului precedent (nefiltrat)
# 3. decimati semnalul astfel incat sa aveti frecventa finala de esantionare f_audio = 44100 Hz pentru a-l putea asculta
ntaps2 = 32
fmaxa = 15000
fcut2 = 0 # calculati frecventa de cutoff pentru semnalele intre 0 si 15 kHz (tot in raport cu fs/2, atentie acum Fs_new)
b2 = signal.firwin(ntaps2, fcut2, window='hamming')
x3f = signal.lfilter(b2, 1, x3)

yff = fft(x3f)
nff = len(yff)
nff2 = int(nff/2)
xff = np.linspace(0, Fs_new, nff)
plt.figure()
plt.plot(xff[0:nff2], 2.0/nff * np.abs(yff[0:nff2]))
plt.title('spectrul (fft) semnalului fm demodulat dupa filtrare pe canalul mono')
plt.grid()
plt.show(block=False)
# sau salvati imaginea ca png, daca nu va merge plt.show
# plt.savefig('fig6.png')

f_audio = 44100
dec_audio = 0 # calculati factorul de decimare pentru a face subsampling de la Fs_new la f_audio (tot un întreg)
dec_audio = int(Fs_new / f_audio)
Fs_audio = Fs_new / dec_audio
xa = signal.decimate(x3f, dec_audio)

input("Press Enter to play audio from signal...")

# Redarea semnalului si salvarea lui ca wav 
xs = np.int16(xa/np.max(np.abs(xa)) * 32767)
sd.play(xs, f_audio) # daca nu va merge, comentati linia
write('x1.wav', f_audio, xs) # si ascultati intr-un media player fisierul acesta (ar trebui să se înțeleagă ceva)

# TODO Bonus: încercați să îmbunătățiți calitatea semnalului!

print("Ați terminat! Felicitări!\n")
input("Press Enter to continue...")