import json
import sys
import time

import numpy as np

sys.path.append("../swig/")
import rpt

# ==================================================================================================


def main():
    print("")

    # Test camera structure
    camera = rpt.Camera()
    camera.name = "Camera 1"
    camera.K = [[1, 0, 0], [0, 1, 0], [0, 0, 1]]
    camera.DC = [0, 0, 0, 0, 0]
    camera.R = [[1, 0, 0], [0, 1, 0], [0, 0, 1]]
    camera.T = [[1], [2], [3]]
    camera.width = 640
    camera.height = 480
    print(camera)
    print("")

    # Load input data
    roomparams = [[4.8, 6.0, 2.0], [0, 0, 1.0]]
    joint_names = [
        "nose",
        "eye_left",
        "eye_right",
        "ear_left",
        "ear_right",
        "shoulder_left",
        "shoulder_right",
        "elbow_left",
        "elbow_right",
        "wrist_left",
        "wrist_right",
        "hip_left",
        "hip_right",
        "knee_left",
        "knee_right",
        "ankle_left",
        "ankle_right",
        "hip_middle",
        "shoulder_middle",
        "head",
    ]
    cpath = "/RapidPoseTriangulation/data/h1/sample.json"
    ppath = "/RapidPoseTriangulation/tests/poses_h1.json"
    with open(cpath, "r") as file:
        cdata = json.load(file)
    with open(ppath, "r") as file:
        pdata = json.load(file)
    cams = cdata["cameras"]
    poses_2d = pdata["2D"]
    cameras = rpt.convert_cameras(cams)

    # Run triangulation
    triangulator = rpt.Triangulator(min_match_score=0.95)
    stime = time.time()
    poses_3d = triangulator.triangulate_poses(
        poses_2d, cameras, roomparams, joint_names
    )
    print("3D time:", time.time() - stime)
    print(np.array(poses_3d))
    print("")

    # Load input data
    roomparams = [[5.6, 6.4, 2.4], [0, -0.5, 1.2]]
    cpath = "/RapidPoseTriangulation/data/p1/sample.json"
    ppath = "/RapidPoseTriangulation/tests/poses_p1.json"
    with open(cpath, "r") as file:
        cdata = json.load(file)
    with open(ppath, "r") as file:
        pdata = json.load(file)
    cams = cdata["cameras"]
    poses_2d = pdata["2D"]
    cameras = rpt.convert_cameras(cams)

    # Run triangulation
    triangulator.reset()
    stime = time.time()
    poses_3d = triangulator.triangulate_poses(
        poses_2d, cameras, roomparams, joint_names
    )
    print("3D time:", time.time() - stime)
    print(np.array(poses_3d))
    print("")

    # Run again to test last pose cache
    stime = time.time()
    poses_3d = triangulator.triangulate_poses(
        poses_2d, cameras, roomparams, joint_names
    )
    print("3D time:", time.time() - stime)
    print(np.array(poses_3d))
    print("")

    # Load input data
    roomparams = [[6.0, 5.0, 2.0], [1.5, 1.0, -0.5]]
    cpath = "/RapidPoseTriangulation/data/e1/sample.json"
    ppath = "/RapidPoseTriangulation/tests/poses_e1.json"
    with open(cpath, "r") as file:
        cdata = json.load(file)
    with open(ppath, "r") as file:
        pdata = json.load(file)
    cams = cdata["cameras"]
    poses_2d = pdata["2D"]
    cameras = rpt.convert_cameras(cams)

    # Run triangulation
    triangulator.reset()
    stime = time.time()
    poses_3d = triangulator.triangulate_poses(
        poses_2d, cameras, roomparams, joint_names
    )
    print("3D time:", time.time() - stime)
    print(np.array(poses_3d))
    print("")

    triangulator.print_stats()
    print("")


# ==================================================================================================

if __name__ == "__main__":
    main()
