import json
import os

import cv2
import matplotlib.pyplot as plt
import numpy as np

from skelda import utils_pose, utils_view

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

filepath = os.path.dirname(os.path.realpath(__file__)) + "/"

core_triangs = [
    [
        [0.287, -0.282, 1.264, 1.000],
        [0.504, -0.052, 1.272, 1.000],
        [0.276, -0.160, 0.764, 1.000],
        [0.443, -0.099, 0.768, 1.000],
        [0.258, -0.313, 0.999, 1.000],
        [0.513, -0.009, 1.008, 1.000],
        [0.204, -0.126, 0.439, 1.000],
        [0.422, -0.132, 0.436, 1.000],
        [0.195, -0.265, 0.807, 1.000],
        [0.415, 0.039, 0.823, 1.000],
        [0.113, -0.103, 0.096, 1.000],
        [0.389, -0.175, 0.097, 1.000],
    ],
    [
        [0.322, -0.192, 1.349, 1.000],
        [0.268, -0.594, 1.336, 1.000],
        [0.272, -0.100, 0.882, 1.000],
        [0.281, -0.379, 0.870, 1.000],
        [0.336, -0.104, 1.124, 1.000],
        [0.249, -0.578, 1.089, 1.000],
        [0.229, 0.009, 0.571, 1.000],
        [0.269, -0.345, 0.553, 1.000],
        [0.289, -0.016, 0.951, 1.000],
        [0.216, -0.327, 0.908, 1.000],
        [0.188, 0.128, 0.268, 1.000],
        [0.267, -0.273, 0.243, 1.000],
    ],
    [
        [0.865, 1.058, 1.613, 1.000],
        [0.862, 0.870, 1.604, 1.000],
        [0.927, 1.562, 1.491, 1.000],
        [0.954, 1.505, 1.486, 1.000],
        [0.908, 1.309, 1.542, 1.000],
        [0.905, 1.170, 1.525, 1.000],
        [0.968, 1.911, 1.454, 1.000],
        [1.019, 1.919, 1.457, 1.000],
        [0.921, 1.542, 1.514, 1.000],
        [0.931, 1.539, 1.506, 1.000],
        [1.008, 2.230, 1.455, 1.000],
        [1.071, 2.271, 1.460, 1.000],
    ],
    [
        [-0.260, 0.789, 1.316, 1.000],
        [0.039, 1.073, 1.322, 1.000],
        [-0.236, 0.798, 0.741, 1.000],
        [-0.048, 0.952, 0.759, 1.000],
        [-0.315, 0.734, 0.995, 1.000],
        [0.080, 1.026, 1.046, 1.000],
        [-0.291, 0.721, 0.339, 1.000],
        [-0.101, 0.887, 0.366, 1.000],
        [-0.300, 0.600, 0.742, 1.000],
        [0.066, 0.768, 0.897, 1.000],
        [-0.381, 0.685, -0.113, 1.000],
        [-0.169, 0.775, -0.040, 1.000],
    ],
    [
        [-0.199, 0.854, 1.414, 1.000],
        [-0.401, 0.566, 1.409, 1.000],
        [-0.242, 0.818, 0.870, 1.000],
        [-0.343, 0.654, 0.856, 1.000],
        [-0.176, 0.903, 1.140, 1.000],
        [-0.398, 0.480, 1.132, 1.000],
        [-0.245, 0.812, 0.492, 1.000],
        [-0.380, 0.642, 0.471, 1.000],
        [-0.145, 0.817, 0.912, 1.000],
        [-0.251, 0.396, 0.973, 1.000],
        [-0.255, 0.879, 0.107, 1.000],
        [-0.383, 0.633, 0.116, 1.000],
    ],
    [
        [0.641, 1.796, 1.681, 1.000],
        [0.603, 1.719, 1.680, 1.000],
        [0.711, 2.000, 1.518, 1.000],
        [0.706, 1.970, 1.515, 1.000],
        [0.689, 1.920, 1.588, 1.000],
        [0.651, 1.784, 1.585, 1.000],
        [0.786, 2.190, 1.448, 1.000],
        [0.780, 2.167, 1.444, 1.000],
        [0.747, 1.994, 1.531, 1.000],
        [0.720, 1.783, 1.546, 1.000],
        [0.868, 2.432, 1.427, 1.000],
        [0.849, 2.341, 1.410, 1.000],
    ],
]

