# Copyright (c) 2018-present, Facebook, Inc.
# All rights reserved.
#
# This source code is licensed under the license found in the
# LICENSE file in the root directory of this source tree.
#

import copy
import numpy as np
from lib.skeleton.skeleton import Skeleton
from lib.dataset.mocap_dataset import MocapDataset
from lib.camera.camera import CameraInfoPacket

h36m_skeleton = Skeleton(parents=[-1, 0, 1, 2, 3, 4, 0, 6, 7, 8, 9, 0, 11, 12, 13, 14, 12,
                                  16, 17, 18, 19, 20, 19, 22, 12, 24, 25, 26, 27, 28, 27, 30],
                         joints_left=[6, 7, 8, 9, 10, 16, 17, 18, 19, 20, 21, 22, 23],
                         joints_right=[1, 2, 3, 4, 5, 24, 25, 26, 27, 28, 29, 30, 31])

h36m_cameras_intrinsic_params = [
    {
        'id': '54138969',
        'center': [512.541504956548, 515.4514869776],
        'focal_length': [1145.04940458804, 1143.78109572365],
        'radial_distortion': [-0.207098910824901, 0.247775183068982, -0.00307515035078854],
        'tangential_distortion': [-0.00142447157470321, -0.000975698859470499],
        'res_w': 1000,
        'res_h': 1002,
        'azimuth': 70,  # Only used for visualization
    },
    {
        'id': '55011271',
        'center': [508.848621645943, 508.064917088557],
        'focal_length': [1149.67569986785, 1147.59161666764],
        'radial_distortion': [-0.194213629607385, 0.240408539138292, 0.00681997559022603],
        'tangential_distortion': [-0.0027408943961907, -0.001619026613787],
        'res_w': 1000,
        'res_h': 1000,
        'azimuth': -70,  # Only used for visualization
    },
    {
        'id': '58860488',
        'center': [519.815837182153, 501.402658888552],
        'focal_length': [1149.14071676148, 1148.7989685676],
        'radial_distortion': [-0.208338188251856, 0.255488007488945, -0.00246049749891915],
        'tangential_distortion': [-0.000759999321030303, 0.00148438698385668],
        'res_w': 1000,
        'res_h': 1000,
        'azimuth': 110,  # Only used for visualization
    },
    {
        'id': '60457274',
        'center': [514.968197319863, 501.882018537695],
        'focal_length': [1145.51133842318, 1144.77392807652],
        'radial_distortion': [-0.198384093827848, 0.218323676298049, -0.00894780704152122],
        'tangential_distortion': [-0.00181336200488089, -0.000587205583421232],
        'res_w': 1000,
        'res_h': 1002,
        'azimuth': -110,  # Only used for visualization
    },
]

