import numpy as np
from scipy import signal
import matplotlib.pyplot as plt

fs=2*np.pi*120 / 512 # sampling frequency (rad/s / samples)
ts = 1 / 120 # sampling period

""" Filter order """
N, Wn = signal.buttord(wp=0.1, 
                       ws=0.3, 
                       gpass=3, 
                       gstop=40, 
                       analog=False, 
                       fs=fs)

""" Filter """
b, a = signal.butter(N, Wn, 'low',                           
                     analog=False)

""" Frequencies """
w, h = signal.freqz(b, a, 
                    worN=512,
                    whole=True, 
                    fs=fs)
""" Rayleigh series"""

sigma = 0.12 # variance
leng = 10000 # number of samples

""" I and Q components """
ii = np.random.randn(leng, 1) * sigma
qq = np.random.randn(leng, 1) * sigma

rayleigh = ii + 1j * qq
rayleigh = [i[0] for i in rayleigh]

""" Filtering """
ryfaux = signal.filtfilt(b, a, list(rayleigh))

ry_mod = abs(ryfaux)
t_axis = np.array([i for i in range(leng)])*ts

""" Plotting """

plt.subplots(figsize=(12, 7))

ax1 = plt.subplot(212)
ax1.margins(2, 2)
ax1.plot(t_axis, 20*np.log10(ry_mod))
ax1.axis([10, 12, -50, -10])
ax1.set_xlabel('Time (s)')
ax1.set_ylabel('Relative signal level (dB)')
ax1.set_title('Rayleigh fading series (filtered)')
ax1.grid(which='both', axis='both')

ax2 = plt.subplot(221)
ax2.plot(w * 512 / (2*np.pi), 20*np.log10(abs(h)))
ax2.axis([0, 31, -71, 0])
ax2.margins(.05) 
ax2.set_title('Doppler filter responce (Butterworth lowpass digital filter)')
ax2.set_xlabel('Frequence (Hz)')
ax2.set_ylabel('Magnitude of frequency response (dB)')
ax2.grid(which='both', axis='both')

ax3 = plt.subplot(222)
ax3.margins(x=0, y=-0.25) 
ax3.plot(t_axis, 20*np.log10(abs(np.array(rayleigh))))
ax3.axis([10, 12, -50, -10])
ax3.set_xlabel('Time (s)')
ax3.set_ylabel('Relative signal level (dB)')
ax3.set_title('Rayleigh fading series')
ax3.grid(which='both', axis='both')


plt.savefig("rayleigh_fading.png")