import json

import cv2
import matplotlib.pyplot as plt
from mpl_toolkits.mplot3d import Axes3D  # <--- This is important for 3d plotting

import numpy as np
import h5py
import scipy.io
import sys
from pathlib import Path

sys.path.append(str(Path("../").resolve()))

from utils.projection import projectToCamera
from utils.links import joints_mpi_inf_3dhp, links_mpi_inf_3dhp


def main():
    file_name = "/home/daniela/catkin_ws/src/hpe/images/mpi-inf-3dhp/S1/Seq1/annot.mat"
    file_name_calibration = "/home/daniela/catkin_ws/src/hpe/images/mpi-inf-3dhp/S1/Seq1/camera.calibration"

    cameras = {}

    with open(file_name_calibration, 'r') as calib:
        name = ''
        for line in calib.readlines():
            if line[0:4] == 'name':
                name = int(line[11:])
                cameras[name] = {}
            if line[2:11] == 'intrinsic':
                intrinsics = line[11:].split()
                print(np.reshape([float(x) for x in intrinsics], (4, 4)))
                cameras[name]['intrinsics'] = np.reshape([float(x) for x in intrinsics], (4, 4))
            if line[2:11] == 'extrinsic':
                intrinsics = line[11:].split()
                cameras[name]['extrinsics'] = np.reshape([float(x) for x in intrinsics], (4, 4))
            if line[2:6] == 'size':
                height, width = line[11:].split()
                cameras[name]['height'] = int(height)
                cameras[name]['width'] = int(width)

    # exit(0)

    mat = scipy.io.loadmat(file_name)
    gt_3d = mat['annot3']

    camera_id = 0
    ground_truth = {}

    camera = gt_3d[camera_id][0]
    # print(camera)

    frame_idx = 0
    for frame in camera:
        ground_truth[frame_idx] = {}
        joint_idx = 0
        for n in range(0, len(frame), 3):
            pts_in_sensor = np.ndarray((4, 1), dtype=float)
            pts_in_sensor[0][0] = frame[n]  # x
            pts_in_sensor[1][0] = frame[n + 1]  # y
            pts_in_sensor[2][0] = frame[n + 2]  # z
            pts_in_sensor[3][0] = 1

            points_in_world = np.dot(np.linalg.inv(cameras[camera_id]['extrinsics']), pts_in_sensor)

            for joint_key, joint in joints_mpi_inf_3dhp.items():
                if joint['idx'] == joint_idx:
                    ground_truth[frame_idx][joint_key] = {}
                    ground_truth[frame_idx][joint_key]['pose'] = {}
                    ground_truth[frame_idx][joint_key]['pose']['x'] = points_in_world[0][0]
                    ground_truth[frame_idx][joint_key]['pose']['y'] = points_in_world[1][0]
                    ground_truth[frame_idx][joint_key]['pose']['z'] = points_in_world[2][0]

            joint_idx += 1
        frame_idx += 1

    print(ground_truth)
    output_path = '/home/daniela/catkin_ws/src/hpe/images/mpi-inf-3dhp/S1/Seq1/ground_truth.txt'

    with open(output_path, "w") as fp:
        # Load the dictionary from the file
        json.dump(ground_truth, fp)
        print("Ground truth file saved successfully!")

if __name__ == "__main__":
    main()