h36m_cameras_extrinsic_params = {
    'S1': [
        {
            'translation': [-346.05078140028075, 546.9807793144001, 5474.481087434061],
            'R':
                [
                    [-0.9153617321513369, 0.40180836633680234, 0.02574754463350265],
                    [0.051548117060134555, 0.1803735689384521, -0.9822464900705729],
                    [-0.399319034032262, -0.8977836111057917, -0.185819527201491]
                ]
        },
        {
            'translation': [251.42516271750836, 420.9422103702068, 5588.195881837821],
            'R':
                [
                    [0.9281683400814921, 0.3721538354721445, 0.002248380248018696],
                    [0.08166409428175585, -0.1977722953267526, -0.976840363061605],
                    [-0.3630902204349604, 0.9068559102440475, -0.21395758897485287]
                ]
        },
        {
            'translation': [480.482559565337, 253.83237471361554, 5704.2076793704555],
            'R':
                [
                    [-0.9141549520542256, -0.4027780222811878, -0.045722952682337906],
                    [-0.04562341383935875, 0.21430849526487267, -0.9756999400261069],
                    [0.40278930937200774, -0.889854894701693, -0.214287280609606]
                ]
        },
        {
            'translation': [51.88347637559197, 378.4208425426766, 4406.149140878431],
            'R':
                [
                    [0.9141562410494211, -0.40060705854636447, 0.061905989962380774],
                    [-0.05641000739510571, -0.2769531972942539, -0.9592261660183036],
                    [0.40141783470104664, 0.8733904688919611, -0.2757767409202658]
                ]
        },
    ],
    'S2': [
        {},
        {},
        {},
        {},
    ],
    'S3': [
        {},
        {},
        {},
        {},
    ],
    'S4': [
        {},
        {},
        {},
        {},
    ],
    'S5': [
        {
            'translation': [-219.3059666108619, 544.4787497640639, 5518.740477016156],
            'R':
                [
                    [-0.9042074184788829, 0.42657831374650107, 0.020973473936051274],
                    [0.06390493744399675, 0.18368565260974637, -0.9809055713959477],
                    [-0.4222855708380685, -0.8856017859436166, -0.1933503902128034]
                ]
        },
        {
            'translation': [103.90282067751986, 395.67169468951965, 5767.97265758172],
            'R':
                [
                    [0.9222116004775194, 0.38649075753002626, 0.012274293810989732],
                    [0.09333184463870337, -0.19167233853095322, -0.9770111982052265],
                    [-0.3752531555110883, 0.902156643264318, -0.21283434941998647]
                ]
        },
        {
            'translation': [520.3272318446208, 283.3690958234795, 5591.123958858676],
            'R':
                [
                    [-0.9258288614330635, -0.3728674116124112, -0.06173178026768599],
                    [-0.023578112500148365, 0.220000562347259, -0.9752147584905696],
                    [0.3772068291381898, -0.9014264506460582, -0.21247437993123308]
                ]
        },
        {
            'translation': [-79.116431351199, 425.59047114848386, 4454.481629705836],
            'R':
                [
                    [0.9222815489764817, -0.3772688722588351, 0.0840532119677073],
                    [-0.021177649402562934, -0.26645871124348197, -0.9636136478735888],
                    [0.3859381447632816, 0.88694303832152, -0.25373962085111357]
                ]
        },
    ],
    'S6': [
        {
            'translation': [-239.5182864132218, 545.8141831785044, 5523.931578633363],
            'R':
                [
                    [-0.9149503344107554, 0.4034864343564006, 0.008036345687245266],
                    [0.07174776353922047, 0.1822275975157708, -0.9806351824867137],
                    [-0.3971374371533952, -0.896655898321083, -0.19567845056940925]
                ]
        },
        {
            'translation': [169.02510061389722, 409.6671223380997, 5714.338002825065],
            'R':
                [
                    [0.9197364689900042, 0.39209901596964664, 0.018525368698999664],
                    [0.101478073351267, -0.19191459963948, -0.9761511087296542],
                    [-0.37919260045353465, 0.899681692667386, -0.21630030892357308]
                ]
        },
        {
            'translation': [521.9864793089763, 286.28272817103516, 5643.2724406159],
            'R':
                [
                    [-0.916577698818659, -0.39393483656788014, -0.06856140726771254],
                    [-0.01984531630322392, 0.21607069980297702, -0.9761760169700323],
                    [0.3993638509543854, -0.8933805444629346, -0.20586334624209834]
                ]
        },
        {
            'translation': [-56.29675276801464, 420.29579722027506, 4499.322693551688],
            'R':
                [
                    [0.9182950552949388, -0.3850769011116475, 0.09192372735651859],
                    [-0.015534985886560007, -0.26706146429979655, -0.9635542737695438],
                    [0.3955917790277871, 0.8833990913037544, -0.25122338635033875]
                ]
        },
    ],
    'S7': [
        {
            'translation': [-323.9118424584857, 541.7715234126381, 5506.569132699328],
            'R':
                [
                    [-0.9055764231419416, 0.42392653746206904, 0.014752378956221508],
                    [0.06862812683752326, 0.18074371881263407, -0.9811329615890764],
                    [-0.41859469903024304, -0.8874784498483331, -0.19277053457045695]
                ]
        },
        {
            'translation': [178.6238708832376, 403.59193467821774, 5694.8801003668095],
            'R':
                [
                    [0.9212640765077017, 0.3886011826562522, 0.01617473877914905],
                    [0.09922277503271489, -0.1946115441987536, -0.9758489574618522],
                    [-0.3760682680727248, 0.9006194910741931, -0.21784671226815075]
                ]
        },
        {
            'translation': [441.1064712697594, 271.91614362573955, 5660.120611352617],
            'R':
                [
                    [-0.9245069728829368, -0.37555597339631824, -0.06515034871105972],
                    [-0.018955014220249332, 0.21601110989507338, -0.9762068980691586],
                    [0.38069353097569036, -0.9012751584550871, -0.20682244613440448]
                ]
        },
        {
            'translation': [25.768533743836343, 431.05581759025813, 4461.872981411145],
            'R':
                [
                    [0.9228353966173104, -0.3744001545228767, 0.09055029013436408],
                    [-0.014982084363704698, -0.269786590656035, -0.9628035794752281],
                    [0.3849030629889691, 0.8871525910436372, -0.25457791009093983]
                ]
        },
    ],
    'S8': [
        {
            'translation': [-82.70216069652597, 552.1896311377282, 5557.353609418419],
            'R':
                [
                    [-0.9115694669712032, 0.4106494283805017, 0.020202818036194434],
                    [0.060907749548984036, 0.1834736632003901, -0.9811359034082424],
                    [-0.40660958293025334, -0.8931430243150293, -0.19226072190306673]
                ]
        },
        {
            'translation': [-209.06289992510443, 375.0691429434037, 5818.276676972416],
            'R':
                [
                    [0.931016282525616, 0.3647626932499711, 0.01252434769597448],
                    [0.08939715221301257, -0.19463753190599434, -0.9767929055586687],
                    [-0.35385990285476776, 0.9105297407479727, -0.2138194574051759]
                ]
        },
        {
            'translation': [623.0985110132146, 290.9053651845054, 5534.379001592981],
            'R':
                [
                    [-0.9209075762929309, -0.3847355178017309, -0.0625125368875214],
                    [-0.02568138180824641, 0.21992027027623712, -0.9751797482259595],
                    [0.38893405939143305, -0.8964450100611084, -0.21240678280563546]
                ]
        },
        {
            'translation': [-178.36705625795474, 423.4669232560848, 4421.6448791590965],
            'R':
                [
                    [0.927667052235436, -0.3636062759574404, 0.08499597802942535],
                    [-0.01666268768012713, -0.26770413351564454, -0.9633570738505596],
                    [0.37303645269074087, 0.8922583555131325, -0.2543989622245125]
                ]
        },
    ],
    'S9': [
        {
            'translation': [-321.2078335720134, 467.13452033013084, 5514.330338522134],
            'R':
                [
                    [-0.9033486204435297, 0.4269119782787646, 0.04132109321984796],
                    [0.04153061098352977, 0.182951140059007, -0.9822444139329296],
                    [-0.4268916470184284, -0.8855930460167476, -0.18299857527497945]
                ]
        },
        {
            'translation': [19.193095609487138, 404.22842728571936, 5702.169280033924],
            'R':
                [
                    [0.9315720471487059, 0.36348288012373176, -0.007329176497134756],
                    [0.06810069482701912, -0.19426747906725159, -0.9785818524481906],
                    [-0.35712157080642226, 0.911120377575769, -0.20572758986325015]
                ]
        },
        {
            'translation': [455.40107288876885, 273.3589338272866, 5657.814488280711],
            'R':
                [
                    [-0.9269344193869241, -0.3732303525241731, -0.03862235247246717],
                    [-0.04725991098820678, 0.218240494552814, -0.9747500127472326],
                    [0.37223525218497616, -0.901704048173249, -0.21993345934341726]
                ]
        },
        {
            'translation': [-69.271255294384, 422.1843366088847, 4457.893374979773],
            'R':
                [
                    [0.915460708083783, -0.39734606500700814, 0.06362229623477154],
                    [-0.04940628468469528, -0.26789167566119776, -0.9621814117644814],
                    [0.39936288133525055, 0.8776959352388969, -0.26487569589663096]
                ]
        },
    ],
    'S11': [
        {
            'translation': [-234.7208032216618, 464.34018262882194, 5536.652631113797],
            'R':
                [
                    [-0.9059013006181885, 0.4217144115102914, 0.038727105014486805],
                    [0.044493184429779696, 0.1857199061874203, -0.9815948619389944],
                    [-0.4211450938543295, -0.8875049698848251, -0.1870073216538954]
                ]
        },
        {
            'translation': [-11.934348472090557, 449.4165893644565, 5541.113551868937],
            'R':
                [
                    [0.9216646531492915, 0.3879848687925067, -0.0014172943441045224],
                    [0.07721054863099915, -0.18699239961454955, -0.979322405373477],
                    [-0.3802272982247548, 0.9024974149959955, -0.20230080971229314]
                ]
        },
        {
            'translation': [781.127357651581, 235.3131620173424, 5576.37044019807],
            'R':
                [
                    [-0.9063540572469627, -0.42053101768163204, -0.04093880896680188],
                    [-0.0603212197838846, 0.22468715090881142, -0.9725620980997899],
                    [0.4181909532208387, -0.8790161246439863, -0.2290130547809762]
                ]
        },
        {
            'translation': [-155.13650339749012, 422.16256306729633, 4435.416222660868],
            'R':
                [
                    [0.91754082476548, -0.39226322025776267, 0.06517975852741943],
                    [-0.04531905395586976, -0.26600517028098103, -0.9629057236990188],
                    [0.395050652748768, 0.8805514269006645, -0.2618476013752581]
                ]
        },
    ],
}

