"""
Handles pre-rendered rooms generated by data.top_down
RenderedScene loads the pre-rendered data
along with other precomputed information,
and creates RenderedComposite, which combines those to 
create the multi-channel top-down view used in the pipeline
"""
from torch.utils import data
from data import ObjectCategories, House, Obj
import random
import numpy as np
import math
import pickle
import os
import json
import copy
import torch
import utils

class RenderedScene():
    """
    Loading a rendered room

    Attributes
    ----------
    category_map (ObjectCategories): object category mapping
        that should be the same across all instances of the class
    categories (list[string]): all categories present in this room type.
        Loaded once when the first room is loaded to reduce disk access.
    cat_to_index (dict[string, int]): maps a category to corresponding index
    current_data_dir (string): keep track of the current data directory, if 
        it changes, then categories and cat_to_index should be recomputed
    """
    category_map = ObjectCategories()
    categories = None
    cat_to_index = None
    current_data_dir = None

    def __init__(self, index, data_dir, data_root_dir=None, \
                 shuffle=True, load_objects=True, seed=None):
        """
        Load a rendered scene from file

        Parameters
        ----------
        index (int): room number
        data_dir (string): location of the pre-rendered rooms
        data_root_dir (string or None, optional): if specified,
            use this as the root directory
        shuffle (bool, optional): If true, randomly order the objects
            in the room. Otherwise use the default order as written
            in the original dataset
        load_objects (bool, optional): If false, only load the doors
            and windows. Otherwise load all objects in the room
        seed (int or None, optional): if set, use a fixed random seed
            so we can replicate a particular experiment
        """
        if seed:
            random.seed(seed)

        if not data_root_dir:
            data_root_dir = utils.get_data_root_dir()
        
        if RenderedScene.categories is None or RenderedScene.current_data_dir != data_dir:
            with open(f"{data_root_dir}/{data_dir}/final_categories_frequency", "r") as f:
                lines = f.readlines()
                cats = [line.split()[0] for line in lines]

            RenderedScene.categories = [cat for cat in cats if cat not in set(['window', 'door'])]
            RenderedScene.cat_to_index = {RenderedScene.categories[i]:i for i in range(len(RenderedScene.categories))}
            RenderedScene.current_data_dir = data_dir

        with open(f"{data_root_dir}/{data_dir}/{index}.pkl", "rb") as f:
            (self.floor, self.wall, nodes), self.room = pickle.load(f)

        self.object_nodes = []
        self.door_window_nodes = []
        for node in nodes:
            category = RenderedScene.category_map.get_final_category(node["modelId"])
            if category in ["door", "window"]:
                node["category"] = category
                self.door_window_nodes.append(node)
            elif load_objects:
                node["category"] = RenderedScene.cat_to_index[category]
                self.object_nodes.append(node)
        
        if shuffle:
            random.shuffle(self.object_nodes)
       
    def create_composite(self):
        """
        Create a initial composite that only contains the floor,
        wall, doors and windows. See RenderedComposite for how
        to add more objects
        """
        r = RenderedComposite(RenderedScene.categories, self.floor, self.wall, self.door_window_nodes)
        return r

