import random
import glob
import json
import os
import sys
import shutil
from collections import OrderedDict

import bpy
from mathutils import Vector

print('#' * 30)

WORKING_DIR = r"D:\Mingming\synthe_dripe"
os.chdir(WORKING_DIR)
sys.path.append(WORKING_DIR)

from scripts import random_pose
from scripts import utils
from scripts import camera_proj
from scripts import human

import importlib

importlib.reload(random_pose)
importlib.reload(utils)
importlib.reload(camera_proj)
importlib.reload(human)

from scripts.human import HumanLoader


def abs_path(rel_path):
    return os.path.join(os.path.dirname(os.path.dirname(__file__)), rel_path)


random.seed()

C = bpy.context
D = bpy.data

CONTINUE = False

# Clean scene
try:
    bpy.ops.object.mode_set(mode='OBJECT')
    utils.unselect_all()
except RuntimeError:
    pass

for col in D.collections:
    D.collections.remove(col)

for bpy_data_iter in (
        D.objects,
        D.meshes,
        D.lights,
        D.cameras,
        D.armatures,
        D.images,
):
    for id_data in bpy_data_iter:
        bpy_data_iter.remove(id_data)

for ob in D.objects:
    C.scene.collection.objects.unlink(ob)

# Import car
car_collection = D.collections.new("Cars")
C.scene.collection.children.link(car_collection)

cars = []

for src_path, name in {
    r"car_models\suv_car\car.blend": 'SUV',
    # r"car_models\red_car\red.blend": 'Red',
    # r"car_models\pickup_car\pickup.blend": 'PickUp',
    # r"car_models\family_car\family_car.blend": 'Family',
    # r"car_models\coupe_car\coupe_car.blend": 'Coupe',
    # r"car_models\truck\truck_open.blend": 'Truck',

}.items():
    with D.libraries.load(abs_path(src_path)) as (data_from, data_to):
        data_to.objects = data_from.objects

    for obj in data_to.objects:
        car_collection.objects.link(obj)

    D.objects['Car'].name = name
    cars.append(D.objects[name])
car_picker = utils.Randomizer(cars)

# import humans
human_loader = HumanLoader('mh_models/exports')

# Creation scene
# add camera
C.scene.render.engine = 'BLENDER_EEVEE'

camera_data = D.cameras.new(name='Camera')
camera_data.type = 'PERSP'
camera_data.lens_unit = 'FOV'
camera_data.angle = utils.r(68)

C.scene.render.resolution_x = 640
C.scene.render.resolution_y = 480

camera_object = D.objects.new('Camera', camera_data)
C.scene.collection.objects.link(camera_object)

camera_object.location = [-.97, -0.11, 0.68]
camera_object.rotation_euler = utils.r([72, 8, -82])

# set background
back_folder = abs_path("backgrounds")
back_imgs = {}
for key in ['night', 'day']:
    list_imgs = glob.glob(os.path.join(back_folder, f'{key}_*'))
    back_imgs[key] = []
    for img in list_imgs:
        img_name = os.path.basename(img)
        bpy.ops.image.open(filepath=img, directory=back_folder,
                           files=[{"name": img_name}], relative_path=False, show_multiview=False)
        back_imgs[key].append(img_name)

# Create image holder
bpy.ops.import_image.to_plane(
    directory=back_folder,
    files=[{"name": "default_green.png"}],
    shader='SHADELESS',
    use_transparency=False,
    offset=False,
    height=round(10 / 2.25, 1),
    align_axis="X-"
)

image_holder = C.active_object
image_holder.name = 'Image_holder'
image_holder.location = (4, 1.5, 1.3)
image_holder.rotation_euler.z = utils.r(-100)
image_holder.active_material.shadow_method = 'NONE'

# add light
sun_collection = D.collections.new("Sun")
C.scene.collection.children.link(sun_collection)

light_params = {
    'day': {
        'energy_bounds': (10, 35),
        'back_color_bounds': (0.4, 0.65),
        'sun_color': (1, 1, 1)
    },
    'night': {
        'energy_bounds': (5, 15),
        'back_color_bounds': (0.05, 0.15),
        'sun_color': (1, 0.5, 0.2)
    }

}

light_data = D.lights.new(name="Sun", type='SUN')
light_data.energy = 20
light_object = D.objects.new(name="Sun", object_data=light_data)

sun_collection.objects.link(light_object)
light_object.location = (5, 0, 3)

# Sun position
C.scene.sun_pos_properties.usage_mode = 'NORMAL'
C.scene.sun_pos_properties.sun_object = light_object
C.scene.sun_pos_properties.object_collection = sun_collection
C.scene.sun_pos_properties.object_collection_type = 'DIURNAL'
C.scene.sun_pos_properties.co_parser = "41°22′14″N 2°09′00″E"
C.scene.sun_pos_properties.sun_distance = 3
C.scene.sun_pos_properties.use_day_of_year = True
C.scene.sun_pos_properties.year = 2022
C.scene.sun_pos_properties.day_of_year = 182

fp = abs_path(r"output")
fp_img = os.path.join(fp, 'images')
fp_ann_2D = os.path.join(fp, 'annots_2D')
fp_ann_3D = os.path.join(fp, 'annots_3D')

info_path = os.path.join(fp, 'infos.json')
scenes_ids = OrderedDict()
if not CONTINUE or not os.path.isfile(info_path):
    if os.path.isdir(fp):
        shutil.rmtree(fp)
    os.mkdir(fp)

    os.mkdir(fp_img)
    os.mkdir(fp_ann_2D)
    os.mkdir(fp_ann_3D)

