import os, sys
import numpy as np
import h5py


# Change this line if you want to use Mask-RCNN or SSD bounding boxes instead of H36M's "ground truth".
# BBOXES_SOURCE = 'GT'  # or 'MRCNN' or 'SSD'

# retval = {
#     'subject_names': ['S1', 'S5', 'S6', 'S7', 'S8', 'S9', 'S11'],
#     'camera_names': ['54138969', '55011271', '58860488', '60457274'],
#     'action_names': [
#         'Directions-1', 'Directions-2',
#         'Discussion-1', 'Discussion-2',
#         'Eating-1', 'Eating-2',
#         'Greeting-1', 'Greeting-2',
#         'Phoning-1', 'Phoning-2',
#         'Posing-1', 'Posing-2',
#         'Purchases-1', 'Purchases-2',
#         'Sitting-1', 'Sitting-2',
#         'SittingDown-1', 'SittingDown-2',
#         'Smoking-1', 'Smoking-2',
#         'TakingPhoto-1', 'TakingPhoto-2',
#         'Waiting-1', 'Waiting-2',
#         'Walking-1', 'Walking-2',
#         'WalkingDog-1', 'WalkingDog-2',
#         'WalkingTogether-1', 'WalkingTogether-2']
# }
def getCamParamsHuman36m(subject_names, camera_names, action_names):
    retval = {
        'subject_names': subject_names,
        'camera_names': camera_names,
        'action_names': action_names
    }

    h36m_root = '/home/daniela/catkin_ws/src/hpe/images/human36m/'
    una_dinosauria_root = h36m_root + 'extra/una-dinosauria-data/h36m/'

    cameras_params = h5py.File(una_dinosauria_root + "cameras.h5")

    cameras = {}
    for subject_idx, subject in enumerate(retval['subject_names']):
        for camera_idx, camera in enumerate(retval['camera_names']):
            cameras[camera] = {}
            assert len(cameras_params[subject.replace('S', 'subject')]) == 4
            camera_params = cameras_params[subject.replace('S', 'subject')]['camera%d' % (camera_idx + 1)]

            T = np.zeros((4, 4))
            T[0:3, 0:3] = np.array(camera_params['R']).T
            T[0:3, 3] = (-np.array(camera_params['R']).T @ camera_params['T'] / 1000).T
            T[3, 3] = 1

            cameras[camera]['R'] = np.array(camera_params['R']).T
            cameras[camera]['t'] = -np.array(camera_params['R']).T @ camera_params['T'] / 1000
            cameras[camera]['T'] = T

            cameras[camera]['K'] = np.zeros((3, 3), np.float32)
            cameras[camera]['K'][:2, 2] = camera_params['c'][:, 0]
            cameras[camera]['K'][0, 0] = camera_params['f'][0]
            cameras[camera]['K'][1, 1] = camera_params['f'][1]
            cameras[camera]['K'][2, 2] = 1.0

            cameras[camera]['dist'] = np.zeros((5), np.float32)
            cameras[camera]['dist'][:2] = camera_params['k'][:2, 0]
            cameras[camera]['dist'][2:4] = camera_params['p'][:, 0]
            cameras[camera]['dist'][4] = camera_params['k'][2, 0]

    return cameras


def getCamParamsMPI(section, sequence, camera_names):
    mpi_root = "/home/daniela/catkin_ws/src/hpe/images/mpi-inf-3dhp/"
    file_name_calibration = mpi_root + section + '/' + sequence + '/' 'camera.calibration'
    cameras = {}

    with open(file_name_calibration, 'r') as calib:
        # print(calib.readlines())
        name = ''
        for line in calib.readlines():
            # print(line)
            if line[0:4] == 'name':
                name = int(line[11:])
                name = "camera_" + str(name)
                if name in camera_names:
                    cameras[name] = {}
                    cameras[name]['dist'] = np.zeros((5), np.float32)

            if line[2:11] == 'intrinsic':
                intrinsics = line[11:].split()
                if name in camera_names:
                    cameras[name]['K'] = np.reshape([float(x) for x in intrinsics], (4, 4))[:3, :3]

            if line[2:11] == 'extrinsic':
                extrinsic = line[11:].split()

                if name in camera_names:
                    cameras[name]['T'] = np.reshape([float(x) for x in extrinsic], (4, 4))
                    # print(extrinsic)
                    # print(cameras[name]['T'][:3,3]/1000)
                    # cameras[name]['T'][:3,3]=cameras[name]['T'][:3,3]/1000
                #     print(cameras[name]['T'])
                # exit(0)

            if line[2:6] == 'size':
                height, width = line[11:].split()
                if name in camera_names:
                    cameras[name]['height'] = int(height)
                    cameras[name]['width'] = int(width)




            #
            # cameras[camera]['R'] = np.array(camera_params['R']).T
            # cameras[camera]['t'] = -np.array(camera_params['R']).T @ camera_params['T'] / 1000
            #
            # cameras[camera]['K'] = np.zeros((3, 3), np.float32)
            # cameras[camera]['K'][:2, 2] = camera_params['c'][:, 0]
            # cameras[camera]['K'][0, 0] = camera_params['f'][0]
            # cameras[camera]['K'][1, 1] = camera_params['f'][1]
            # cameras[camera]['K'][2, 2] = 1.0
    # print(cameras)
    return cameras