core_joints = [
    "shoulder_left",
    "shoulder_right",
    "hip_left",
    "hip_right",
    "elbow_left",
    "elbow_right",
    "knee_left",
    "knee_right",
    "wrist_left",
    "wrist_right",
    "ankle_left",
    "ankle_right",
]

poses_2d = [
    [
        [
            [383.443, 144.912, 0.923],
            [382.629, 135.143, 0.83],
            [374.488, 134.329, 1.0],
            [349.251, 136.771, 0.478],
            [343.552, 139.213, 1.0],
            [356.578, 201.899, 0.73],
            [323.2, 201.899, 0.825],
            [357.392, 282.494, 0.663],
            [324.014, 289.821, 0.854],
            [378.558, 339.481, 0.621],
            [355.764, 356.577, 0.821],
            [370.417, 357.391, 0.714],
            [332.155, 359.834, 0.71],
            [391.584, 452.641, 0.768],
            [331.341, 458.34, 0.789],
            [414.379, 547.076, 0.864],
            [332.969, 550.333, 0.9],
            [351.286, 358.613, 0.71],
            [339.889, 201.899, 0.73],
            [346.402, 137.992, 0.478],
        ],
        [
            [640.948, 116.443, 0.908],
            [650.057, 100.249, 0.788],
            [642.972, 100.249, 0.681],
            [682.445, 103.285, 0.862],
            [684.469, 100.249, 0.518],
            [707.748, 181.219, 0.836],
            [693.578, 180.207, 0.705],
            [702.688, 290.528, 0.867],
            [664.227, 276.358, 0.786],
            [662.202, 375.547, 0.847],
            [605.523, 319.88, 0.881],
            [692.566, 373.522, 0.758],
            [679.409, 372.51, 0.739],
            [679.409, 500.038, 0.844],
            [672.324, 494.978, 0.837],
            [679.409, 635.663, 0.913],
            [659.166, 599.226, 0.894],
            [685.987, 373.016, 0.739],
            [700.663, 180.713, 0.705],
            [683.457, 101.767, 0.518],
        ],
    ],
    [
        [
            [495.125, 304.671, 0.581],
            [492.338, 301.885, 0.462],
            [502.091, 301.885, 0.295],
            [495.125, 308.851, 0.92],
            [528.562, 306.064, 0.754],
            [477.013, 359.703, 0.822],
            [557.819, 355.523, 0.855],
            [466.564, 431.452, 0.855],
            [565.481, 425.879, 0.836],
            [458.902, 480.911, 0.85],
            [544.583, 464.889, 0.596],
            [491.642, 490.663, 0.741],
            [539.707, 492.056, 0.746],
            [480.496, 569.379, 0.779],
            [531.348, 577.041, 0.784],
            [464.475, 646.005, 0.872],
            [518.809, 661.33, 0.913],
            [515.675, 491.36, 0.741],
            [517.416, 357.613, 0.822],
            [511.843, 307.458, 0.754],
        ],
        [
            [472.982, 273.983, 0.911],
            [477.875, 266.645, 0.848],
            [464.421, 266.645, 0.896],
            [483.378, 268.48, 0.599],
            [448.521, 268.48, 0.88],
            [493.163, 308.841, 0.753],
            [425.894, 311.899, 0.837],
            [502.336, 363.268, 0.625],
            [417.944, 368.16, 0.847],
            [499.278, 407.91, 0.495],
            [438.736, 410.357, 0.844],
            [484.602, 426.868, 0.682],
            [448.521, 427.48, 0.681],
            [485.825, 504.534, 0.745],
            [441.794, 505.757, 0.796],
            [489.494, 571.803, 0.688],
            [442.405, 577.918, 0.867],
            [466.561, 427.174, 0.681],
            [459.528, 310.37, 0.753],
            [465.95, 268.48, 0.599],
        ],
        [
            [702.349, 208.747, 0.215],
            [705.862, 207.944, 0.212],
            [700.341, 207.944, 0.209],
            [708.472, 196.399, 0.182],
            [699.538, 196.299, 0.193],
            [708.071, 196.7, 0.194],
            [696.927, 196.098, 0.22],
            [709.175, 206.137, 0.191],
            [696.526, 206.94, 0.188],
            [707.368, 210.052, 0.128],
            [699.738, 209.751, 0.145],
            [704.658, 215.172, 0.172],
            [701.445, 215.273, 0.168],
            [705.159, 224.007, 0.179],
            [703.654, 225.211, 0.185],
            [705.059, 225.914, 0.23],
            [704.457, 230.532, 0.241],
            [703.051, 215.223, 0.168],
            [702.499, 196.399, 0.194],
            [704.005, 196.349, 0.182],
        ],
    ],
]

