import json
import sys
import time

import numpy as np

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

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


def main():
    print("")

    # Test camera structure
    camera = spt.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 = [[0, 0, 1.0], [4.8, 6.0, 2.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 = "/SimplePoseTriangulation/data/h1/sample.json"
    ppath = "/SimplePoseTriangulation/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 = spt.convert_cameras(cams)

    # Run triangulation
    triangulator = spt.Triangulator(min_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))

    # Load input data
    roomparams = [[0, -0.5, 1.2], [5.6, 6.4, 2.4]]
    cpath = "/SimplePoseTriangulation/data/p1/sample.json"
    ppath = "/SimplePoseTriangulation/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 = spt.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))


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

if __name__ == "__main__":
    main()