#! /usr/bin/env python

import os
import numpy as np
import cv2
from PIL import Image
import cvui
import sys
import copy
import argparse


def main(file_path, window_size_width, window_size_height):

    WINDOW_NAME = 'Simple Fisheye Calibrator'
    cvui.init(WINDOW_NAME)

    cap = None
    FILE_PATH = file_path
    if not file_path.isdecimal():
        dir_name, file_name = os.path.split(FILE_PATH)
        file_name_short, ext = os.path.splitext(file_name)
        image = Image.open(FILE_PATH)
        image = np.array(image, dtype=np.uint8)
        h, w = image.shape[0], image.shape[1]
    else:
        cap = cv2.VideoCapture(int(FILE_PATH))
        w = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))
        h = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))

    display_scale = 1.00
    required_size_of_image_disp_area_w = int(600+w*display_scale)
    required_size_of_image_disp_area_h = int(70+h*display_scale)

    if required_size_of_image_disp_area_h > window_size_height or \
        required_size_of_image_disp_area_w > window_size_width:

        window_size_height = required_size_of_image_disp_area_h + 10
        window_size_width = required_size_of_image_disp_area_w + 10

    frame = np.zeros((window_size_height, window_size_width, 3), np.uint8)

    view_scale_format = '%.2Lf'
    SCALE_VALUE_DEFAULT = [1.30]
    scale_value = copy.deepcopy(SCALE_VALUE_DEFAULT)
    PARAM1_DEFAULT = [-0.78]
    PARAM2_DEFAULT = [-0.01]
    PARAM3_DEFAULT = [0.00]
    PARAM4_DEFAULT = [0.00]
    PARAM5_DEFAULT = [0.53]
    param1 = copy.deepcopy(PARAM1_DEFAULT)
    param2 = copy.deepcopy(PARAM2_DEFAULT)
    param3 = copy.deepcopy(PARAM3_DEFAULT)
    param4 = copy.deepcopy(PARAM4_DEFAULT)
    param5 = copy.deepcopy(PARAM5_DEFAULT)

    while True:
        img = None
        if not cap:
            img = image.copy()
        else:
            ret, img = cap.read()
            if not ret:
                continue
            img = img[:,:,::-1].astype(np.uint8)

        frame[:] = (60, 60, 60)

        cvui.beginRow(frame, 10, 20, 100, 50)
        cvui.text('FilePath:', 1.4)
        cvui.text(FILE_PATH, 1.4)
        cvui.endRow()

        cvui.beginRow(frame, 10, 70, 100, 50)
        cvui.text('Scale   :', 1.4)
        cvui.trackbar(300, scale_value, 0.01, 2.00, 1, view_scale_format)
        cvui.endRow()

        cvui.beginRow(frame, 10, 120, 100, 50)
        cvui.endRow()

        cvui.beginRow(frame, 10, 170, 100, 50)
        cvui.text('Distortion Parameters', 1.4)
        cvui.endRow()

        cvui.beginRow(frame, 10, 220, 100, 50)
        cvui.text('Param1(k1): ', 1.0)
        cvui.trackbar(300, param1, -2.00, 2.00, 1, view_scale_format)
        cvui.endRow()

        cvui.beginRow(frame, 10, 270, 100, 50)
        cvui.text('Param2(k2): ', 1.0)
        cvui.trackbar(300, param2, -2.00, 2.00, 1, view_scale_format)
        cvui.endRow()

        cvui.beginRow(frame, 10, 320, 100, 50)
        cvui.text('Param3(p1): ', 1.0)
        cvui.trackbar(300, param3, -1.00, 1.00, 1, view_scale_format)
        cvui.endRow()

        cvui.beginRow(frame, 10, 370, 100, 50)
        cvui.text('Param4(p2): ', 1.0)
        cvui.trackbar(300, param4, -1.00, 1.00, 1, view_scale_format)
        cvui.endRow()

        cvui.beginRow(frame, 10, 420, 100, 50)
        cvui.text('Param5(k3): ', 1.0)
        cvui.trackbar(300, param5, -2.00, 2.00, 1, view_scale_format)
        cvui.endRow()

        mtx = np.array([[h,    0.,  w / 2],
                        [0.,   w,   h / 2],
                        [0.,   0.,  1    ]])
        dist = np.array(
            [
                param1,
                param2,
                param3,
                param4,
                param5
            ]
        )
        n_mtx = cv2.getOptimalNewCameraMatrix(
            mtx,
            dist,
            (img.shape[1], img.shape[0]),
            0
        )[0]
        map = cv2.initUndistortRectifyMap(
            mtx,
            dist,
            np.eye(3),
            n_mtx,
            (img.shape[1], img.shape[0]),
            cv2.CV_32FC1
        )

        mapx = map[0]*scale_value[0] + (1-scale_value[0])*w/2
        mapy = map[1]*scale_value[0] + (1-scale_value[0])*h/2

        img = cv2.remap(
            img,
            mapx,
            mapy,
            cv2.INTER_CUBIC
        )

        view_image = cv2.resize(
            img,
            None,
            fx=display_scale,
            fy=display_scale,
            interpolation=cv2.INTER_NEAREST
        )
        cvui.image(frame, 600, 70, view_image[:,:,::-1])

        if cvui.button(frame, 80, 500, 100, 30, "QUIT"):
            break

        if cvui.button(frame, 230, 500, 100, 30, "RESET"):
            scale_value = copy.deepcopy(SCALE_VALUE_DEFAULT)
            param1 = copy.deepcopy(PARAM1_DEFAULT)
            param2 = copy.deepcopy(PARAM2_DEFAULT)
            param3 = copy.deepcopy(PARAM3_DEFAULT)
            param4 = copy.deepcopy(PARAM4_DEFAULT)
            param5 = copy.deepcopy(PARAM5_DEFAULT)

        if cvui.button(frame, 380, 500, 100, 30, "SAVE"):
            if dir_name:
                path = f'{dir_name}/{file_name_short}_{param1[0]}_{param2[0]}_{param3[0]}_{param4[0]}_{param5[0]}{ext}'
            else:
                path = f'./{file_name_short}_{param1[0]}_{param2[0]}_{param3[0]}_{param4[0]}_{param5[0]}{ext}'
            print(path)
            cv2.imwrite(path, img[:,:,::-1])

        cvui.imshow(WINDOW_NAME, frame)
        cvui.update()

        # ESC
        if cv2.waitKey(20) == 27:
            break

if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    parser.add_argument('--file_path', type=str, default='fisheye_test.jpg', help='File path of the still image (e.g. xxx.jpg) or device number of the camera (e.g. 0)')
    parser.add_argument('--window_size_width', type=int, default=1500, help='Default window size width')
    parser.add_argument('--window_size_height', type=int, default=700, help='Default window size height')
    args = parser.parse_args()

    file_path = args.file_path
    window_size_width = args.window_size_width
    window_size_height = args.window_size_height

    main(file_path, window_size_width, window_size_height)