#
# Code for plotting the effect of running DFT/FFT on finite length data
#
# Prerequisites: Python+numpy+matplotlib
#
#
# How to run it:
# >> python 3470_kap8_dft_final_length_data.py
#


# Import all the functions we need
from numpy import *
from pylab import *

# Define a function that computes dB
def db(val):
   return 20*log10(val)


N     = 48      # The signal periodicity
Nc    = 25*N    # The total number of periods that we care to plot
ns    = 5       # Number of samples difference from N when we take DFT

N1    = N-ns    # For the DFT of one period. We "miss" by 5 samples
N2    = 11*N-ns # For the DFT of several periods. Still miss by 5 samples

F0    = 3.0/N   # Signal frequency (normalized)

# Time axes, and corresponding signals
n     = arange(N)                   # Creates n = 0,1,...,N-1
nc    = arange(-(Nc-1)/2,(Nc-1)/2)  # Creates negative/positive indexes for the full signal
n1    = arange(N1)
n2    = arange(N2)

# Create time 
x     = 0.5*sin(2*pi*F0*(n))        # Finite length signal of length N (one period)
xc    = 0.5*sin(2*pi*F0*(nc))       # Underlying signal (lots of periods)
x1    = 0.5*sin(2*pi*F0*(n1))       # Finite length signal of length N-5 (less than one period)
x2    = 0.5*sin(2*pi*F0*(n2))       # Finite length signal of length 11N-5 (still no multiple of
                                    # the period, but now 11 periods are considered.
xp    = 0.5*sin(2*pi*F0*(n2+ns))

# Frequency axes for the DFT
F     = linspace(-0.5,0.5,N)-0.5/N
F1    = linspace(-0.5,0.5,N1)
F2    = linspace(-0.5,0.5,N2)

########################
# PLOT SIGNALS IN TIME #
########################

fn = figure()
ax = fn.add_subplot(311)

# Plot the underlying signal, give it a grey appearance, and add legend (description text box)
markerline, stemlines, baseline = ax.stem(nc,xc, linefmt='0.8', markerfmt='.', basefmt='')
setp(markerline, 'markerfacecolor', '0.8', 'markeredgecolor', '0.8')
legend((r'$x_p[n] = 0.5\sin(2\pi F_0 n)$',), loc='upper center')

# Plot the red box to indicate signal extent, then plot the signal
ax.plot([0]+list(n)+[N-1],[0.0]+list(ones((N,)))+[0.0],'r-')
h = ax.stem(n,x, linefmt='b-', markerfmt='b.', basefmt='')

# Turn on a grid, specify extent of the x- and y-axes, and add descriptive texts
ax.grid(True)
ylim([-1.5,1.5])
xlim([-N/2,N+N/2])
ylabel('Amplitude')
legend(h,(r'$N = 3\frac{F_s}{F_0}$',), loc='lower center')

# That was the first signal. Now we do the same for the next two cases:
ax = fn.add_subplot(312)
markerline, stemlines, baseline = ax.stem(nc,xc, linefmt='b-', markerfmt='b.', basefmt='')
setp(markerline, 'markerfacecolor', '0.8', 'markeredgecolor', '0.8')
setp(stemlines,  'color', '0.8')
ax.plot([0]+list(n1)+[N1-1], [0.0]+list(ones((N1,)))+[0.0],'r-')
h = ax.stem(n1,x1, linefmt='b-', markerfmt='b.', basefmt='')
ax.grid(True)
ylim([-1.5,1.5])
xlim([-N/2,N+N/2])
ylabel('Amplitude')
legend(h,(r'$N = 3\frac{F_s}{F_0}$-5',), loc='lower center')

ax = fn.add_subplot(313)
markerline, stemlines, baseline = ax.stem(nc,xc, linefmt='b-', markerfmt='b.', basefmt='')
setp(markerline, 'markerfacecolor', '0.8', 'markeredgecolor', '0.8')
setp(stemlines,  'color', '0.8')
ax.plot([0]+list(n2)+[N2-1], [0.0]+list(ones((N2,)))+[0.0],'r-')
h = ax.stem(n2,x2, linefmt='b-', markerfmt='b.', basefmt='')
ax.grid(True)
ylim([-1.5,1.5])
xlim([-N/2,(11+0.5)*N])
ylabel('Amplitude')
legend(h,(r'$N = 11\frac{F_s}{F_0}$-5',), loc='lower center')

# And finally save this figure to a pdf-file
savefig('final_length_data.pdf', bbox_in2hes='tight', pad_in2hes=0, format='pdf')


############################
# PLOT FREQUENCY RESPONSES #
############################

fn = figure()
ax = fn.add_subplot(311)

# Perform N-point DFT/FFT, then shift it to be symmetric around F=0, and plot absolute value:
ax.stem(F, abs(fftshift(fft(x,N)/(N))),linefmt='b-', markerfmt='.', basefmt='')

# Turn on the grid and specify x-axis extent and add legend
grid(True)
xlim([-0.5,0.5])
ax.legend((r'$N = 3\frac{F_s}{F_0}$',), loc='upper right')

# Specify and tag the values on the x- and y-axes manually (to mark the interesting ones)
xticks( [-0.5, -F0, 0, F0, 0.5],
        (r'-$\frac{F_s}{2}$', r'-$F_0$', r'0', r'$F_0$', r'$\frac{F_s}{2}$'))
yticks( [0, 0.25/2, 0.25, 0.25*1.1],
        (r'0', r'$\frac{1}{4}$', r'$\frac{1}{2}$', r''))
ylabel('Magnitude')


# Then repeat for the next to signals
ax = fn.add_subplot(312)
ax.stem(F1, abs(fftshift(fft(x1,N1)/(N1))),linefmt='b-', markerfmt='.', basefmt='')
grid(True)
xticks( [-0.5, -F0, 0, F0, 0.5],
(r'-$\frac{F_s}{2}$', r'-$F_0$', r'0', r'$F_0$', r'$\frac{F_s}{2}$'))
xlim([-0.5,0.5])
yticks( [0, 0.25/2, 0.25, 0.25*1.1],
(r'0', r'$\frac{1}{4}$', r'$\frac{1}{2}$', r''))
ylabel('Magnitude')
ax.legend((r'$N = 3\frac{F_s}{F_0}$-5',), loc='upper right')

ax = fn.add_subplot(313)
ax.stem(F2, abs(fftshift(fft(x2,N2)/(N2))),linefmt='b-', markerfmt='.', basefmt='')
grid(True)
xticks( [-0.5, -F0, 0, F0, 0.5],
(r'-$\frac{F_s}{2}$', r'-$F_0$', r'0', r'$F_0$', r'$\frac{F_s}{2}$'))
xlim([-0.5,0.5])
yticks( [0, 0.25/2, 0.25, 0.25*1.1],
(r'0', r'$\frac{1}{4}$', r'$\frac{1}{2}$', r''))
ylabel('Magnitude')
ax.legend((r'$N = 11\frac{F_s}{F_0}$-5',), loc='upper right')

# And finally save the figure to a pdf-file
savefig('final_length_spectrum.pdf', bbox_in2hes='tight', pad_in2hes=0, format='pdf')