#!/usr/bin/env python3
# -*- coding: utf-8 -*-
## License: Apache 2.0. See LICENSE file in root directory.
## Copyright(c) 2019 Intel Corporation. All Rights Reserved.

import pyrealsense2 as rs
import argparse
import sys

serial_number = "845412110431"
K1 = [[286.95726457967322, 0, 428.89434809446971], [0, 287.23657369198514, 401.3023457694253]] # row-major
D1 = [-0.0057473528100385664, 0.041866358798935424, -0.04271161675972359, 0.0087687435847667791]
R_p_1 = [0.99998915195465088, 0.00018631802231539041, -0.0046501094475388527, 0.00019517501641530544, -0.99999845027923584, 0.0019043015781790018, -0.0046497462317347527, -0.0019051884301006794, -0.99998760223388672] # column-major
t_p_1 = [-0.032806921750307083, -0.00018981023458763957, -0.00022083433577790856]
K2 = [[287.45813871045891, 0, 429.41319687073536], [0, 287.9870499155972, 401.53394554329969]] # row-major
D2 = [-0.0077079909534496262, 0.04443351597428985, -0.043068452506983802, 0.0083743519521316048]
R_p_2 = [0.99998915195465088, -0.00019517501641530544, 0.00464974669739604, -0.00018631802231539041, -0.99999845027923584, -0.0019051881972700357, 0.00465010991320014, 0.0019043013453483582, -0.99998760223388672] # column-major
t_p_2 = [0.032806918025016785, 0.00018981023458763957, 0.00022083433577790856]

def query_tm2():
    tm2 = None
    ctx = rs.context()
    devs = ctx.query_devices()
    for dev in devs:
        tm2 = dev.as_tm2()
    if not tm2:
        print("No device found")

    assert tm2.get_info(rs.camera_info.serial_number) == serial_number, "Serial number of camera found does not match the calibration"

    return tm2

def lrs_intrinsics(K, D):
    fe_intrinsics = rs.intrinsics()  # width: 0, height: 0, ppx: 0, ppy: 0, fx: 0, fy: 0, model: None, coeffs: [0, 0, 0, 0, 0]
    fe_intrinsics.width = 848
    fe_intrinsics.height = 800
    fe_intrinsics.fx = K[0][0]
    fe_intrinsics.fy = K[1][1]
    fe_intrinsics.ppx = K[0][2]
    fe_intrinsics.ppy = K[1][2]
    fe_intrinsics.model = rs.distortion.kannala_brandt4  # kb4
    fe_intrinsics.coeffs = D + [0]
    return fe_intrinsics

def lrs_extrinsics(R,t):
    ext = rs.extrinsics()
    ext.rotation = R
    ext.translation = t
    return ext

def write_extrinsics(tm2, R_p_1, t_p_1, R_p_2, t_p_2):
    ext_fe1 = lrs_extrinsics(R_p_1, t_p_1)
    ext_fe2 = lrs_extrinsics(R_p_2, t_p_2)
    tm2.set_extrinsics(rs.stream.fisheye, 1, rs.stream.pose, 0, ext_fe1)
    tm2.set_extrinsics(rs.stream.fisheye, 2, rs.stream.pose, 0, ext_fe2)

def reset_calibration():
    tm2 = query_tm2()
    if tm2 is not None:
        print("Resetting...")
        tm2.reset_to_factory_calibration()
        print("Finished")

if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    parser.add_argument('--write', default=False, help='write calibration to device', action='store_true')
    parser.add_argument('--confirm', default=False, help='skip write prompt', action='store_true')
    parser.add_argument('--reset', default=False, help='reset calibration to factory default', action='store_true')

    if len(sys.argv) == 1:
        parser.print_help(sys.stderr)
        sys.exit(1)

    args = parser.parse_args()

    # utility function to reset calibration to factory calibration
    if args.reset:
        reset_calibration()
        sys.exit()

    # write to device
    if args.write:
        print()
        key = None
        while key not in ['y', 'n'] and not args.confirm:
            key = input("Write to device? [y/n]: ")
        if key == 'n' and not args.confirm:
            sys.exit()
        else:
            tm2 = query_tm2()
            if tm2 is not None:
                print("Writing to device...")
                fe_intrinsics = lrs_intrinsics(K1, D1)
                tm2.set_intrinsics(1, fe_intrinsics)

                fe_intrinsics = lrs_intrinsics(K2, D2)
                tm2.set_intrinsics(2, fe_intrinsics)

                # extrinsics
                write_extrinsics(tm2, R_p_1, t_p_1, R_p_2, t_p_2)

                # flush to eeprom
                tm2.write_calibration()

                print("Finished")