from matplotlib import pyplot as plt
%matplotlib inline
import math
import numpy as np
import cv2

def normalize_img(s):
    start = 0
    end = 255
    width = end - start
    res = (s - s.min())/(s.max() - s.min()) * width + start
    return res

# Vertical
A = 0.5
f = 7
t = np.linspace(0,1,50) # for loop 1-100 and / 100 each time
phi = np.pi/4
x = normalize_img(A*np.cos(2*np.pi*f*t + phi)) # 1d sine
sinusoid = np.array([x]*len(t)).astype(np.float32) # cv::Mat 100 * 100, 1-row copy of x, 2-row

dft = cv2.dft(sinusoid, flags = cv2.DFT_COMPLEX_OUTPUT) # Computed the 2-d discrete Fourier Transform
dft_shift = np.fft.fftshift(dft) # Shift the zero-frequency component to the center of the spectrum.

magnitude_spectrum = 20*np.log(cv2.magnitude(dft_shift[:,:,0],dft_shift[:,:,1])+1) # compute magnitude spectrum

fig, (ax1, ax2) = plt.subplots(figsize=(9, 5), nrows=1, ncols=2)

ax1.imshow(sinusoid, cmap = 'gray')
ax1.set_title('Sinusoid 1')
ax1.set_xticks([])
ax1.set_yticks([])

ax2.imshow(magnitude_spectrum, cmap = 'gray')
ax2.set_title('Frequency Spectrum 1')
ax2.set_xticks([])
ax2.set_yticks([])
fig.savefig('vertical.jpg')

# Horizontal
sinusoid2 = sinusoid.transpose()

dft = cv2.dft(sinusoid2, flags = cv2.DFT_COMPLEX_OUTPUT) # Computed the 2-d discrete Fourier Transform
dft_shift = np.fft.fftshift(dft) # Shift the zero-frequency component to the center of the spectrum.

magnitude_spectrum = 20*np.log(cv2.magnitude(dft_shift[:,:,0],dft_shift[:,:,1])+1) # compute magnitude spectrum


fig, (ax1, ax2) = plt.subplots(figsize=(9, 5), nrows=1, ncols=2)
ax1.imshow(sinusoid2, cmap = 'gray')
ax1.set_title('Sinusoid 2')
ax1.set_xticks([])
ax1.set_yticks([])

ax2.imshow(magnitude_spectrum, cmap = 'gray')
ax2.set_title('Frequency Spectrum 2')
ax2.set_xticks([])
ax2.set_yticks([])

fig.savefig('horizontal.jpg')

# Linear sum
combine_signal = cv2.add(sinusoid, sinusoid2)

dft = cv2.dft(combine_signal, flags = cv2.DFT_COMPLEX_OUTPUT) # Computed the 2-d discrete Fourier Transform
dft_shift = np.fft.fftshift(dft) # Shift the zero-frequency component to the center of the spectrum.

magnitude_spectrum = 20*np.log(cv2.magnitude(dft_shift[:,:,0],dft_shift[:,:,1])+1) # compute magnitude spectrum

fig, (ax1, ax2) = plt.subplots(figsize=(9, 5), nrows=1, ncols=2)

ax1.imshow(combine_signal, cmap = 'gray')
ax1.set_title('Sinusoid 1 + Sinusoid 2')
ax1.set_xticks([])
ax1.set_yticks([])

ax2.imshow(magnitude_spectrum, cmap = 'gray')
ax2.set_title('Frequency Spectrum 1 + Frequency Spectrum 2')
ax2.set_xticks([])
ax2.set_yticks([])

fig.savefig('Linear_sum.jpg')