import sys
import os
import glob
import numpy as np
import cv2
import imageio
from scipy import interpolate
import matplotlib
import matplotlib.pyplot as plt


if __name__=="__main__":

    scale = 1

    coord_file = sys.argv[1]
    base_image_path = sys.argv[2]
    depth_channel = 0
    if len(sys.argv) > 3:
        depth_channel = int(sys.argv[3])

    # parse the coord file: frame-x-y
    with open(coord_file, 'r') as f:
        lines = f.readlines()

    frame_num = []
    x = []
    y = []

    for line in lines:
        bits = line.split('-')
        frame_num.append(int(bits[0]))
        x.append(int(bits[1]))
        y.append(int(bits[2]))


    frame_num = np.array(frame_num)
    x = np.array(x)
    y = np.array(y)

    frame_num = frame_num[::2]
    x1 = x[::2]
    x2 = x[1::2]
    y1 = y[::2]
    y2 = y[1::2]

    # upsample the tracks to the length of frames
    image_paths = glob.glob(os.path.join(base_image_path, '*.tiff'))
    nimages = len(image_paths)

    frames = range(0, nimages)
    x1f = interpolate.interp1d(frame_num, x1, kind="cubic", fill_value="extrapolate")
    x2f = interpolate.interp1d(frame_num, x2, kind="cubic", fill_value="extrapolate")
    y1f = interpolate.interp1d(frame_num, y1, kind="cubic", fill_value="extrapolate")
    y2f = interpolate.interp1d(frame_num, y2, kind="cubic", fill_value="extrapolate")

    x1i = x1f(frames)
    x2i = x2f(frames)
    y1i = y1f(frames)
    y2i = y2f(frames)

    rot_angle = np.arctan2(x2i - x1i, y2i - y1i) * 180 / np.pi
    rot_poly = np.polyfit(frames, rot_angle, 27)
    rot_poly_f = np.poly1d(rot_poly)
    fig, ax = plt.subplots()

    ax.plot(frames, rot_angle)
    ax.plot(frames, rot_poly_f(frames))
    #ax.plot(frames, y1i)

    plt.show()

    rot_smooth = rot_poly_f(frames)

    # Shift images and save to dir

    for ind, image_path in enumerate(image_paths):
        print(image_path)
        im = imageio.imread(image_path)
        if depth_channel > 0:
            im = im[:, :, depth_channel]
        cent = im.shape[0]/2
        del_x = cent - x1i[ind]/scale
        del_y = cent - y1i[ind]/scale

        # Translate image
        M = np.float32([[1, 0, del_y], [0, 1, del_x]])
        dst = cv2.warpAffine(im, M, im.shape[0:2]) # assumes square image

        # Rotate Image

        #rot_angle = np.arctan2(x2i[ind]-x1i[ind], y2i[ind]-y1i[ind])*180/np.pi

        rows, cols = dst.shape[0:2]
        M = cv2.getRotationMatrix2D((cols / 2, rows / 2), rot_smooth[ind], 1)
        dst = cv2.warpAffine(dst, M, (cols, rows))

        #cv2.imshow('shifted', dst.astype('uint8'))
        #cv2.waitKey(1)
        cv2.imwrite(image_path[:-4] + '-shifted.tif', dst)