import bpy
import numpy as np


def get_camera_intrinsics(mode='simple'):
    scene = bpy.context.scene

    scale = scene.render.resolution_percentage / 100
    width_px = scene.render.resolution_x * scale
    height_px = scene.render.resolution_y * scale

    camdata = scene.camera.data

    K = np.zeros((3, 3), dtype=np.float32)

    if mode == 'simple':
        aspect_ratio = width / height
        K[0][0] = width_px / 2 / np.tan(camdata.angle / 2)
        K[1][1] = height_px / 2. / np.tan(camdata.angle / 2) * aspect_ratio
        K[0][2] = width_px / 2.
        K[1][2] = height_px / 2.
        K[2][2] = 1.
        K.transpose()

    if mode == 'complete':

        focal_mm = camdata.lens
        sensor_width_mm = camdata.sensor_width
        sensor_height_mm = camdata.sensor_height
        pixel_aspect_ratio = scene.render.pixel_aspect_x / scene.render.pixel_aspect_y

        if camdata.sensor_fit == 'VERTICAL':
            # the sensor height is fixed (sensor fit is horizontal), 
            # the sensor width is effectively changed with the pixel aspect ratio
            s_u = width_px / sensor_width_mm / pixel_aspect_ratio
            s_v = height_px / sensor_height_mm
        else:  # 'HORIZONTAL' and 'AUTO'
            # the sensor width is fixed (sensor fit is horizontal), 
            # the sensor height is effectively changed with the pixel aspect ratio
            s_u = width_px / sensor_width_mm
            s_v = height_px * pixel_aspect_ratio / sensor_height_mm

        # parameters of intrinsic calibration matrix K
        alpha_u = focal_mm * s_u
        alpha_v = focal_mm * s_v
        assert camdata.shift_x == 0 and camdata.shift_y == 0
        u_0 = width_px / 2
        v_0 = height_px / 2
        skew = 0  # only use rectangular pixels

        K = np.array([
            [alpha_u, skew, u_0],
            [0, alpha_v, v_0],
            [0, 0, 1]
        ], dtype=np.float32)

    return K


# Returns camera rotation and translation matrices from Blender.
def get_camera_extrinsics():
    scene = bpy.context.scene
    stereo_camdata = scene.camera.data.stereo

    if stereo_camdata.convergence_mode == 'PARALLEL':
        translation = stereo_camdata.interocular_distance  # translation between cameras
        translation = np.array([translation, 0, 0])
        # scene.camera.rotation_euler
        rotation = np.eye(3)

        return translation, rotation

    elif stereo_camdata.convergence_mode == 'TOE':
        raise NotImplementedError
    else:  # OFF-AXIS
        raise NotImplementedError
        # There are 3 coordinate systems involved:
        #    1. The World coordinates: "world"
        #       - right-handed
        #    2. The Blender camera coordinates: "bcam"
        #       - x is horizontal
        #       - y is up
        #       - right-handed: negative z look-at direction
        #    3. The desired computer vision camera coordinates: "cv"
        #       - x is horizontal
        #       - y is down (to align to the actual pixel coordinates
        #         used in digital images)
        #       - right-handed: positive z look-at direction
        #
        #  cam = scene.camera
        #
        # R_bcam2cv = np.array([
        #     [1, 0, 0],
        #     [0, -1, 0],
        #     [0, 0, -1]
        # ], dtype=np.float32)
        #
        # # Transpose since the rotation is object rotation,
        # # and we want coordinate rotation
        # # R_world2bcam = cam.rotation_euler.to_matrix().transposed()
        # # T_world2bcam = -1*R_world2bcam * location
        # #
        # # Use matrix_world instead to account for all constraints
        # location, rotation = cam.matrix_world.decompose()[0:2]
        # R_world2bcam = rotation.to_matrix().transposed()
        #
        # # Convert camera location to translation vector used in coordinate changes
        # # T_world2bcam = -1*R_world2bcam*cam.location
        # # Use location from matrix_world to account for constraints:
        # T_world2bcam = -1*R_world2bcam @ location
        #
        # # Build the coordinate transform matrix from world to computer vision camera
        # # NOTE: Use * instead of @ here for older versions of Blender
        # # TODO: detect Blender version
        # R_world2cv = R_bcam2cv@R_world2bcam
        # T_world2cv = R_bcam2cv@T_world2bcam
        #
        # # put into 3x4 matrix
        # RT = np.ndarray([
        #     R_world2cv[0][:] + (T_world2cv[0],),
        #     R_world2cv[1][:] + (T_world2cv[1],),
        #     R_world2cv[2][:] + (T_world2cv[2],)
        #     ])
        # return RT


def get_frame_size():
    render = bpy.context.scene.render
    return render.resolution_x, render.resolution_y


def write_par_file():
    K = get_camera_intrinsics('complete')
    T, R = get_camera_extrinsics()
    width, height = get_frame_size()

    with open('StereoData.par', 'wt') as fid:
        print("left.A=[", file=fid)
        print([" " + str(a) for r in K for a in r], file=fid)
        print("]\n", file=fid)

        print("left.k1= 0.0", file=fid)
        print("left.k2= 0.0", file=fid)

        print("right.A=[", file=fid)
        print([" " + str(a) for row in K for a in row], file=fid)
        print("]\n", file=fid)

        print("right.k1= 0.0", file=fid)
        print("right.k2= 0.0", file=fid)

        print("T=[", file=fid)
        print([" " + str(a) for a in T], file=fid)
        print("]\n", file=fid)

        print("R=[", file=fid)
        print([" " + str(a) for row in R for a in row], file=fid)
        print("]\n", file=fid)

        print(f"video.width= {width}", file=fid)
        print(f"video.heigth= {height}", file=fid)