import numpy as np
import matplotlib.pyplot as plt
import scipy.fftpack
import cv2
import face_tracking
import EVM


def display_fft(video_name, fps):
    video = cv2.VideoCapture(video_name)

    # Magnitudes of the RGB signal
    R = []
    G = []
    B = []

    # Timescale and frame count
    T = 1/fps
    N = 0


    while video.isOpened():
        ret, image = video.read()
        if ret == True:
            img_h, img_w, img_c = image.shape

            R_sum = np.sum(image[:,:,0])
            G_sum = np.sum(image[:,:,1])
            B_sum = np.sum(image[:,:,2])
            
            R.append(R_sum/(img_w*img_h))
            G.append(G_sum/(img_w*img_h))
            B.append(B_sum/(img_w*img_h))
            N += 1
        else:
            break

    # Apply Fourier transform and display
    R_f = scipy.fftpack.fft(R)
    G_f = scipy.fftpack.fft(G)
    B_f = scipy.fftpack.fft(B)
    R_f[0:5] = 0
    G_f[0:5] = 0
    B_f[0:5] = 0

    # Plot the calculated heart rate
    f = np.linspace(0.0, 1.0*60/(2.0*T), N//2)
    figR, aR = plt.subplots(2)
    figG, aG = plt.subplots(2)
    figB, aB = plt.subplots(2)

    # Name the plots
    figR.suptitle('Red channel signal in time and frequency domain')
    figG.suptitle('Green channel signal in time and frequency domain')
    figB.suptitle('Blue channel signal in time and frequency domain')

    # Name axes
    aR.flat[0].set(xlabel='Frekvence [BPM]', ylabel='Amplituda')
    aR.flat[1].set(xlabel='Čas [s]', ylabel='Amplituda')
    aG.flat[0].set(xlabel='Frekvence [BPM]', ylabel='Amplituda')
    aG.flat[1].set(xlabel='Čas [s]', ylabel='Amplituda')
    aB.flat[0].set(xlabel='Frekvence [BPM]', ylabel='Amplituda')
    aB.flat[1].set(xlabel='Čas [s]', ylabel='Amplituda')

    # Plot the frequency domain
    aR[0].plot(f, 2.0/N * np.abs(np.real(R_f[:N//2])))
    aG[0].plot(f, 2.0/N * np.abs(np.real(G_f[:N//2])))
    aB[0].plot(f, 2.0/N * np.abs(np.real(B_f[:N//2])))

    # Plot the time signal of red channel
    t = np.linspace(0.0, N*T, N)
    aR[1].plot(t, R)
    aG[1].plot(t, G)
    aB[1].plot(t, B)

    # Find the peak
    R_location = np.argmax(np.abs(np.real(R_f[:N//2])))
    G_location = np.argmax(np.abs(np.real(G_f[:N//2])))
    B_location = np.argmax(np.abs(np.real(B_f[:N//2])))
    RHR = f[R_location]
    GHR = f[G_location]
    BHR = f[B_location]

    print('Measured heart rate: R = ' + str(RHR))
    print('Measured heart rate: G = ' + str(GHR))
    print('Measured heart rate: B = ' + str(BHR))
    plt.show()

    return 0

video_name = '14-36-47_moving_30fps_lowres'
cutout_size = 16
amplification = 80
fps = 30

face_tracking.detect_face(video_name+'.mp4', video_name+'_'+str(cutout_size)+'_out.mp4', fps, cutout_size = cutout_size)
EVM.magnify_color(video_name+'_'+str(cutout_size)+'_out.mp4', video_name+'_'+str(cutout_size)+'_'+str(amplification)+'_out.avi', 75/60, 120/60, 3, amplification)
display_fft(video_name+'_'+str(cutout_size)+'_'+str(amplification)+'_out.avi', fps)