import json
import cv2
import numpy as np
from pathlib import Path


def load_calibration(calibration_json_filename: Path):
    # with open(dataset + '/calibration-chessboard4.json', 'r') as fp:
    # with open(dataset + '/calibration-chessboard4.json', 'r') as fp:
    with open(calibration_json_filename, 'r') as fp:
        calibdata = json.load(fp)
    print(calibdata)
    if 'type' in calibdata and calibdata['type'] == 'fisheye':
        dim1 = calibdata['dim1']
        dim2 = calibdata['dim2']
        dim3 = calibdata['dim3']
        K = np.array(calibdata['K'])
        D = np.array(calibdata['D'])
        new_K = np.array(calibdata['new_K'])
        scaled_K = np.array(calibdata['scaled_K'])
        balance = calibdata['balance']
        
        # print(f"{K=}")
            
        map1, map2 = cv2.fisheye.initUndistortRectifyMap(scaled_K, D, np.eye(3), new_K, dim3, cv2.CV_16SC2)

        undistorted_img = lambda img: cv2.remap(img, map1, map2, interpolation=cv2.INTER_LINEAR, borderMode=cv2.BORDER_CONSTANT)
        undistort_points = lambda distorted_img_points: cv2.fisheye.undistortPoints (np.array([distorted_img_points]).astype(np.float32), K=scaled_K, D=D, R=np.eye(3), P=new_K)
        # undistorted_img = cv2.remap(img, map1, map2, interpolation=cv2.INTER_LINEAR, borderMode=cv2.BORDER_CONSTANT)
        # new_points = cv2.fisheye.undistortPoints (np.array([distorted_img_points]).astype(np.float32), K=scaled_K, D=D, R=np.eye(3), P=new_K)
        w, h = dim1[0], dim1[1]
    else:
        mtx = np.array(calibdata['camera_matrix'])
        dist = np.array(calibdata['dist_coeff'])
        w, h = calibdata['dim']['width'], calibdata['dim']['height']




        # # Refining the camera matrix using parameters obtained by calibration
        # if we don't set this, the new image will be cropped to the minimum size
        # this way, no cropping occurs
        # w, h = 2560, 1440
        newcameramtx, roi = cv2.getOptimalNewCameraMatrix(mtx, dist, (w,h), 1, (w,h))

        undistorted_img = lambda img: cv2.undistort(img, mtx, dist, None, newcameramtx)
        undistort_points = lambda distorted_img_points: cv2.undistortPoints(np.array(distorted_img_points).astype('float32'), mtx, dist, None, P=newcameramtx)
    
    return w, h, undistorted_img, undistort_points