import sys
import os
import argparse
import numpy as np
import cv2
from tqdm import tqdm

sys.path.append(os.path.abspath(sys.path[0] + '/../../'))

parser = argparse.ArgumentParser()
parser.add_argument('-s', '--start', type=int)
parser.add_argument('-t', '--duration', type=int)
parser.add_argument('--fps', type=str, required=True)
parser.add_argument('datadir', type=str)
args = parser.parse_args()

os.chdir(args.datadir)

rawK = np.array([
    [1369.757446, 0., 1838.643555, 0., 1369.757446, 1524.068604, 0., 0., 1.],
    [1367.517944, 0., 1840.157837, 0., 1367.517944, 1536.036133, 0., 0., 1.],
    [1369.830322, 0., 1827.990723, 0., 1369.830322, 1514.463135, 0., 0., 1.],
    [1368.966187, 0., 1829.976196, 0., 1368.966187, 1512.734375, 0., 0., 1.],
    [1373.654297, 0., 1838.130859, 0., 1373.654297, 1534.985840, 0., 0., 1.],
    [1365.853027, 0., 1835.100830, 0., 1365.853027, 1533.032959, 0., 0., 1.]
]).reshape(-1, 3, 3)
D = np.array([[-0.044752], [-0.006285], [0.000000], [0.000000]])

mean_focal = np.mean(rawK[:, 0, 0])
K = None # 1369.2632038333334, 1500.0, 1900.0

for i in range(6):
    # Extract frames from video
    os.makedirs(f"raw_images/{i + 1}", exist_ok=True)
    extra_args = []
    if args.start is not None:
        extra_args.append(f"-ss {args.start}")
    if args.duration is not None:
        extra_args.append(f"-t {args.duration}")
    extra_args = ' '.join(extra_args)
    os.system(f"ffmpeg -i raw_video/{i + 1:02d}.mov {extra_args} -f image2 -q:v 2 -vf fps={args.fps} "
              f"raw_images/{i + 1}/image%03d.png")
    
    # Undistort frames and collect
    os.makedirs(f"images", exist_ok=True)
    raw_image_files = os.listdir(f"raw_images/{i + 1}")
    map1, map2 = None, None
    for raw_file in tqdm(raw_image_files):
        raw_im = cv2.imread(f"raw_images/{i + 1}/{raw_file}")
        if K is None:
            K = np.array([[mean_focal, 0., raw_im.shape[1] / 2],
                          [0., mean_focal, raw_im.shape[0] / 2],
                          [0., 0., 1.]])
            tqdm.write(
                f"Intrinsic parameters: {mean_focal}, {raw_im.shape[0] / 2}, {raw_im.shape[1] / 2}")
        if map1 is None:
            map1, map2 = cv2.fisheye.initUndistortRectifyMap(
                rawK[i], D, None, K, (raw_im.shape[1], raw_im.shape[0]), cv2.CV_16SC2)
        im = cv2.remap(raw_im, map1, map2, interpolation=cv2.INTER_LINEAR,
                       borderMode=cv2.BORDER_CONSTANT)
        im = cv2.rotate(im, cv2.ROTATE_90_COUNTERCLOCKWISE)
        cv2.imwrite(f"images/image{i}{raw_file[5:]}", im)