import blenderproc as bproc
"""
  renderBOPdataset
  Общая задача: common pipeline
  Реализуемая функция: создание датасета в формате BOP с заданными параметрами рандомизации
  Используется модуль blenderproc

  19.04.2024 @shalenikol release 0.1
"""
import numpy as np
import argparse
import random
import os
import shutil
import json

VHACD_PATH = "blenderproc_resources/vhacd"
# DIR_BOP = "bop_data"
DIR_MODELS = "models"
FILE_LOG_SCENE = "res.txt"
FILE_RBS_INFO = "rbs_info.json"
FILE_GT_COCO = "scene_gt_coco.json"

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

def _get_path_model(name_model: str) -> str:
    # TODO on name_model find path for mesh (model.fbx)
    # local_path/assets/mesh/
    return os.path.join(rnd_par.output_dir, "assets/mesh/"+name_model+".fbx")
    # , d: dict
    # return d["model"]

def _get_path_object(name_obj: str) -> str:
    # TODO on name_obj find path for scene object (object.fbx)
    return os.path.join(rnd_par.output_dir, "assets/mesh/"+name_obj+".fbx")
    # , d: dict
    # return d["path"]

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

def render() -> int:
    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(VHACD_PATH)

    objs = all_meshs + rnd_par.scene.objs

    log_txt = os.path.join(rnd_par.output_dir, FILE_LOG_SCENE)
    with open(log_txt, "w") as fh:
        for i,o in enumerate(objs):
            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
    ls = []
    for l in rnd_par.scene.light_data:
        light = bproc.types.Light(name=f"l{l['id']}")
        light.set_type(l["type"])
        light.set_location(l["loc_xyz"]) #[5, -5, 5])
        light.set_rotation_euler(l["rot_euler"]) #[-0.063, 0.6177, -0.1985])
        ls += [light]

    # define the camera intrinsics
    bproc.camera.set_intrinsics_from_blender_params(1,
                                                    rnd_par.image_size_wh[0],
                                                    rnd_par.image_size_wh[1],
                                                    lens_unit="FOV")

    # add segmentation masks (per class and per instance)
    bproc.renderer.enable_segmentation_output(map_by=["category_id", "instance", "name"])

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

    res_dir = os.path.join(rnd_par.output_dir, rnd_par.ds_name)
    if os.path.isdir(res_dir):
        shutil.rmtree(res_dir)
    # Цикл рендеринга
    # Do multiple times: Position the shapenet objects using the physics simulator and render X images with random camera poses
    for r in range(rnd_par.n_series):
        # один случайный объект в кадре / все заданные объекты
        random_obj = random.choice(range(rnd_par.scene.n_obj))
        meshs = []
        for i,o in enumerate(all_meshs): #objs
            if rnd_par.single_object and i != random_obj:
                continue
            meshs += [o]
            rnd_mat = rnd_par.scene.obj_data[i]["material_randomization"]
            mats = o.get_materials() #[0]
            for mat in mats:
                val = rnd_mat["specular"]
                mat.set_principled_shader_value("Specular", random.uniform(val[0], val[1]))
                val = rnd_mat["roughness"]
                mat.set_principled_shader_value("Roughness", random.uniform(val[0], val[1]))
                val = rnd_mat["base_color"]
                mat.set_principled_shader_value("Base Color", np.random.uniform(val[0], val[1]))
                val = rnd_mat["metallic"]
                mat.set_principled_shader_value("Metallic", random.uniform(val[0], val[1]))

        # Randomly set the color and energy
        for i,l in enumerate(ls):
            current = rnd_par.scene.light_data[i]
            l.set_color(np.random.uniform(current["color_range_low"], current["color_range_high"]))
            energy = current["energy_range"]
            l.set_energy(random.uniform(energy[0], energy[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(rnd_par.loc_range_low, rnd_par.loc_range_high)) #[-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 + rnd_par.scene.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(log_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(rnd_par.n_cam_pose):
            # Sample location
            location = bproc.sampler.shell(center=rnd_par.center_shell,
                                        radius_min=rnd_par.radius_range[0],
                                        radius_max=rnd_par.radius_range[1],
                                        elevation_min=rnd_par.elevation_range[0],
                                        elevation_max=rnd_par.elevation_range[1])
            # координата, по которой будем сэмплировать положение камеры
            j = random.randint(0, 2)
            # разовый сдвиг по случайной координате
            d = (coord_max[j] - coord_min[j]) / rnd_par.n_sample_on_pose
            if location[j] < 0:
                d = -d
            for _ in range(rnd_par.n_sample_on_pose):
                # 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
        # 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=rnd_par.image_format,
                            append_to_existing_output = (r>0),
                            save_world2cam = False) # world coords are arbitrary in most real BOP datasets
        #                       dataset="robo_ds",

    models_dir = os.path.join(res_dir, DIR_MODELS)
    os.mkdir(models_dir)

    data = []
    for i,objn in enumerate(rnd_par.models.names):
        rec = {}
        rec["id"] = i+1
        rec["name"] = objn
        rec["model"] = os.path.join(DIR_MODELS, os.path.split(rnd_par.models.filenames[i])[1]) # путь относительный
        t = [obj.get_bound_box(local_coords=True).tolist() for obj in all_meshs if obj.get_name() == objn]
        rec["cuboid"] = t[0]
        data.append(rec)
        # ff = os.path.join(args.obj_path, rnd_par.models.filenames[i]) # путь к исходному файлу
        # shutil.copy2(ff, models_dir)
        shutil.copy2(rnd_par.models.filenames[i], models_dir)
        f = (os.path.splitext(rnd_par.models.filenames[i]))[0] + ".mtl" # файл материала
        if os.path.isfile(f):
            shutil.copy2(f, models_dir)
    
    with open(os.path.join(res_dir, FILE_RBS_INFO), "w") as fh:
        json.dump(data, fh, indent=2)

    """
    !!! 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]

    поэтому заменим наименование категории в аннотации
    """
    def change_categories_name(dir: str):
        coco_file = os.path.join(dir,FILE_GT_COCO)
        with open(coco_file, "r") as fh:
            data = json.load(fh)
        cats = data["categories"]

        for i,cat in enumerate(cats):
            cat["name"] = rnd_par.models.names[i] #obj_names[i]

        with open(coco_file, "w") as fh:
            json.dump(data, fh, indent=0)

    def explore(path: str):
        if not os.path.isdir(path):
            return
        folders = [
            os.path.join(path, o)
            for o in os.listdir(path)
            if os.path.isdir(os.path.join(path, o))
        ]
        for path_entry in folders:
            print(path_entry)
            if os.path.isfile(os.path.join(path_entry,FILE_GT_COCO)):
                change_categories_name(path_entry)
            else:
                explore(path_entry)

    if Not_Categories_Name:
        explore(res_dir)
    return 0 # success

def _get_models(par, data) -> int:
    global all_meshs

    par.models = lambda: None
    par.models.n_item = len(data)
    if par.models.n_item == 0:
        return 0 # no models

    # загрузим объекты
    par.models.names = [] #list(map(lambda x: x["name"], data)) # obj_names
    par.models.filenames = [] #list(map(lambda x: x["model"], data)) #obj_filenames
    i = 1
    for f in data:
        nam = f
        par.models.names.append(nam)
        ff = _get_path_model(nam)
        # ff = f["model"] # путь к файлу объекта
        par.models.filenames.append(ff)
        if not os.path.isfile(ff):
            print(f"Error: no such file '{ff}'")
            return -1
        obj = bproc.loader.load_obj(ff)
        all_meshs += obj
        obj[0].set_cp("category_id", i) #f["id"]) # начиная с 1
        i += 1
    return par.models.n_item

def _get_scene(par, data) -> int:
    # load scene
    par.scene = lambda: None
    objs = data["objects"]
    par.scene.n_obj = len(objs)
    if par.scene.n_obj == 0:
        return 0 # empty scene
    lights = data["lights"]
    par.scene.n_light = len(lights)
    if par.scene.n_light == 0:
        return 0 # no lighting
    
    par.scene.objs = []
    par.scene.collision_objects = []
    for f in objs:
        ff = _get_path_object(f["name"]) # f["path"]
        if not os.path.isfile(ff):
            print(f"Error: no such file '{ff}'")
            return -1
        obj = bproc.loader.load_obj(ff)
        obj[0].set_cp("category_id", 999)
        coll = f["collision_shape"]
        if len(coll) > 0:
            obj[0].enable_rigidbody(False, collision_shape=coll)
            par.scene.collision_objects += obj
        par.scene.objs += obj #bproc.loader.load_blend(args.scene, data_blocks=["objects"])

    if not par.scene.collision_objects:
        print("Collision objects not found in the scene")
        return 0
    par.scene.obj_data = objs
    par.scene.light_data = lights
    return par.scene.n_obj

if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    parser.add_argument("--cfg", required=True, help="Json-string with dataset parameters")
    args = parser.parse_args()

    if args.cfg[-5:] == ".json":
        if not os.path.isfile(args.cfg):
            print(f"Error: no such file '{args.cfg}'")
            exit(-1)
        with open(args.cfg, "r") as f:
            j_data = f.read()
    else:
        j_data = args.cfg
    try:
        cfg = json.loads(j_data)
    except json.JSONDecodeError as e:
        print(f"JSon error: {e}")
        exit(-2)

    ds_cfg = cfg["formBuilder"]["output"] # dataset config
    generation = ds_cfg["generation"]
    cam_pos = ds_cfg["camera_position"]
    models_randomization = ds_cfg["models_randomization"]

    rnd_par = lambda: None
    rnd_par.single_object = True
    rnd_par.ds_name = cfg["name"]
    rnd_par.output_dir = cfg["local_path"]
    rnd_par.dataset_objs = cfg["dataSetObjects"]
    rnd_par.n_cam_pose = generation["n_cam_pose"]
    rnd_par.n_sample_on_pose = generation["n_sample_on_pose"]
    rnd_par.n_series = generation["n_series"]
    rnd_par.image_format = generation["image_format"]
    rnd_par.image_size_wh = generation["image_size_wh"]
    rnd_par.center_shell = cam_pos["center_shell"]
    rnd_par.radius_range = cam_pos["radius_range"]
    rnd_par.elevation_range = cam_pos["elevation_range"]
    rnd_par.loc_range_low = models_randomization["loc_range_low"]
    rnd_par.loc_range_high = models_randomization["loc_range_high"]

    if not os.path.isdir(rnd_par.output_dir):
        # os.mkdir(rnd_par.output_dir)
        print(f"Error: invalid path '{rnd_par.output_dir}'")
        exit(-3)

    bproc.init()
    
    all_meshs = []
    ret = _get_models(rnd_par, rnd_par.dataset_objs)
    if ret <= 0:
        print("Error: no models in config")
        exit(-4)
    if _get_scene(rnd_par, ds_cfg["scene"]) == 0:
        print("Error: empty scene in config")
        exit(-5)
    exit(render())