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

import matplotlib
matplotlib.rc('figure.subplot', hspace=.35, wspace=.35) #adjust subplot layout
matplotlib.rc('font', size=8) # adjust font size of plot
import matplotlib.pyplot as plt
plt.close('all')
import numpy as np

def filter_bc(x, dt, f0):
    """
    The simple-most low-pass/boxcar filter

    @param x: Time series
    @param dt: Sampling intervall
    @param f0: Frequency of lowpass
    @return: (freq, H, x_filt) with x_f frequency range, 
             H consturcted filter, x_filt filtered time series
    """
    npts = len(x)
    
    # Initialize
    nyq = 1.0/(2*dt) # Nyquist frequency
    freq = np.linspace(0, nyq, npts//2+1) #// == integer devision

    # Construct the Boxcar
    # generate array with 0 everywhere but 1 where freq < f0
    H = np.where(freq<f0, 1.0, 0.0)
    H[0] = 1 # set offset for sure to one

    # Filter in the Fourier Domain
    x_f = np.fft.rfft(x)
    return freq, H, np.fft.irfft(x_f*H)





######################################################################
#                Start of Main Program                               #
######################################################################


# Number of points (here must be even for fft) and sampling_rate
npts=1024
#npts=64
dt=0.05

# Determining nyquist frequency and asking for cut off frequency f0
fmax=1.0/(2*dt) # nyquist
f0 = float(raw_input('Give cut-off below Nyquist: fmax = %4.1f Hz ' % fmax))
print 'f0 =', f0

# Spike at npts/2
y = np.zeros(npts,dtype='float')
y[npts/2] = 1

# Uncomment from random points
#y = np.random.rand(npts) - .5 # uniform random numbers, zero mean

# Filter with filter_bc
freq, H, y_filt = filter_bc(y,dt,f0)

# Just for the plot, frequency domain representaion of y
y_f = np.fft.rfft(y)


#
# Plot the whole filtering process
#
time = np.arange(0,npts)*dt
#
plt.subplot(231)
plt.plot(time, y)
plt.title('Original data')
plt.xlabel('Time [s]')
plt.ylabel('Amplitude')
plt.ylim([plt.ylim()[0]-0.2, plt.ylim()[1]+0.2]) 

plt.subplot(232)
plt.plot(freq, abs(y_f))
plt.title('Amplitude spectrum')
plt.xlabel('Frequency [Hz]')
plt.ylabel('Arbitrary Amplitude')
plt.ylim([plt.ylim()[0]-0.2, plt.ylim()[1]+0.2]) 

plt.subplot(233), 
plt.plot(freq, np.angle(y_f))
plt.title('Phase spectrum')
plt.xlabel('Frequency [Hz]')
plt.ylabel('Phase shift [rad]')

plt.subplot(234)
plt.plot(freq, H)
plt.title('The filter')
plt.xlabel('Frequency [Hz]')
plt.ylabel('Filter amplitude')
plt.ylim([plt.ylim()[0]-0.2, plt.ylim()[1]+0.2]) 

plt.subplot(235)
plt.plot(freq, abs(y_f*H))
plt.title('The filtered spectrum')
plt.xlabel('Time [s]')
plt.ylabel('Frequency [Hz]')
plt.ylim([plt.ylim()[0]-0.2, plt.ylim()[1]+0.2]) 

plt.subplot(236)
plt.plot(time, y_filt)
plt.title('The filtered signal')
plt.xlabel('Time [s]')
plt.ylabel('Amplitude')

plt.show()