import blenderproc as bproc
"""
  objs2BOPdataset
  Общая задача: распознавание 6D позы объекта (6D pose estimation)
  Реализуемая функция: создание датасета в формате BOP для серии заданных объектов (*.obj) в заданной сцене (*.blend)
  Используется модуль blenderproc

  29.08.2023 @shalenikol release 0.1
  12.10.2023 @shalenikol release 0.2
"""
import sys
import numpy as np
import argparse
import random
import os
import shutil
import json

Not_Categories_Name = True # наименование категории в аннотации отсутствует

def convert2relative(height, width, bbox):
    """
    YOLO format use relative coordinates for annotation
    """
    x, y, w, h = bbox
    x += w/2
    y += h/2
    return x/width, y/height, w/width, h/height

parser = argparse.ArgumentParser()
parser.add_argument('scene', nargs='?', default="resources/sklad.blend", help="Path to the scene object.")
parser.add_argument('obj_path', nargs='?', default="resources/in_obj", help="Path to the object files.")
parser.add_argument('output_dir', nargs='?', default="output", help="Path to where the final files, will be saved")
parser.add_argument('vhacd_path', nargs='?', default="blenderproc_resources/vhacd", help="The directory in which vhacd should be installed or is already installed.")
parser.add_argument('-single_object', nargs='?', type= bool, default=True, help="One object per frame.")
parser.add_argument('--imgs', default=2, type=int, help="The number of times the objects should be rendered.")
args = parser.parse_args()

if not os.path.isdir(args.obj_path):
    print(f"{args.obj_path} : no object directory")
    sys.exit()

if not os.path.isdir(args.output_dir):
    os.mkdir(args.output_dir)

single_object = args.single_object

bproc.init()

# ? загрузим свет из сцены
#cam = bproc.loader.load_blend(args.scene, data_blocks=["cameras"])
#lights = bproc.loader.load_blend(args.scene, data_blocks=["lights"])

# загрузим объекты
list_files = os.listdir(args.obj_path)
obj_names = []
obj_filenames = []
all_meshs = []
nObj = 0
for f in list_files:
    if (os.path.splitext(f))[1] == ".obj":
        f = os.path.join(args.obj_path, f) # путь к файлу объекта
        if os.path.isfile(f):
            obj = bproc.loader.load_obj(f)
            all_meshs += obj
            obj_names += [obj[0].get_name()]
            obj_filenames += [f]
            nObj += 1

if nObj == 0:
    print("Objects not found")
    sys.exit()

for i,obj in enumerate(all_meshs):
    #print(f"{i} *** {obj}")
    obj.set_cp("category_id", i+1)

# загрузим сцену
scene = bproc.loader.load_blend(args.scene, data_blocks=["objects"])

# найдём объекты коллизии (пол и т.д.)
obj_type = ["floor", "obj"]
collision_objects = []
#floor = None
for o in scene:
    o.set_cp("category_id", 999)
    s = o.get_name()
    for type in obj_type:
        if s.find(type) >= 0:
            collision_objects += [o]
            o.enable_rigidbody(False, collision_shape='BOX')
if not collision_objects:
    print("Collision objects not found in the scene")
    sys.exit()

#floor.enable_rigidbody(False, collision_shape='BOX')

for obj in all_meshs:
    # Make the object actively participate in the physics simulation
    obj.enable_rigidbody(active=True, collision_shape="COMPOUND")
    # Also use convex decomposition as collision shapes
    obj.build_convex_decomposition_collision_shape(args.vhacd_path)

objs = all_meshs + scene

with open(os.path.join(args.output_dir,"res.txt"), "w") as fh:
#    fh.write(str(type(scene[0]))+"\n")
    i = 0
    for o in objs:
        i += 1
        loc = o.get_location()
        euler = o.get_rotation_euler()
        fh.write(f"{i} : {o.get_name()}  {loc} {euler}  category_id = {o.get_cp('category_id')}\n")

# define a light and set its location and energy level
light = bproc.types.Light()
light.set_type("POINT")
light.set_location([5, -5, 5])
#light.set_energy(900)
#light.set_color([0.7, 0.7, 0.7])

light1 = bproc.types.Light(name="light1")
light1.set_type("SUN")
light1.set_location([0, 0, 0])
light1.set_rotation_euler([-0.063, 0.6177, -0.1985])
#light1.set_energy(7)
light1.set_color([1, 1, 1])

# define the camera intrinsics
bproc.camera.set_intrinsics_from_blender_params(1, 640, 480, lens_unit="FOV")

# add segmentation masks (per class and per instance)
bproc.renderer.enable_segmentation_output(map_by=["category_id", "instance", "name"])
#bproc.renderer.enable_segmentation_output(map_by=["category_id", "instance", "name", "bop_dataset_name"],
#                                          default_values={"category_id": 0, "bop_dataset_name": None})

# activate depth rendering
bproc.renderer.enable_depth_output(activate_antialiasing=False)

res_dir = os.path.join(args.output_dir, "bop_data")
if os.path.isdir(res_dir):
    shutil.rmtree(res_dir)