frame_rate = 25
nb_scene = 1
nb_pose = 10
human_loader.max_len = min(human_loader.max_len, nb_scene)
ratio_conf_man = int(nb_scene / len(human_loader.human_paths))
if CONTINUE:
    with open(info_path) as f_info:
        scene_ids = json.load(f_info, object_pairs_hook=OrderedDict)['id_max_scenes']

    human_loader.human_paths = [hp for hp in human_loader.human_paths if hp not in scene_ids]

man = None
for sc in range(nb_scene):
    # Random car
    car = car_picker()
    car_targets = {side: [ch for ch in car.children if f'Target_{side.upper()}' in ch.name] for side in 'lr'}
    nb_targets = sum([len(v) for v in car_targets.values()])
    # Random personne
    if ratio_conf_man < 1:
        if not sc % 10:
            human_loader.load_next()
        man = human_loader(car=car)
    else:
        if not sc % ratio_conf_man:
            man = human_loader.next(car=car)
        else:
            human.set_bounds(man, car)

    human_path = human_loader.paths[man]
    scenes_ids.setdefault(human_path, -1)
    scenes_ids[human_path] += 1

    # Random time
    C.scene.sun_pos_properties.north_offset = utils.r(random.randint(-179, 180))
    time_day = random.randint(0, 23)
    if 6 <= time_day <= 21:
        day_night = 'day'
        C.scene.sun_pos_properties.time = time_day
    else:
        day_night = 'night'
        C.scene.sun_pos_properties.time = (time_day + 12) % 24

    light_param = light_params[day_night]
    light_data.energy = random.randint(*light_param['energy_bounds'])
    light_data.color = light_param['sun_color']
    back_val = random.uniform(*light_param['back_color_bounds'])
    bpy.data.worlds["World"].node_tree.nodes["Background.001"].inputs[0].default_value = \
        (back_val, back_val, back_val, 1)

    # Random background
    back_img = random.choice(back_imgs[day_night])
    image_holder.active_material.node_tree.nodes['Image Texture'].image = D.images[back_img]
    image_holder.location.y = 1.5 + random.uniform(-0.3, 0.3)

    # Camera movement
    camera_object.rotation_euler = utils.r([72 + random.randint(-2, 2), 8, -82])

    C.scene.render.filepath = fp
    C.scene.render.image_settings.file_format = 'PNG'
    C.scene.camera = camera_object

    P, K, RT = camera_proj.get_3x4_P_matrix_from_blender(camera_object)

    file_root_name = f'{list(scenes_ids).index(human_path)}_{scenes_ids[human_path]}'

    with open(os.path.join(fp, 'cameras.json'), 'a+') as f_cam:
        previous_cameras = f_cam.read()
        if previous_cameras:
            previous_cameras = json.loads(previous_cameras)
        else:
            previous_cameras = {}
        previous_cameras[file_root_name] = {
            'P': utils.mat_to_list(P),
            'K': utils.mat_to_list(K),
            'RT': utils.mat_to_list(RT),
        }
        json.dump(previous_cameras, f_cam, indent=4)

    man.animation_data_clear()

    # Exemple: 150k / 200 / 2 = 1500 poses
    with open(os.path.join(fp_ann_2D, f'annotations_{file_root_name}.csv'), 'w') as annot_file_2D, \
            open(os.path.join(fp_ann_3D, f'annotations_{file_root_name}.csv'), 'w') as annot_file_3D:
        bone_lbls = list(man.pose.bones.keys())
        annot_file_2D.write(';'.join([lbl for bone in bone_lbls for lbl in [bone + k for k in ['_x', '_y']]]) + '\n')
        annot_file_3D.write(
            ';'.join([lbl for bone in bone_lbls for lbl in [bone + k for k in ['_X', '_Y', '_Z']]]) + '\n')

        for po in range(nb_pose):
            C.scene.frame_set(po * frame_rate)
            use_targets = nb_pose - po <= (nb_targets // 2) ** 2
            # use_targets = False
            human.switch_constraints(man, enable=not use_targets)
            random_pose.random_pose_ik(man, targets=car_targets if use_targets else None)

            bpy.ops.object.mode_set(mode='OBJECT')

            man.keyframe_insert(data_path="location", index=-1)
            man.keyframe_insert(data_path="rotation_euler", index=-1)

            bpy.ops.object.mode_set(mode='POSE')

            for bone in man.pose.bones:
                bone.keyframe_insert(data_path="rotation_euler", index=-1)
                if bone.name[-3:] == '_IK':
                    bone.keyframe_insert(data_path="location", index=-1)

            bpy.ops.object.mode_set(mode='OBJECT')

            # set output path so render won't get overwritten
            C.scene.render.filepath = os.path.join(fp_img, f"{file_root_name}_{po}" + f'_drive' if use_targets else '')
            bpy.ops.render.render(write_still=True)  # render still

            annotations_2D = []
            annotations_3D = []
            for lbl in bone_lbls:
                bone_3d = utils.get_head_pose(lbl, man)
                annotations_3D.append(f"{bone_3d[0]:.3f};{bone_3d[1]:.3f};{bone_3d[2]:.3f}")

                bone_2d = P @ bone_3d
                bone_2d /= bone_2d[-1]
                annotations_2D.append(f"{bone_2d[0]:.2f};{bone_2d[1]:.2f}")
            annot_file_2D.write(';'.join(annotations_2D) + '\n')
            annot_file_3D.write(';'.join(annotations_3D) + '\n')

C.scene.frame_end = int(frame_rate * (nb_pose - 0.5))
utils.select_only(man)
bpy.ops.object.mode_set(mode='POSE')

with open(info_path, 'w') as f:
    json.dump({
        'models': list(scenes_ids),
        'id_max_scenes': scenes_ids
    }, f, indent=4)

print('Done', '#' * 25)