import copy
import json
import os

import matplotlib
import numpy as np

import utils_pipeline
from skelda import utils_pose, utils_view
from skelda.writers import json_writer

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

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

whole_body = {
    "foots": False,
    "face": False,
    "hands": False,
}
config = {
    "min_match_score": 0.94,
    "min_group_size": 1,
    "min_bbox_score": 0.3,
    "min_bbox_area": 0.1 * 0.1,
    "batch_poses": True,
    "whole_body": whole_body,
    "take_interval": 1,
}

joint_names_2d = utils_pipeline.get_joint_names(whole_body)
joint_names_3d = list(joint_names_2d)


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


def update_sample(sample, new_dir=""):
    sample = copy.deepcopy(sample)

    # Rename image paths
    sample["imgpaths"] = [
        os.path.join(new_dir, os.path.basename(v)) for v in sample["imgpaths"]
    ]

    # Add placeholders for missing keys
    if not "scene" in sample:
        sample["scene"] = "default"
    if not "id" in sample:
        sample["id"] = "0"
    if not "index" in sample:
        sample["index"] = 0
    for cam in sample["cameras"]:
        if not "type" in cam:
            cam["type"] = "pinhole"

    return sample


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


def main():

    for dirname in sorted(os.listdir(test_img_dir)):
        dirpath = os.path.join(test_img_dir, dirname)

        if not os.path.isdir(dirpath):
            continue

        if (dirname[0] not in ["p", "h", "e", "q"]) or len(dirname) != 2:
            continue

        # Load sample infos
        print("\n" + dirpath)
        with open(os.path.join(dirpath, "sample.json"), "r", encoding="utf-8") as file:
            sample = json.load(file)
        sample = update_sample(sample, dirpath)

        if len(sample["imgpaths"]) == 1:
            # At least two images are required
            continue

        # Save dataset
        labels = [sample]
        tmp_export_dir = "/tmp/rpt/"
        for label in labels:
            if "splits" in label:
                label.pop("splits")
        json_writer.save_dataset(labels, tmp_export_dir)

        # Save config
        config_path = tmp_export_dir + "config.json"
        utils_pipeline.save_json(config, config_path)

        # Call the CPP binary
        os.system("/RapidPoseTriangulation/scripts/test_skelda_dataset.bin")

        # Load the results
        print("Loading exports ...")
        res_path = tmp_export_dir + "results.json"
        results = utils_pipeline.load_json(res_path)
        poses_3d = results["all_poses_3d"][0]
        poses_2d = results["all_poses_2d"][0]
        joint_names_3d = results["joint_names_3d"]

        # Visualize the 2D results
        fig1 = utils_view.draw_many_images(
            sample["imgpaths"], [], [], poses_2d, joint_names_2d, "2D detections"
        )
        fig1.savefig(os.path.join(dirpath, "2d-k.png"), dpi=fig1.dpi)

        # Visualize the 3D results
        print("Detected 3D poses:")
        poses_3d = np.array(poses_3d)
        print(poses_3d.round(3))
        if len(poses_3d) == 0:
            utils_view.show_plots()
            continue
        camparams = sample["cameras"]
        roomparams = {
            "room_size": sample["room_size"],
            "room_center": sample["room_center"],
        }
        poses_2d_proj = []
        for cam in camparams:
            poses_2d_cam, _ = utils_pose.project_poses(poses_3d, cam)
            poses_2d_proj.append(poses_2d_cam)
        fig2 = utils_view.draw_poses3d(poses_3d, joint_names_3d, roomparams, camparams)
        fig3 = utils_view.draw_many_images(
            sample["imgpaths"], [], [], poses_2d_proj, joint_names_3d, "2D projections"
        )
        fig2.savefig(os.path.join(dirpath, "3d-p.png"), dpi=fig2.dpi)
        fig3.savefig(os.path.join(dirpath, "2d-p.png"), dpi=fig3.dpi)
        utils_view.show_plots()


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

if __name__ == "__main__":
    main()