import os

import tqdm
import torch

import numpy as np
import matplotlib.pyplot as plt
import geopandas as gpd

from pyrr import Matrix44
from simple_3dviz import Mesh, Scene

from shapely.geometry import box
from shapely.affinity import rotate
from scene_diffusion.datasets.threed_front import ThreedFront
from scene_diffusion.datasets.threed_future_dataset import ThreedFutureDataset
from scene_diffusion.datasets import filter_function, get_dataset_raw_and_encoded
from scripts.training_utils import load_config
from scene_diffusion.networks import build_network
from scripts.utils import floor_plan_from_scene, export_scene, get_textured_objects
from scripts.utils import render as render_top2down


# config_file = "./config/diffusion/diffusion_bedrooms_dim512_nomask_instancond_objfeats_lat32.yaml"
config_file = "./config/diffusion/diffusion_livingrooms_dim512_nomask_instancond_objfeats_lat32.yaml"
config = load_config(config_file)
# weight_file = '/anonymous/workspace/pretrain/pretrained/bedrooms_uncond/bedrooms_uncond.pt'
weight_file = '/anonymous/workspace/pretrain/pretrained/livingrooms_uncond/livingrooms_uncond.pt'
device = torch.device("cuda:0")
raw_dataset, dataset = get_dataset_raw_and_encoded(
        config["data"],
        filter_fn=filter_function(
            config["data"],
            split=config["training"].get("splits", ["train", "val"])
        ),
        split=config["training"].get("splits", ["train", "val"])
    )
network, _, _ = build_network(
        dataset.feature_size, dataset.n_classes,
        config, weight_file, device=device
    )
network.eval()
# objects_dataset = ThreedFutureDataset.from_pickled_dataset('/anonymous/workspace/pickle_data/threed_future_model_bedroom.pkl')
objects_dataset = ThreedFutureDataset.from_pickled_dataset('/anonymous/workspace/pickle_data/threed_future_model_livingroom.pkl')
scene_top2down = Scene(size=(256, 256), background=[1,1,1,1])

scene_top2down.up_vector = (0,0,-1)
scene_top2down.camera_target = (0, 0, 0)
scene_top2down.camera_position = (0,4,0)
scene_top2down.light = (0,4,0)
scene_top2down.camera_matrix = Matrix44.orthogonal_projection(
    left=-3.1, right=3.1,
    bottom=3.1, top=-3.1,
    near=0.1, far=6
)
category2color = {
    'cabinet': '#FFFF00', 'children_cabinet': '#FFFF00', 'coffee_table': '#CCFF99', 'desk': '#A52A2A',
    'double_bed': '#FF0000', 'dressing_table': '#A52A2A',
    'kids_bed': '#FF0000', 'nightstand': '#F08080', 'single_bed': '#FF0000', 'sofa': '#FF9933', 'stool': '#CC6600',
    'table': '#A52A2A', 'tv_stand': '#FFCC99',
    'wardrobe': '#FFFF00'
}
n_sequences = len(raw_dataset)


def generate_bboxes(network, scene_idx, raw_dataset, dataset, path_to_floor_plan_textures, device):
    scene_idx = scene_idx if scene_idx < len(raw_dataset) else np.random.choice(len(raw_dataset))
    current_scene = raw_dataset[scene_idx]
    samples = dataset[scene_idx]
    floor_plan, tr_floor, room_mask = floor_plan_from_scene(current_scene, path_to_floor_plan_textures, no_texture=True)
    bbox_params = network.generate_layout(
        room_mask=room_mask.to(device),
        num_points=config["network"]["sample_num_points"],
        point_dim=config["network"]["point_dim"],
        text=torch.from_numpy(samples['desc_emb'])[None, :].to(device) if 'desc_emb' in samples.keys() else None,
        device=device,
        clip_denoised=True,
        batch_seeds=torch.arange(scene_idx, scene_idx + 1), batch_size=1
    )

    boxes = dataset.post_process(bbox_params)
    return boxes


def get_box_polys(boxes, classes, objects_dataset):
    batch_size = boxes['angles'].shape[0]
    polys = [[]] * batch_size
    colors = [[]] * batch_size
    for batch_idx in range(0, batch_size):
        for poly_idx in range(boxes['angles'][batch_idx].shape[0]):
            label_idx = boxes['class_labels'][batch_idx, poly_idx].argmax()
            label_str = classes[label_idx.item()]
            if label_str not in category2color:
                continue
            size = boxes['sizes'][batch_idx, poly_idx].tolist()
            pos = boxes['translations'][batch_idx, poly_idx].tolist()
            angle = boxes['angles'][batch_idx, poly_idx].tolist()[0]
            poly = box(pos[0] - size[0], pos[2] - size[2], pos[0] + size[0], pos[2] + size[2])
            poly = rotate(poly, angle, use_radians=True)
            polys[batch_idx].append(poly)
            colors[batch_idx].append(category2color[label_str])
    return polys, colors


def plot_polys(ax, polys, colors, alpha=0.1):
    if polys is None:
        polys = []
    if len(polys) < 1:
        return
    gpdf = gpd.GeoDataFrame({'geometry': polys, 'colors': colors})
    gpdf.plot(ax=ax, color=gpdf['colors'], alpha=alpha, aspect='equal', edgecolor='black')


def plot_bboxes(scene_id, bboxes, out_path):
    fig_size, dpi = 256, 72
    fig, ax = plt.subplots(figsize=(fig_size / dpi, fig_size / dpi), dpi=dpi)
    plt.gca().set_aspect('equal', adjustable='box')
    plt.gca().invert_yaxis()
    plt.axis('off')
    polys, colors = get_box_polys(bboxes, dataset.class_labels, objects_dataset)
    for i in range(len(polys)):
        plot_polys(ax, polys[i], colors[i], alpha=0.3)
        plt.savefig(os.path.join(out_path, f'{scene_id}.png'), pad_inches=0)
        plt.cla()
    plt.close()


def render_bboxes(scene_id, boxes, out_path):
    bbox_params_t = torch.cat([boxes["class_labels"], boxes["translations"], boxes["sizes"], boxes["angles"]],
                              dim=-1).cpu().numpy()
    renderables, trimesh_meshes, model_jids = get_textured_objects(bbox_params_t, objects_dataset, dataset.class_labels,
                                                                   diffusion=True, no_texture=True)
    r = render_top2down(scene_top2down, renderables, color=None, mode="shading",
                        frame_path=os.path.join(out_path, f'{scene_id}.png'))


def generate_plot_bboxes(network, scene_idx, raw_dataset, dataset, path_to_floor_plan_textures, device, out_path):
    bboxes = generate_bboxes(network, scene_idx, raw_dataset, dataset, path_to_floor_plan_textures, device)
    plot_bboxes(raw_dataset[scene_idx].scene_id, bboxes, os.path.join(out_path, 'plot'))
    render_bboxes(raw_dataset[scene_idx].scene_id, bboxes, os.path.join(out_path, 'render'))


# out_path = '/anonymous/workspace/room_train_val_pred/'
out_path = '/anonymous/workspace/living_room_train_val_pred/'
if not os.path.exists(out_path):
    os.makedirs(out_path)
for x in ['plot', 'render']:
    if not os.path.exists(os.path.join(out_path, x)):
        os.makedirs(os.path.join(out_path, x))
path_to_floor_plan_textures = './demo/floor_plan_texture_images'
for i in tqdm.tqdm(range(0, n_sequences)):
    generate_plot_bboxes(network, i, raw_dataset, dataset, path_to_floor_plan_textures, device, out_path)
