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

  02.05.2024 @shalenikol release 0.1
  02.07.2024 @shalenikol release 0.2
  28.10.2024 @shalenikol release 0.3
"""
import numpy as np
import argparse
import random
import os
import shutil
import json
from pathlib import Path

import bpy

VHACD_PATH = "blenderproc_resources/vhacd"
DIR_MODELS = "models"
DIR_MESH = "assets/libs/objects/"
FILE_LOG_SCENE = "res.txt"
FILE_RBS_INFO = "rbs_info.json"
FILE_GT_COCO = "scene_gt_coco.json"
EXT_MODELS = ".fbx"
TEXTURE_TMPL = "*.jpg"

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

def _get_list_texture(rel_path: str) -> list:
    # local_path/texture/
    loc = os.path.dirname(os.path.dirname(rnd_par.output_dir))
    path = os.path.join(loc, rel_path)
    return list(Path(path).absolute().rglob(TEXTURE_TMPL))

# def _get_path_model(name_model: str) -> str:
#     # TODO on name_model find path for mesh (model.fbx)
#     # local_path/assets/libs/objects   # assets/mesh/
#     loc = os.path.dirname(os.path.dirname(rnd_par.output_dir))
#     return os.path.join(loc, DIR_MESH + name_model + EXT_MODELS)

def _get_path_object(name_obj: str) -> str:
    # TODO on name_obj find path for scene object (object.fbx)
    # loc = os.path.dirname(os.path.dirname(rnd_par.output_dir))
    # return os.path.join(loc, DIR_MESH + name_obj + EXT_MODELS)
    return os.path.join(rnd_par.details_dir, name_obj + EXT_MODELS)

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:
    i = 0
    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)
        i += 1
        # print(f"{i} : {obj.get_name()}")

    objs = all_meshs + rnd_par.scene.objs

    log_txt = os.path.join(os.path.dirname(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)
    res_dir = rnd_par.output_dir
    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):
        print(f"********** Series : {r+1}")
        is_texture = True if "texture_path" in rnd_par.models_randomization else False
        if is_texture:
            val = rnd_par.models_randomization["texture_path"]
            l_texture = _get_list_texture(val)
            image = bpy.data.images.load(filepath=str(l_texture[r % len(l_texture)]))
        # один случайный объект в кадре / все заданные объекты
        random_obj = random.choice(range(rnd_par.models.n_item))
        meshs = []
        for i,o in enumerate(all_meshs): # активные модели
            if rnd_par.single_object and i != random_obj:
                continue
            meshs += [o]
            if is_texture:
                mats = o.get_materials()
                for mat in mats:
                    # image = bpy.data.images.load(filepath=str(random.choice(l_texture)))
                    mat.set_principled_shader_value("Base Color", image)

        for i,o in enumerate(rnd_par.scene.objs): # объекты сцены
            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["metallic"]
                mat.set_principled_shader_value("Metallic", random.uniform(val[0], val[1]))
                if "texture_path" in rnd_mat: # путь к текстурам (*.jpg)
                    val = rnd_mat["texture_path"]
                    val = _get_list_texture(val)
                    image = bpy.data.images.load(filepath=str(random.choice(val)))
                    mat.set_principled_shader_value("Base Color", image)
                else:
                    val = rnd_mat["base_color"]
                    mat.set_principled_shader_value("Base Color", np.random.uniform(val[0], val[1]))
                # mat.set_principled_shader_value("Base Color", image)

        # 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)
        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=1)

    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 = []     # obj_names
    par.models.filenames = [] # obj_filenames
    i = 1
    for f in data:
        nam = f["name"]
        par.models.names.append(nam)
        ff = f["fbx"] # _get_path_model(nam)
        par.models.filenames.append(ff)
        if not os.path.isfile(ff):
            print(f"Error: no such file '{ff}'")
            return -1
        # !!! dir with meshs
        par.details_dir = os.path.split(ff)[0]

        obj = bproc.loader.load_obj(ff)
        all_meshs += obj
        obj[0].set_cp("category_id", i) # начиная с 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
    if len(rnd_par.details_dir) == 0:
        return 0 # no path to details

    par.scene.objs = []
    par.scene.collision_objects = []
    for f in objs:
        ff = _get_path_object(f["name"])
        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

    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("--form", required=True, help="Json-string with dataset parameters")
    parser.add_argument("--path", required=True, help="Output path")
    args = parser.parse_args()

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

    # output_dir = args.path

    ds_cfg = cfg["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.output_dir = args.path # cfg["local_path"]

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

    rnd_par.single_object = False # True
    rnd_par.details_dir = ""
    # rnd_par.ds_name = os.path.split(rnd_par.output_dir)[1] # cfg["name"]
    rnd_par.dataset_objs = ds_cfg["datasetObjects"]["details"] # ["knight"]
    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.models_randomization = models_randomization
    rnd_par.loc_range_low = models_randomization["loc_range_low"]
    rnd_par.loc_range_high = models_randomization["loc_range_high"]

    bproc.init()
    
    all_meshs = []
    if _get_models(rnd_par, rnd_par.dataset_objs) <= 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())