import numpy as np
import matplotlib.pyplot as plt

N = 32
omega = np.exp(-1j*2*np.pi/N)
W = np.ones((N,N), dtype=complex)
for i in range(1,N):
  for j in range(i,N):
    W[i,j] = omega**(i*j)
    W[j,i] = W[i,j]

diff = np.linalg.inv(W) - np.conj(W)/N
np.linalg.norm(diff)

plt.figure(figsize=(12,3), dpi=100)
plt.subplot(131)
plt.imshow(W.real, cmap='gray')
plt.title('Re(W)')
plt.colorbar()
plt.subplot(132);
plt.imshow(W.imag, cmap='gray')
plt.title('Im(W)')
plt.colorbar()
plt.subplot(133);
plt.imshow(np.angle(W)*180/np.pi, cmap='gray')
plt.title('Phase(W)')
plt.colorbar()
plt.show()

plt.figure(figsize=(12,12), dpi=100)
for i in range(N):
  plt.subplot(N,1,i+1)
  if i==0:
    plt.title(f'{N} point DFT basis vectors \
    (Solid: real part, Dashed: imaginary part)')
  plt.plot(W[i,:].real, label='Re')
  plt.plot(W[i,:].imag, ':', label='Im')
  plt.yticks([])
  if i!=N-1:
    plt.xticks([])
  plt.ylim(-1.2,1.2)
  plt.ylabel(f'{i+1}')
plt.show()



# Create a sound
fs = 44100
dt = 1/fs
t = np.arange(0,2,dt)
xl = 2*np.sin(2*np.pi*220*t)
xr = np.sin(2*np.pi*225*t)
#t = t[:fs]
#xl=xl[:fs]
#xr=xr[:fs]
x = xl + xr

plt.figure(figsize=(12,12), dpi=100)
plt.subplot(311)
plt.plot(t,xl, label=r'$x_l$')
plt.plot(t,xr, label=r'$x_r$')
plt.xlim(0,0.05)
plt.xlabel(r'$t$ (sec)')
plt.ylabel(r'$x_l$ and $x_r$')
plt.title('Original sound signal')
plt.grid()
plt.legend()
plt.subplot(312)
plt.plot(t,x)
plt.xlabel(r'$t$ (sec)')
plt.ylabel(r'$x=x_l + x_r$')
plt.grid()
plt.subplot(313)
plt.plot(t,x)
plt.xlim(0,0.25)
plt.xlabel(r'$t$ (sec)')
plt.ylabel(r'$x=x_l + x_r$')
plt.grid()
plt.show()

from IPython.display import Audio

# play in mono
Audio(x, rate=fs, autoplay=True)

# play in stereo
Audio([xl, xr], rate=fs, autoplay=True)

N, 1/dt

N = len(x)

df = fs/N                 # df = fs/N
f = np.arange(0,N)*df     # f = [0, df, 2*df,..., (N-1)*df]

X = np.fft.fft(x)*dt

plt.figure(figsize=(12,8), dpi=100)
plt.subplot(211)
plt.plot(f[0:int(N/2+1)],np.abs(X[0:int(N/2+1)]))
plt.xlabel('freq (hz)')
plt.ylabel(r'|X|')
plt.title('Fourier transform of original signal')
plt.grid()
plt.subplot(212)
plt.plot(f[0:int(N/2+1)],np.abs(X[0:int(N/2+1)]))
plt.xlabel('frequency (hz)')
plt.ylabel(r'|X|')
plt.grid()
plt.xlim(210,230)
plt.show()

np.random.seed(1)
x_c = x + 10*np.random.randn(len(x))

plt.figure(figsize=(12,8), dpi=100)
plt.subplot(211)
plt.plot(t,x_c)
plt.xlabel(r'$t$ (sec)')
plt.ylabel(r'$x_c$')
plt.title('Corrupted sound signal')
plt.grid()
plt.subplot(212)
plt.plot(t,x_c, label='Corrupted signal')
plt.plot(t,x, label='Original signal')
plt.xlim(0,0.25)
plt.xlabel(r'$t$ (sec)')
plt.ylabel(r'$x_c$')
plt.grid()
plt.legend()
plt.show()

# play corrupted signal
Audio(x_c, rate=fs, autoplay=True)

df = fs/N                 # df = fs/N
f = np.arange(0,N)*df     # f = [0, df, 2*df,..., (N-1)*df]

X_c = np.fft.fft(x_c)*dt

plt.figure(figsize=(12,8), dpi=100)
plt.subplot(211)
plt.plot(f[0:int(N/2+1)],np.abs(X_c[0:int(N/2+1)]))
plt.xlabel('freq (hz)')
plt.ylabel(r'$|X_c|$')
plt.title('Fourier transform of corrupted signal')
plt.grid()
plt.subplot(212)
plt.plot(f[0:int(N/2+1)],np.abs(X_c[0:int(N/2+1)]))
plt.xlabel('frequency (hz)')
plt.ylabel(r'$|X_c|$')
plt.grid()
plt.xlim(210,230)
plt.show()

X_recon = np.copy(X_c)
X_recon[np.abs(X_recon)<0.1*np.max(np.abs(X_recon))] = 0

plt.figure(figsize=(12,8), dpi=100)
plt.subplot(211)
plt.plot(f[0:int(N/2+1)],np.abs(X_c[0:int(N/2+1)]), label='corrupted')
plt.plot(f[0:int(N/2+1)],np.abs(X_recon[0:int(N/2+1)]), label='reconstructed')
plt.xlabel('freq (hz)')
plt.ylabel(r'$|X|$')
plt.legend()
plt.title('Fourier transform of corrupted signal');
plt.grid()
plt.subplot(212)
plt.plot(f[0:int(N/2+1)],np.abs(X_c[0:int(N/2+1)]), label='corrupted')
plt.plot(f[0:int(N/2+1)],np.abs(X_recon[0:int(N/2+1)]), label='reconstructed')
plt.xlabel('frequency (hz)')
plt.ylabel(r'$|X|$')
plt.legend()
plt.grid()
plt.xlim(210,230)
plt.show()

x_recon = np.fft.ifft(X_recon).real*fs

plt.figure(figsize=(12,8), dpi=100)
plt.subplot(211)
plt.plot(t,x_recon)
plt.xlabel(r'$t$ (sec)')
plt.ylabel(r'$x_{reconstructed}$')
plt.title('Reconstructed sound signal')
plt.subplot(212)
plt.plot(t,x_recon, label='Reconstructed signal')
plt.plot(t,x, label='Original signal')
plt.xlim(0,0.25)
plt.xlabel(r'$t$ (sec)')
plt.ylabel(r'$x_{reconstructed}$')
plt.grid()
plt.legend()
plt.show()

# play reconstructed signal
Audio(x_recon, rate=fs, autoplay=True)