class RenderedComposite():
    """
    Multi-channel top-down composite, used as input to NN
    """
    def __init__(self, categories, floor, wall, door_window_nodes=None):
        #Optional door_window just in case
        self.size = floor.shape[0]
        
        self.categories = categories
        
        self.room_mask = (floor + wall)
        self.room_mask[self.room_mask != 0] = 1
        
        self.wall_mask = wall.clone()
        self.wall_mask[self.wall_mask != 0] = 0.5
        
        self.height_map = torch.max(floor, wall)
        self.cat_map = torch.zeros((len(self.categories),self.size,self.size))

        self.sin_map = torch.zeros((self.size,self.size))
        self.cos_map = torch.zeros((self.size,self.size))

        self.door_map = torch.zeros((self.size, self.size))
        self.window_map = torch.zeros((self.size, self.size))

        if door_window_nodes:
            for node in door_window_nodes:
                h = node["height_map"]
                xsize, ysize = h.size()
                xmin = math.floor(node["bbox_min"][0])
                ymin = math.floor(node["bbox_min"][2])
                if xmin < 0: xmin = 0
                if ymin < 0: ymin = 0
                to_add = torch.zeros((self.size, self.size))
                to_add[xmin:xmin+xsize,ymin:ymin+ysize] = h
                update = to_add > self.height_map
                self.height_map[update] = to_add[update]
                self.wall_mask[to_add>0] = 1
                to_add[to_add>0] = 0.5
                if node["category"] == "door":
                    self.door_map = self.door_map + to_add
                else:
                    self.window_map = self.window_map + to_add
    
    def get_transformation(self, transform):
        """
        Bad naming, really just getting the sin and cos of the
        angle of rotation.
        """
        a = transform[0]
        b = transform[8]
        scale = (a**2+b**2)**0.5
        return (b/scale, a/scale)

    def add_height_map(self, to_add, category, sin, cos):
        """
        Add a new object to the composite. 
        Height map, category, and angle of rotation are
        all the information required.
        """
        update = to_add>self.height_map
        self.height_map[update] = to_add[update]
        mask = torch.zeros(to_add.size())
        mask[to_add>0] = 0.5
        self.cat_map[category] = self.cat_map[category] + mask
        self.sin_map[update] = (sin + 1) / 2
        self.cos_map[update] = (cos + 1) / 2

    def add_node(self, node):
        """
        Add a new object to the composite.
        Computes the necessary information and calls
        add_height_map
        """
        h = node["height_map"]
        category = node["category"]
        xsize, ysize = h.shape
        xmin = math.floor(node["bbox_min"][0])
        ymin = math.floor(node["bbox_min"][2])
        to_add = torch.zeros((self.size, self.size))
        to_add[xmin:xmin+xsize,ymin:ymin+ysize] = h
        sin, cos = self.get_transformation(node["transform"])
        self.add_height_map(to_add, category, sin, cos)

    def add_nodes(self, nodes):
        for node in nodes:
            self.add_node(node)
    
    def get_cat_map(self):
        return self.cat_map.clone()

    def add_and_get_composite(self, to_add, category, sin, cos, \
                              num_extra_channels=1, temporary=True):
        """
        Sometimes we need to create a composite to test
        if some objects should be added, without actually
        fixing the object to the scene. This method allows doing so.
        See get_composite.
        """
        if not temporary:
            raise NotImplementedError
        update = to_add>self.height_map
        mask = torch.zeros(to_add.size())
        mask[to_add>0] = 0.5
        composite = torch.zeros((len(self.categories)+num_extra_channels+8, self.size, self.size))
        composite[0] = self.room_mask
        composite[1] = self.wall_mask
        composite[2] = self.cat_map.sum(0) + mask
        composite[3] = self.height_map
        composite[3][update] = to_add[update]
        composite[4] = self.sin_map
        composite[4][update] = (sin + 1) / 2
        composite[5] = self.cos_map
        composite[5][update] = (cos + 1) / 2
        composite[6] = self.door_map
        composite[7] = self.window_map
        for i in range(len(self.categories)):
            composite[i+8] = self.cat_map[i]
        composite[8+category] += mask
        
        return composite


    def get_composite(self, num_extra_channels=1, ablation=None):
        """
        Create the actual multi-channel representation.
        Which is a N x img_size x img_size tensor.
        See the paper for more information.
        Current channel order:
            -0: room mask
            -1: wall mask
            -2: object mask
            -3: height map
            -4, 5: sin and cos of the angle of rotation
            -6, 7: single category channel for door and window
            -8~8+C: single category channel for all other categories

        Parameters
        ----------
        num_extra_channels (int, optional): number of extra empty 
            channels at the end. 1 for most tasks, 0 for should continue
        ablation (string or None, optional): if set, return a subset of all
            the channels for ablation study, see the paper for more details
        """
        if ablation is None:
            composite = torch.zeros((len(self.categories)+num_extra_channels+8, self.size, self.size))
            composite[0] = self.room_mask
            composite[1] = self.wall_mask
            composite[2] = self.cat_map.sum(0)
            composite[3] = self.height_map
            composite[4] = self.sin_map
            composite[5] = self.cos_map
            composite[6] = self.door_map
            composite[7] = self.window_map
            for i in range(len(self.categories)):
                composite[i+8] = self.cat_map[i]
        elif ablation == "depth":
            composite = torch.zeros((1+num_extra_channels, self.size, self.size))
            composite[0] = self.height_map
        elif ablation == "basic":
            composite = torch.zeros((6+num_extra_channels, self.size, self.size))
            composite[0] = self.room_mask
            composite[1] = self.wall_mask
            composite[2] = self.cat_map.sum(0)
            composite[3] = self.height_map
            composite[4] = self.sin_map
            composite[5] = self.cos_map
        else:
            raise NotImplementedError

        return composite