import cv2
import numpy as np
import scipy.signal as signal
import scipy.fftpack as fftpack

def im2double(im):
    info = np.iinfo(im.dtype) # Get the data type of the input image
    return im.astype(np.float) / info.max # Divide all values by the largest possible value in the datatype

def im2uint8(im):
    return ((im*255).astype(np.int))

#convert RBG to YIQ
def rgb2ntsc(src):
    [rows,cols]=src.shape[:2]
    dst=np.zeros((rows,cols,3),dtype=np.float64)
    T = np.array([[0.114, 0.587, 0.298], [-0.321, -0.275, 0.596], [0.311, -0.528, 0.212]])#T = np.array([[0.299, 0.587, 0.114], [0.596, -0.274, -0.322], [0.211, -0.523, 0.312]])
    for i in range(rows):
        for j in range(cols):
            dst[i, j]=np.dot(T,src[i,j])
    return dst

#convert YIQ to RBG
def ntsc2rbg(src):
    [rows, cols] = src.shape[:2]
    dst=np.zeros((rows,cols,3),dtype=np.float64)
    T = np.array([[1, -1.108, 1.705], [1, -0.272, -0.647], [1, 0.956, 0.620]])
    for i in range(rows):
        for j in range(cols):
            dst[i, j]=np.dot(T,src[i,j])
    return dst

# Build a Gaussian pyramid of the image
def build_gaussian_pyramid(img, level = 3):
    a = img.copy()
    pyramid = [a]

    for i in range(level):
        a = cv2.pyrDown(a)
        pyramid.append(a)

    return pyramid

# Build a Laplacian pyramid for motion magnification
def build_laplacian_pyramid(img, levels = 3):
    gaussianPyramid = build_gaussian_pyramid(img, levels)
    pyramid = []

    for i in range(levels, 0, -1):
        GE = cv2.pyrUp(gaussianPyramid[i])
        L = cv2.subtract(gaussianPyramid[i-1], GE)
        pyramid.append(L)

    return pyramid

def load_video(filename):
    cap = cv2.VideoCapture(filename)
    frame_cnt = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))
    width, height = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH)), int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
    fps = int(cap.get(cv2.CAP_PROP_FPS))
    video_tensor = np.zeros((frame_cnt, height, width,3), dtype='uint8')

    x = 0
    while cap.isOpened():
        ret, frame = cap.read()
        if ret is True:
            video_tensor[x] = frame
            x += 1
        else:
            break
    
    return video_tensor, fps

# Temporal bandpass filter
# DEBUG
def temporal_filter(tensor, low_f, high_f, fps, axis = 0):
    fft = fftpack.fft(tensor, axis=axis)
    freqs = fftpack.fftfreq(tensor.shape[0], d = 1.0/fps)
    bound_low = (np.abs(freqs - low_f)).argmin()
    bound_high = (np.abs(freqs - high_f)).argmin()
    fft[:bound_low] = 0
    fft[bound_high:-bound_low] = 0
    fft[-bound_low:] = 0
    iff = fftpack.ifft(fft, axis=axis)
    return np.real(iff)

# Build a Gaussian pyramid for the video
def gaussian_video(video_tensor, levels = 3):
    for i in range(0, video_tensor.shape[0]):
        frame = im2double(video_tensor[i])#cv2.normalize(video_tensor[i].astype('float'), None, 0.0, 1.0, cv2.NORM_MINMAX)
        frame = rgb2ntsc(frame)
        pyr = build_gaussian_pyramid(frame, level = levels)
        gaussian_frame = pyr[-1]
        if i == 0:
            vid_data = np.zeros((video_tensor.shape[0], gaussian_frame.shape[0], gaussian_frame.shape[1], 3))
        vid_data[i] = gaussian_frame
    return vid_data

# Amplify the video
def amplify_video(gaussian_vid, amplification = 50):
    return (gaussian_vid*amplification)

# Collapse the pyramid and combine with original video
def reconstruct_video(amp_video, origin_video, levels = 3):
    final_video = np.zeros(origin_video.shape)

    for i in range(0, amp_video.shape[0]):
        img = amp_video[i]
        for x in range(levels):
            img = cv2.pyrUp(img)
        img = img + rgb2ntsc(im2double(origin_video[i]))
        img = ntsc2rbg(img)
        for x in range(0,img.shape[0]):
            for y in range(0,img.shape[1]):
                for z in range(0,img.shape[2]):
                    if img[x,y,z] >= 1.0:
                        img[x,y,z] = 1.0
                    elif img[x,y,z] <= 0.0:
                        img[x,y,z] = 0.0

        final_video[i] = im2uint8(img)

    return final_video

# Save video into a file
# TODO optional frame rate and video name
def save_video(video_tensor, video_name):
    fourcc = cv2.VideoWriter_fourcc('M', 'J', 'P', 'G')
    [height, width] = video_tensor[0].shape[0:2]
    writer = cv2.VideoWriter(video_name, fourcc, 30, (width, height), 1)

    for i in range(0, video_tensor.shape[0]):
        writer.write(cv2.convertScaleAbs(video_tensor[i]))
    writer.release()

# Magnify color
def magnify_color(video_name, video_out_name, low_f, high_f, levels = 3, amplification = 20):
    t,fps = load_video(video_name)
    gauss_video = gaussian_video(t, levels = levels)
    filtered_tensor = temporal_filter(gauss_video, low_f, high_f, fps)
    amplified_video = amplify_video(filtered_tensor, amplification = amplification)
    final = reconstruct_video(amplified_video, t, levels = levels)
    save_video(final, video_out_name)

# Build a Laplacian pyramid for motion magnification
def laplacian_video(video_tensor, levels = 3):
    tensor_list = []

    for i in range(0, video_tensor.shape[0]):
        frame = video_tensor[i]
        pyr = build_laplacian_pyramid(frame, levels = levels)
        if i == 0:
            for k in range(levels):
                tensor_list.append(np.zeros((video_tensor.shape[0], pyr[k].shape[0], pyr[k].shape[1], 3)))
        for n in range(levels):
            tensor_list[n][i] = pyr[n]

    return tensor_list

# Butterworth bandpass filter
def butter_filter(data, lowcut, highcut, fs, order = 5):
    omega = 0.5*fs
    low = lowcut/omega
    high = highcut/omega
    b, a = signal.butter(order, [low, high], btype = 'band')
    y = signal.lfilter(b, a, data, axis = 0)

    return y

# Collapse the Laplacian pyramid
def reconstruct_from_tensorlist(filter_tensor_list, levels = 3):
    final = np.zeros(filter_tensor_list[-1].shape)
    
    for i in range(filter_tensor_list[0].shape[0]):
        up = filter_tensor_list[0][i]
        for n in range(levels-1):
            up = cv2.pyrUp(up) + filter_tensor_list[n + 1][i]
        final[i] = up

    return final

# Motion magnification
def magnify_motion(video_name, low_f, high_f, levels = 3, amplification = 20):
    t, fps = load_video(video_name)
    lap_video_list = laplacian_video(t, levels = levels)
    filter_tensor_list = []

    for i in range(levels):
        filter_tensor = butter_filter(lap_video_list[i], low_f, high_f, fps)
        filter_tensor *= amplification
        filter_tensor_list.append(filter_tensor_list)

    recon = reconstruct_from_tensorlist(filter_tensor_list)
    final = t + recon
    save_video(final)

#Dev notes: There might be a problem with array.shape[0] and array[0].shape[0]