# coding: utf-8

# # Frequency-domain audio experiments

# In[5]:

import numpy as np
import numpy.linalg as la
import matplotlib.pyplot as pt


# In[6]:

wave = np.load("cs357.npy")


# In[7]:

pt.plot(wave)


# In[9]:

from simple_audio import play
play(wave)


# Now build the Fourier Vandermonde matrix. This is also called the *Discrete Fourier Transform* (DFT) matrix.

# In[10]:

n = 255
assert n % 2 == 1

cos_k = np.arange(0, n//2 + 1, dtype=np.float64)
sin_k = np.arange(1, n//2 + 1, dtype=np.float64)

x = np.linspace(0, 2*np.pi, n, endpoint=False)

DFT = np.zeros((n,n))
DFT[:, ::2] = np.cos(cos_k*x[:, np.newaxis])
DFT[:, 1::2] = np.sin(sin_k*x[:, np.newaxis])

IDFT = la.inv(DFT)


# And transform the signal:

# In[11]:


n_chunks = len(wave)//n

chunked = wave[:n_chunks*n].reshape(-1, n)
coeffs = np.dot(IDFT, chunked.T).T


# In[17]:

pt.figure(figsize=(10,8))
pt.imshow(np.log10(1e-5+np.abs(coeffs)).T)
pt.colorbar()


# * What do you observe around the "s" sounds?

# In[30]:

modified = coeffs.copy()

#modified[:,5:-5] *= 1e-2

#modified *= 1e-2
#modified[:,5:-5] /= 1e-2

nshift = 3
modified.fill(0)
modified[:, nshift:] = coeffs[:, 0:-nshift]

pt.figure(figsize=(8,5))
pt.imshow(np.log10(1e-5+abs(modified)).T)
pt.colorbar()


# In[31]:

result = np.dot(DFT, modified.T).T
result = result.reshape(-1)


# In[32]:

new_wave = np.zeros_like(wave)
new_wave[:len(result)] = result
new_wave = new_wave/np.max(np.abs(new_wave))


# In[33]:

play(new_wave)


# In[30]:

play(wave)


# In[22]:

pt.plot(new_wave)