# --------------------------------------------------------
# Deep Iterative Matching Network
# Licensed under The Apache-2.0 License [see LICENSE for details]
# Written by Gu Wang
# --------------------------------------------------------
# generate rendered data with the poses generated in last step
from __future__ import print_function, division

import sys
import os

cur_dir = os.path.dirname(os.path.abspath(__file__))
sys.path.insert(1, os.path.join(cur_dir, ".."))
from lib.render_glumpy.render_py import Render_Py
import numpy as np
from lib.utils.mkdir_if_missing import mkdir_if_missing
from lib.pair_matching import RT_transform
import cv2
from tqdm import tqdm

idx2class = {
    1: "ape",
    2: "benchvise",
    # 3: 'bowl',
    4: "camera",
    5: "can",
    6: "cat",
    # 7: 'cup',
    8: "driller",
    9: "duck",
    10: "eggbox",
    11: "glue",
    12: "holepuncher",
    13: "iron",
    14: "lamp",
    15: "phone",
}
# config for render machine
width = 640
height = 480
K = np.array([[572.4114, 0, 325.2611], [0, 573.57043, 242.04899], [0, 0, 1]])
ZNEAR = 0.25
ZFAR = 6.0
depth_factor = 1000

LM6d_root = os.path.join(cur_dir, "../data/LINEMOD_6D/LM6d_converted/LM6d_refine")
observed_set_root = os.path.join(LM6d_root, "image_set/observed")
rendered_pose_path = "%s/LM6d_{}_rendered_pose_{}.txt" % (os.path.join(LM6d_root, "rendered_poses"))

# output_path
rendered_root_dir = os.path.join(LM6d_root, "data/rendered")
pair_set_dir = os.path.join(LM6d_root, "image_set")
mkdir_if_missing(rendered_root_dir)
mkdir_if_missing(pair_set_dir)
print("target path: {}".format(rendered_root_dir))
print("target path: {}".format(pair_set_dir))