class Human36mDataset(MocapDataset):
    def __init__(self, path, remove_static_joints=True, universal=False):
        super().__init__(fps=50, skeleton=h36m_skeleton)
        self.universal = universal
        self._cameras = copy.deepcopy(h36m_cameras_extrinsic_params)
        for cameras in self._cameras.values():
            for i, cam in enumerate(cameras):
                cam.update(h36m_cameras_intrinsic_params[i])
                for k, v in cam.items():
                    if k not in ['id', 'res_w', 'res_h']:
                        cam[k] = np.array(v, dtype='float32')

                # Normalize camera frame
                if 'translation' in cam:
                    cam['translation'] = cam['translation'] / 1000  # mm to meters

        camera_info = dict()
        for subject in self._cameras:
            camera_info.setdefault(subject, list())
            for cam in self._cameras[subject]:
                if 'translation' not in cam:
                    continue
                K = np.eye(3, dtype=np.float64)
                K[0, 0] = cam['focal_length'][0]
                K[1, 1] = cam['focal_length'][1]
                K[0, 2] = cam['center'][0]
                K[1, 2] = cam['center'][1]

                R = cam['R']
                dist_coeff = np.concatenate(
                    (cam['radial_distortion'][:2], cam['tangential_distortion'], cam['radial_distortion'][2:])
                ).astype(np.float32).reshape((5,))

                t = np.array(cam['translation'], dtype=np.float64).reshape(3, 1)
                camera_info[subject].append(CameraInfoPacket(P=None, K=K, R=R, t=t,
                                                             res_w=cam['res_w'], res_h=cam['res_h'],
                                                             azimuth=cam['azimuth'], dist_coeff=dist_coeff))
        self.camera_info = camera_info

        # Load serialized dataset
        data = np.load(path, allow_pickle=True)['positions_3d'].item()

        self._data = {}
        for subject, actions in data.items():
            self._data[subject] = {}
            for action_name, positions in actions.items():
                self._data[subject][action_name] = {
                    'positions': positions
                }

        if remove_static_joints:
            if self.universal:
                self.remove_joints([4, 5, 9, 10, 11, 12, 13, 14, 16, 20, 21, 22, 23, 24, 28, 29, 30, 31])
            else:
                # Bring the skeleton to 17 joints instead of the original 32
                self.remove_joints([4, 5, 9, 10, 11, 16, 20, 21, 22, 23, 24, 28, 29, 30, 31])

                # Rewire shoulders to the correct parents
                self._skeleton._parents[11] = 8
                self._skeleton._parents[14] = 8

    def supports_semi_supervised(self):
        return True

    @staticmethod
    def remove_irrelevant_kpts(keypoints, universal=False):
        """

        :param keypoints:
        :return:
        """
        if universal:
            origin_keypoints, origin_keypoints_metadata = keypoints['positions_2d'].item(), keypoints['metadata'].item()
            updated_keypoints, updated_keypoints_metadata = dict(), dict()

            human36m_kpt_index = [0, 1, 2, 3, 4, 5, 6, 10, 11, 12, 13, 14, 15, 16]
            updated_keypoints_metadata['layout_name'] = 'h36m'
            updated_keypoints_metadata['num_joints'] = len(human36m_kpt_index)
            updated_keypoints_metadata['keypoints_symmetry'] = [[4, 5, 6, 8, 9, 10], [1, 2, 3, 11, 12, 13]]

            for subject in origin_keypoints.keys():
                updated_keypoints.setdefault(subject, dict())
                for action in origin_keypoints[subject]:
                    updated_keypoints[subject].setdefault(action, list())
                    for cam_idx, kps in enumerate(origin_keypoints[subject][action]):
                        updated_keypoints[subject][action].append(kps[:, human36m_kpt_index, :])

            return updated_keypoints, updated_keypoints_metadata
        else:
            raise NotImplementedError