# -*- coding: utf-8 -*-
# 2009-12-06 Moritz

import matplotlib
matplotlib.rc('font', size=8) # adjust font size of plot
import matplotlib.pyplot as plt
plt.close('all')
import numpy as np
from obspy.signal import lowpass, lowpassZPHSH, cosTaper
PI = np.pi

# Number of points
#npts=64
nsec=4.0
df=20.0

# Determine Nyquist and set generator frequency
fmax=df/2.0 # nyquist
fg1 = 8.0 # generator frequency I
fg2 = 4.0 # generator frequency II

# Generate f0 Hz sine wave
#time = np.arange(0,npts)/df # in seconds
time = np.linspace(0,nsec,nsec*df) # in seconds
y  = np.sin(2*PI*fg1 * time)#+ PI/5)
y += np.sin(2*PI*fg2 * time)#+ PI/5)

# Downsample to 10Hz by taking every second element
y_2 = y[::2]

# Downsample after Lowpassing the signal
y_l = lowpass(y, 5.0, df = df, corners=4)
y_new = y_l[::2]

# Looking at result in frequency domain
y_f = np.fft.rfft(y)
y_f2 = np.fft.rfft(y_2)
y_fnew = np.fft.rfft(y_l)
freq = np.linspace(0, fmax, len(y_f))


#
# Plot the whole filtering process
#
plt.figure(figsize=(14,10))
plt.subplot(211)
plt.plot(time,y,'k',label="Original data",lw=1)
plt.plot(time[::2], y_2,'r--',label="Every second element",lw=3)
plt.plot(time[::2], y_new,'g',label="Lowpassed below nyquist",lw=3)
plt.legend()
plt.title('Time Domain')
plt.xlabel('Time [s]')
plt.ylabel('Amplitude')

plt.subplot(212)
plt.plot(freq, abs(y_f), 'k', label="Original frequencies", lw=1)
plt.plot(freq, abs(y_fnew), 'g--', label="Lowpassed below nyquist", lw=3)
plt.plot(freq[:len(y_f2)], abs(y_f2), 'r--', label="Just taking half", lw=3)
plt.legend()
plt.title('Frequency Domain')
plt.xlabel('Frequency [Hz]')
plt.ylabel('Arbitrary Amplitude')

plt.show()