import numpy as np
import matplotlib.pyplot as plt
import scipy as sp

f_s = 10
T_s = 1/f_s
t = np.arange(0,1+T_s,T_s)

f = np.linspace(0,f_s,len(t))
omega = 2*np.pi*f

def u(t):
    return np.heaviside(t, 1)

w = u(t) - u(t-0.6) + u(t-0.7) - u(t-1)

W_c = 1j*(np.exp(-1j*omega*0.6) + np.exp(-1j*omega) - np.exp(-1j*omega*0.7) - 1)/omega
W_d = sp.fft.fft(w)

plt.stem(f, W_d)
plt.xlabel("Frequency (Hz)")
plt.show()
plt.bar(t, sp.fft.ifft(W_d), T_s, align='edge')
plt.xlabel("Time (s)")
plt.show()