def main():
    gen_images = True
    for class_idx, class_name in idx2class.items():
        train_pair = []
        val_pair = []
        print("start ", class_idx, class_name)
        if class_name in ["__back_ground__"]:
            continue
        # uncomment here to only generate data for ape
        # if class_name not in ['ape']:
        #     continue

        if gen_images:
            # init render
            model_dir = os.path.join(LM6d_root, "models", class_name)
            render_machine = Render_Py(model_dir, K, width, height, ZNEAR, ZFAR)

        for set_type in ["all"]:
            with open(os.path.join(observed_set_root, "{}_{}.txt".format(class_name, "all")), "r") as f:
                all_observed_list = [x.strip() for x in f.readlines()]
            # with open(
            #         os.path.join(observed_set_root, '{}_{}.txt'.format(
            #             class_name, 'train')), 'r') as f:
            #     train_observed_list = [x.strip() for x in f.readlines()]
            with open(os.path.join(observed_set_root, "{}_{}.txt".format(class_name, "test")), "r") as f:
                test_observed_list = [x.strip() for x in f.readlines()]

            with open(rendered_pose_path.format(set_type, class_name)) as f:
                str_rendered_pose_list = [x.strip().split(" ") for x in f.readlines()]
            rendered_pose_list = np.array([[float(x) for x in each_pose] for each_pose in str_rendered_pose_list])
            rendered_per_observed = 10
            assert len(rendered_pose_list) == 10 * len(all_observed_list), "{} vs {}".format(
                len(rendered_pose_list), len(all_observed_list)
            )
            for idx, observed_index in enumerate(tqdm(all_observed_list)):
                video_name, observed_prefix = observed_index.split("/")
                rendered_dir = os.path.join(rendered_root_dir, class_name)
                mkdir_if_missing(rendered_dir)
                for inner_idx in range(rendered_per_observed):
                    if gen_images:
                        # if gen_images and observed_index in test_observed_list and inner_idx == 0: # only generate my_val_v{}  # noqa:E501
                        image_file = os.path.join(rendered_dir, "{}_{}-color.png".format(observed_prefix, inner_idx))
                        depth_file = os.path.join(rendered_dir, "{}_{}-depth.png".format(observed_prefix, inner_idx))
                        rendered_idx = idx * rendered_per_observed + inner_idx
                        pose_rendered_q = rendered_pose_list[rendered_idx]

                        rgb_gl, depth_gl = render_machine.render(pose_rendered_q[:4], pose_rendered_q[4:])
                        rgb_gl = rgb_gl.astype("uint8")

                        depth_gl = (depth_gl * depth_factor).astype(np.uint16)

                        cv2.imwrite(image_file, rgb_gl)
                        cv2.imwrite(depth_file, depth_gl)

                        pose_rendered_file = os.path.join(
                            rendered_dir, "{}_{}-pose.txt".format(observed_prefix, inner_idx)
                        )
                        text_file = open(pose_rendered_file, "w")
                        text_file.write("{}\n".format(class_idx))
                        pose_rendered_m = np.zeros((3, 4))
                        pose_rendered_m[:, :3] = RT_transform.quat2mat(pose_rendered_q[:4])
                        pose_rendered_m[:, 3] = pose_rendered_q[4:]
                        pose_ori_m = pose_rendered_m
                        pose_str = "{} {} {} {}\n{} {} {} {}\n{} {} {} {}".format(
                            pose_ori_m[0, 0],
                            pose_ori_m[0, 1],
                            pose_ori_m[0, 2],
                            pose_ori_m[0, 3],
                            pose_ori_m[1, 0],
                            pose_ori_m[1, 1],
                            pose_ori_m[1, 2],
                            pose_ori_m[1, 3],
                            pose_ori_m[2, 0],
                            pose_ori_m[2, 1],
                            pose_ori_m[2, 2],
                            pose_ori_m[2, 3],
                        )
                        text_file.write(pose_str)

                    if observed_index in test_observed_list:
                        if inner_idx == 0:
                            val_pair.append(
                                "{} {}/{}_{}".format(observed_index, class_name, observed_prefix, inner_idx)
                            )
                    else:
                        train_pair.append("{} {}/{}_{}".format(observed_index, class_name, observed_prefix, inner_idx))

            train_pair_set_file = os.path.join(pair_set_dir, "train_{}.txt".format(class_name))
            train_pair = sorted(train_pair)
            with open(train_pair_set_file, "w") as text_file:
                for x in train_pair:
                    text_file.write("{}\n".format(x))

            test_pair_set_file = os.path.join(pair_set_dir, "my_val_{}.txt".format(class_name))
            val_pair = sorted(val_pair)
            with open(test_pair_set_file, "w") as text_file:
                for x in val_pair:
                    text_file.write("{}\n".format(x))
        print(class_name, " done")


def check_observed_rendered():
    from lib.utils.utils import read_img
    import matplotlib.pyplot as plt

    observed_dir = os.path.join(LM6d_root, "data/observed")

    for class_idx, class_name in idx2class.items():
        if class_name != "duck":
            continue
        print(class_name)
        observed_list_path = os.path.join(observed_set_root, "{}_train.txt".format(class_name))
        with open(observed_list_path, "r") as f:
            observed_list = [x.strip() for x in f.readlines()]
        for idx, observed_index in enumerate(observed_list):
            print(observed_index)
            prefix = observed_index.split("/")[1]
            color_observed = read_img(os.path.join(observed_dir, observed_index + "-color.png"), 3)
            color_rendered = read_img(os.path.join(rendered_root_dir, class_name, prefix + "_0-color.png"), 3)
            fig = plt.figure()  # noqa:F401
            plt.axis("off")
            plt.subplot(1, 2, 1)
            plt.imshow(color_observed[:, :, [2, 1, 0]])
            plt.axis("off")

            plt.subplot(1, 2, 2)
            plt.imshow(color_rendered[:, :, [2, 1, 0]])
            plt.axis("off")

            plt.show()


if __name__ == "__main__":
    main()
    # check_observed_rendered()
    print("{} finished".format(__file__))
