"""
DESCRIPTION
Python 2.7 script to create a FFT of a square wave with differing duty cycle
ultimately displaying the data in an intensity plot.

AUTHORS
James Barrington, Stephen James
"""


from scipy import signal
import matplotlib.pyplot as plt

import numpy as np

period = 50e-6 # Vary the frequency resolution 
length = 40e-3

per = np.linspace(-length/2, length/2, 100000)
freq = np.fft.fftfreq(np.size(per), length/np.size(per))

spec = []
per = np.linspace(-length/2, length/2, length/period)
duty = np.linspace(0, 1, 10000) # Vary the duty cycle resolution

for i in duty:
    sq = 1+(signal.square(2*np.pi*per*length/(period), i))
    ft = np.fft.fft(sq)
    spec.append(np.fft.fftshift(ft))

spec1 = np.array(spec)

# Production of the intensity plot
plt.figure('Intensity plot 2', facecolor='white', edgecolor='white')
plt.imshow(np.abs(spec1),
           cmap='gray_r', # white and black 0% and 100% power
           vmin=np.min(spec1), # If running linux, comment out
           vmax=np.max(spec1),
           # vmin=np.min(np.float128(spec1)), # If running windows, comment out
           # vmax=np.max(np.float128(spec1)),
           interpolation='nearest',
           aspect='auto',
           origin='lower',
           extent =[np.min(freq/1000), np.max(freq/1000), np.min(duty)*100, np.max(duty)*100])
plt.xlabel("Frequency (a.u.)", fontsize=24)
plt.ylabel("Duty Cycle (%)", fontsize=24)
plt.axhline(25, color='k', linestyle=':')
plt.axhline(75, color='k', linestyle=':')
plt.axhline(50, color='k', linestyle=':')
plt.annotate('1$^{st}$', xy=(88.404, 80), xytext=(43.022, 86), arrowprops=dict(facecolor='black', shrink=0.05), fontsize=18)
plt.annotate('2$^{nd}$', xy=(188.092, 80), xytext=(136.002, 86), arrowprops=dict(facecolor='black', shrink=0.05), fontsize=18)
plt.annotate('3$^{rd}$', xy=(294.425, 80), xytext=(242.335, 86), arrowprops=dict(facecolor='black', shrink=0.05), fontsize=18)
plt.tick_params(labelsize=18)
plt.subplots_adjust(left=0.2, right=0.7, bottom=0.1, top=0.9)


plt.show() # x-axis -350 to 350