import cv2
import numpy as np

has_cuda = cv2.cuda.getCudaEnabledDeviceCount()
print(has_cuda)  # if one cuda is getCudaEnabledDeviceCount

left = "Left.mp4"
delay = 0
clip1 = None

delay = int(1000 / 20)


def straighten_image(image, coefficients):
    DIST_COEFF = np.array([coefficients[0], coefficients[1], coefficients[2], 0, 0])
    camera_matrix = np.array(
        [[3000, 0, image.shape[1] / 2], [0, 3000, image.shape[0] / 2], [0, 0, 1]],
        dtype="double",
    )
    return cv2.undistort(image, camera_matrix, DIST_COEFF)


clip1 = cv2.VideoCapture(left)
fps = clip1.get(cv2.CAP_PROP_FPS)
delay = int(1000 / 20)

height = 0


def init_undistort(image, coefficients):
    DIST_COEFF = np.array([coefficients[0], coefficients[1], coefficients[2], 0, 0])
    camera_matrix = np.array(
        [[3000, 0, image.shape[1] / 2], [0, 3000, image.shape[0] / 2], [0, 0, 1]],
        dtype="double",
    )

    new_K, _ = cv2.getOptimalNewCameraMatrix(
        camera_matrix, DIST_COEFF, (image.shape[0], image.shape[1]), 0
    )

    map1, map2 = cv2.initUndistortRectifyMap(
        camera_matrix,
        DIST_COEFF,
        None,
        camera_matrix,
        (4000, 4000),
        cv2.CV_32FC1,
    )
    return map1, map2


def start():
    index = 0
    COEFFICIENTS_1 = [-0.199, -0.1300, -0.0150]

    r, frame = clip1.read()

    map1, map2 = init_undistort(frame, COEFFICIENTS_1)

    while True:
        index += 1

        ret1, frame1 = clip1.read()

        if not ret1:
            # either 1 of the video's ended, break out of it
            break

        frame1 = cv2.remap(frame1, map1, map2, interpolation=cv2.INTER_LINEAR)
        # frame1 = straighten_image(frame1, COEFFICIENTS_1)

        # Resize final frame and write to final clip
        final_frame = cv2.resize(frame1, (1920, 1080), interpolation=cv2.INTER_LINEAR)

        # final_clip.write(final_frame)
        cv2.imshow("Video", final_frame)
        key = cv2.waitKey(delay) & 0xFF


start()