import os
import torch
import pickle
import itertools

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

from shapely.geometry import box
from shapely.affinity import rotate
from torch.utils.data import DataLoader
import torch.nn.functional as F
from diffusers.training_utils import EMAModel

from src.utils.util import *
from src.utils.visualize import *
from src.data import filter_function, get_encoded_dataset, get_dataset_raw_and_encoded
from src.data.threed_future_dataset import ThreedFutureDataset
from src.models import model_from_config, ObjectFeatureVQVAE


config_file = "./configs/bedroom_sg2sc_diffusion_objfeat.yaml"
config = load_config(config_file)
objects_dataset = ThreedFutureDataset.from_pickled_dataset(
        config["data"]["path_to_pickled_3d_futute_models"]
    )
print(f"Load [{len(objects_dataset)}] 3D-FUTURE models")
if "eval" not in config["data"]["encoding_type"]:
    config["data"]["encoding_type"] += "_eval"
train_dataset = get_encoded_dataset(
        config["data"],  # same encoding type as validation dataset
        filter_function(
            config["data"],
            split=config["training"].get("splits", ["train", "val"])
        ),
        path_to_bounds=None,
        augmentations=None,  # no need for prior statistics computation
        split=config["training"].get("splits", ["train", "val"])
    )
raw_dataset, dataset = get_dataset_raw_and_encoded(
    config["data"],
    filter_fn=filter_function(
        config["data"],
        split=config["training"].get("splits", ["train", "val"])
    ),
    augmentations=None,
    split=config["training"].get("splits", ["train", "val"])
)
print(f"Load [{len(dataset)}] validation scenes with [{dataset.n_object_types}] object types\n")
dataloader = DataLoader(
        dataset,
        batch_size=1,
        num_workers=1,
        pin_memory=False,
        collate_fn=dataset.collate_fn,
        shuffle=False
    )


device = torch.device("cuda:0")
print("Load pretrained VQ-VAE\n")
with open(f"/anonymous/workspace/InstructScene/objfeat_bounds.pkl", "rb") as f:
    kwargs = pickle.load(f)
vqvae_model = ObjectFeatureVQVAE("openshape_vitg14", "gumbel", **kwargs)
vqvae_model.load_state_dict(torch.load("/anonymous/workspace/InstructScene/threedfront_objfeat_vqvae_epoch_01999.pth", map_location="cpu")["model"])
vqvae_model = vqvae_model.to(device)
vqvae_model.eval()
# Initialize the model
model = model_from_config(
    config["model"],
    raw_dataset.n_object_types,
    raw_dataset.n_predicate_types
).to(device)

# Create EMA for the model
ema_config = config["training"]["ema"]
if ema_config["use_ema"]:
    ema_states = EMAModel(model.parameters())
    ema_states.to(device)
else:
    ema_states: EMAModel = None
checkpoint = torch.load('/anonymous/workspace/InstructScene/checkpoints/bedroom_sg2scdiffusion_objfeat_epoch_01999.pth',
                        map_location=device)

model.load_state_dict(checkpoint["model"])
if ema_states is not None:
    ema_states.load_state_dict(checkpoint["ema_states"])

# Evaluate with the EMA parameters if specified
if ema_states is not None:
    print(f"Copy EMA parameters to the model\n")
    ema_states.copy_to(model.parameters())
model.eval()

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'
}


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_uids, bboxes, out_path):
    polys, colors = get_box_polys(bboxes, dataset.class_labels, objects_dataset)
    piou_li = []
    for i in range(len(polys)):
        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')
        plot_polys(ax, polys[i], colors[i], alpha=0.3)
        piou_li.append(calculate_pioc(polys[i]))
        plt.savefig(os.path.join(out_path, f'{scene_uids[i]}.png'), pad_inches=0)
        plt.close()
    return np.mean(piou_li)


def calculate_pioc(polys):
    if len(polys) < 2:
        return 0
    piou_li = []
    for pa, pb in itertools.combinations(polys, 2):
        if pa.intersects(pb):
            current_val = pa.intersection(pb).area / pa.union(pb).area
            piou_li.append(current_val)
        else:
            piou_li.append(0)
    return np.mean(piou_li)


out_path = '/anonymous/workspace/ins_bedroom_train_val_pred/'
if not os.path.exists(out_path):
    os.makedirs(out_path)

classes = np.array(dataset.object_types)
rel_counts, rel_count_errors = 0, 0
piou_li = []
print("Sample boxes from scene graphs with the [diffusion model]")
for batch_idx, batch in enumerate(dataloader):
    # Move everything to the device
    for k, v in batch.items():
        if not isinstance(v, list):
            batch[k] = v.to(device)
    # Unpack the batch parameters
    objs = batch["objs"]
    edges = batch["edges"]
    objfeat_vq_indices = batch["objfeat_vq_indices"]
    obj_masks = batch["obj_masks"]
    # Generate the box parameters
    with torch.no_grad():
        boxes_pred = model.generate_samples(
            objs, edges, objfeat_vq_indices, obj_masks,
            vqvae_model,
            cfg_scale=1
        )

    # Decode objfeat indices to objfeat embeddings
    B, N = objfeat_vq_indices.shape[:2]
    objfeats = vqvae_model.reconstruct_from_indices(
        objfeat_vq_indices.reshape(B * N, -1)
    ).reshape(B, N, -1)
    objfeats = (objfeats * obj_masks[..., None].float()).cpu().numpy()

    objs = objs.cpu()
    obj_masks = obj_masks.cpu()
    boxes_pred = boxes_pred.cpu()  # (B, N, 8)

    bbox_params = {
        "class_labels": F.one_hot(objs, num_classes=raw_dataset.n_object_types + 1).float(),
        # +1 for empty node (not really used)
        "translations": boxes_pred[..., :3],
        "sizes": boxes_pred[..., 3:6],
        "angles": boxes_pred[..., 6:]
    }
    boxes = dataset.post_process(bbox_params)
    avg_piou = plot_bboxes(batch['scene_uids'], boxes, out_path)
    piou_li.append(avg_piou)

print(f'average piou {np.mean(piou_li)}')