joints_2d = [
    "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",
]


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


def main():

    with open(os.path.join(filepath, "sample.json"), "r", encoding="utf-8") as file:
        sample = json.load(file)

    camparams = sample["cameras"]
    roomparams = {
        "room_size": sample["room_size"],
        "room_center": sample["room_center"],
    }

    fig2 = utils_view.draw_poses3d(core_triangs, core_joints, roomparams, camparams)
    fig2.axes[0].view_init(elev=30, azim=0)
    fig2.savefig(os.path.join(filepath, "core-triangs.png"), dpi=fig2.dpi)

    core_projections = []
    for i in range(len(camparams)):
        b2d, _ = utils_pose.project_poses(np.array(core_triangs), camparams[i])
        core_projections.append(b2d)

    img_size = [900, 900]
    scale = 0.66
    fig_size = 2
    plotsize = (35, 30)
    fig, axs = plt.subplots(fig_size, fig_size, figsize=plotsize)
    fig.suptitle("core reprojections", fontsize=20)
    fig.tight_layout(rect=[0, 0, 1, 0.97])

    num_persons = max((len(b2d) for b2d in core_projections))
    colors = plt.cm.hsv(np.linspace(0, 1, num_persons, endpoint=False)).tolist()
    colors = [[int(c[0] * 255), int(c[1] * 255), int(c[2] * 255)] for c in colors]

    for i in range(len(camparams)):
        img = np.ones((img_size[0], img_size[1], 3), dtype=np.uint8) * 255

        for j in range(len(core_projections[i])):
            color = colors[j]
            body = np.array(core_projections[i][j])
            img = utils_view.draw_body_in_image(img, body, core_joints, color)

        for k in range(len(poses_2d[i])):
            body = np.array(poses_2d[i][k])
            cjids = [joints_2d.index(j) for j in core_joints]
            body = body[cjids]
            img = utils_view.draw_body_in_image(
                img, body, core_joints, [0, 0, 0], thickness=2
            )

        img = cv2.resize(img, (int(img.shape[1] * scale), int(img.shape[0] * scale)))
        title = str(i)

        x, y = divmod(i, fig_size)
        axs[x][y].imshow(img)
        axs[x][y].set_title(title)

    # Delete empty plots
    for i in range(2, fig_size**2):
        x, y = divmod(i, fig_size)
        fig.delaxes(axs[x][y])

    fig.savefig(os.path.join(filepath, "core-reprojections.png"), dpi=fig.dpi)


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

if __name__ == "__main__":
    main()