# Цикл рендеринга
n_cam_location = 3 #5 # количество случайных локаций камеры
n_cam_poses = 3 #3    # количество сэмплов для каждой локации камеры
# Do multiple times: Position the shapenet objects using the physics simulator and render X images with random camera poses
for r in range(args.imgs):
    # один случайный объект в кадре / все заданные объекты
    meshs = [random.choice(all_meshs)] if single_object else all_meshs[:]

    # Randomly set the color and energy
    light.set_color(np.random.uniform([0.5, 0.5, 0.5], [1, 1, 1]))
    light.set_energy(random.uniform(500, 1000))
    light1.set_energy(random.uniform(3, 11))

    for i,o in enumerate(meshs): #objs
        mat = o.get_materials()[0]
        mat.set_principled_shader_value("Specular", random.uniform(0, 1))
        mat.set_principled_shader_value("Roughness", random.uniform(0, 1))
        mat.set_principled_shader_value("Base Color", np.random.uniform([0, 0, 0, 1], [1, 1, 1, 1]))
        mat.set_principled_shader_value("Metallic", random.uniform(0, 1))

    # Clear all key frames from the previous run
    bproc.utility.reset_keyframes()

    # Define a function that samples 6-DoF poses
    def sample_pose(obj: bproc.types.MeshObject):
        obj.set_location(np.random.uniform([-1, -1.5, 0.2], [1, 2, 1.2])) #[-1, -1, 0], [1, 1, 2]))
        obj.set_rotation_euler(bproc.sampler.uniformSO3())

    # Sample the poses of all shapenet objects above the ground without any collisions in-between
    #bproc.object.sample_poses(meshs, objects_to_check_collisions = meshs + [floor], sample_pose_func = sample_pose)
    bproc.object.sample_poses(meshs, objects_to_check_collisions = meshs + collision_objects, sample_pose_func = sample_pose)

    # Run the simulation and fix the poses of the shapenet objects at the end
    bproc.object.simulate_physics_and_fix_final_poses(min_simulation_time=4, max_simulation_time=20, check_object_interval=1)

    # Find point of interest, all cam poses should look towards it
    poi = bproc.object.compute_poi(meshs)

    coord_max = [0.1, 0.1, 0.1]
    coord_min = [0., 0., 0.]

    with open(os.path.join(args.output_dir,"res.txt"), "a") as fh:
        fh.write("*****************\n")
        fh.write(f"{r}) poi = {poi}\n")
        i = 0
        for o in meshs:
            i += 1
            loc = o.get_location()
            euler = o.get_rotation_euler()
            fh.write(f"   {i} : {o.get_name()}  {loc} {euler}\n")
            for j in range(3):
                if loc[j] < coord_min[j]:
                    coord_min[j] = loc[j]
                if loc[j] > coord_max[j]:
                    coord_max[j] = loc[j]

    # Sample up to X camera poses
    #an = np.random.uniform(0.78, 1.2) #1. #0.35
    for i in range(n_cam_location):
        # Sample location
        location = bproc.sampler.shell(center=[0, 0, 0],
                                       radius_min=1.1,
                                       radius_max=2.2,
                                       elevation_min=5,
                                       elevation_max=89)
        # координата, по которой будем сэмплировать положение камеры
        j = random.randint(0, 2)
        # разовый сдвиг по случайной координате
        d = (coord_max[j] - coord_min[j]) / n_cam_poses
        if location[j] < 0:
          d = -d
        for k in range(n_cam_poses):
          # Compute rotation based on vector going from location towards poi
          rotation_matrix = bproc.camera.rotation_from_forward_vec(poi - location, inplane_rot=np.random.uniform(-0.7854, 0.7854))
          # Add homog cam pose based on location an rotation
          cam2world_matrix = bproc.math.build_transformation_mat(location, rotation_matrix)
          bproc.camera.add_camera_pose(cam2world_matrix)
          location[j] -= d
        #world_matrix = bproc.math.build_transformation_mat([2.3, -0.4, 0.66], [1.396, 0., an])
        #bproc.camera.add_camera_pose(world_matrix)
        #an += 0.2

    # render the whole pipeline
    data = bproc.renderer.render()
    # Write data to bop format
    bproc.writer.write_bop(res_dir,
                           target_objects = all_meshs, # Optional[List[MeshObject]] = None
                           depths = data["depth"],
                           depth_scale = 1.0,
                           colors = data["colors"],
                           color_file_format='JPEG',
                           append_to_existing_output = (r>0),
                           save_world2cam = False) # world coords are arbitrary in most real BOP datasets
    #                       dataset="robo_ds",
"""
   !!! categories -> name берётся из category_id !!!
       см.ниже
blenderproc.python.writer : BopWriterUtility.py
  class _BopWriterUtility
    def calc_gt_coco
      ...
      CATEGORIES = [{'id': obj.get_cp('category_id'), 'name': str(obj.get_cp('category_id')), 'supercategory':
                          dataset_name} for obj in dataset_objects]

поэтому заменим наименование категории в аннотации
"""
if Not_Categories_Name:
    coco_file = os.path.join(res_dir,"train_pbr/000000/scene_gt_coco.json")
    with open(coco_file, "r") as fh:
        data = json.load(fh)
    cats = data["categories"]
    #print(f"type(cat) = {type(cat)}  cat : {cat}")
    i = 0
    for cat in cats:
        cat["name"] = obj_names[i]
        i += 1
        #print(cat)
    with open(coco_file, "w") as fh:
        json.dump(data, fh, indent=0)