from data import *
import random
import scipy.misc as m
import math
import json
import torch
import torch.nn as nn
import torch.nn.functional as F
import torch.optim as optim
from torchvision import datasets, transforms
from models import *
from models.nade import *
from torch.autograd import Variable
from PIL import Image
import copy
from model_prior import *
from scene_synth import *
from categoryCounts_dataset import CategoryCountsDataset

class SceneSynthArrangementBaseline(SceneSynth):

    def __init__(self, *args, **kwargs):
        super(SceneSynthArrangementBaseline, self).__init__(*args, **kwargs)
    
    def synth_room(self, room_id, trial, size, samples, \
                   save_dir, temperature_cat, temperature_pixel, \
                   min_p, max_collision, save=True, save_heatmap=False):
        room = SynthedRoomArrangementBaseline(room_id, trial, size, samples, self, temperature_cat, temperature_pixel, min_p, max_collision)

        while room.should_continue():
            room.save_top_down_view(save_dir)
            room.save_json(save_dir)
            room.add_node()
        room.save_top_down_view(save_dir, final=True)
        room.save_json(save_dir, final=True)


class SynthedRoomArrangementBaseline(SynthedRoom):
    
    def __init__(self, room_id, trial, size, samples, synthesizer, temperature_cat, temperature_pixel, min_p, max_collision):
        self.__dict__.update(locals())
        del self.self #Of course I don't care about readability

        self.scene = RenderedScene(index = room_id, \
                                   data_dir = synthesizer.data_dir_relative, \
                                   data_root_dir = synthesizer.data_root_dir, \
                                   load_objects = True)

        self._parse_original_nodes(self.scene.object_nodes)
        self.scene.object_nodes = []

        self.composite = self.scene.create_composite()
        self.door_window_nodes = self.scene.door_window_nodes
        self.object_nodes = []
        self.empty_house_json = None
        self.failures = 0
        self.obj_data = ObjectData()

    def _parse_original_nodes(self, nodes):
        category_dict = ObjectCategories()
        self.category_counts = np.zeros(len(self.synthesizer.categories))
        self.remaining_models = [[] for i in range(len(self.synthesizer.categories))]
        for node in nodes:
            modelId = node["modelId"]
            category = node["category"]
            self.category_counts[category] += 1
            self.remaining_models[category].append(modelId)
        #print(self.category_counts)
        #print(self.remaining_models)
            #print(modelId, category, category2, self.synthesizer.categories[category2])
    
    def should_continue(self):
        return self.category_counts.sum() > 0
    
    def add_node(self):
        print(self.category_counts)
        print(self.remaining_models)
        self.location_category_map = None
        self.current_room = self.composite.get_composite()

        self.existing_collisions = self._get_collisions()

        category_counts = self.category_counts
        self.count = 0
        best_x, best_y, best_p, best_r, best_modelId = None, None, -100, None, None
        best_category = None
        while True:
            self.count += 1
            gridx,gridy,category = self._sample_location_category(category_counts)
            
            print(f"Choosing type {self._get_category_name(category)} at grid {gridx}, {gridy}")

            x,y = self._sample_exact_location(gridx, gridy, category)

            #print(f"Try placing an object at image space coordinate {x}, {y}")
                
            modelId, r, p = self._sample_model_rotation(x, y, category)
            print(p)
            if p > best_p:
                best_p = p
                best_x = x
                best_y = y
                best_modelId = modelId
                best_r = r
                best_category = category
            if p > self.min_p or self.count > 10:
                print(f"Choosing model {modelId} rotated by {r} radians")
                if self.count > 10:
                    self.failures += 1000 #bad bad
                break
            else:
                self.failures += 1
                print(f"Best probability is {p}, resample location-category")

        category_counts[best_category] -= 1
        self.remaining_models[best_category].remove(best_modelId)
        new_obj = SynthNode(best_modelId, best_category, best_x, best_y, best_r, self)
        self.composite.add_height_map(new_obj.get_render(), best_category, math.sin(best_r), math.cos(best_r))
        self.object_nodes.append(new_obj)


    def _sample_location_category(self, category_counts):
        if self.location_category_map is None:
            self.location_category_map = self._create_location_category_map(category_counts)

        total_p = self.location_category_map.sum()
        max_p = np.max(self.location_category_map)
        
        x,y,category = self._sample_location_category_helper(self.location_category_map, total_p)
        self.location_category_map[category][x][y] = 0

        return x,y,category 
    
    def _create_location_category_map(self, category_counts):
        location_category_map = super(SynthedRoomArrangementBaseline, self)._create_location_category_map()
        
        for k in range(self.synthesizer.num_categories):
            if category_counts[k] == 0:
                location_category_map[k] = 0

        return location_category_map

    def _sample_model_rotation(self, x, y, category, trials=10, num_orientations=16):
        num_categories = self.synthesizer.num_categories
        synthesizer = self.synthesizer
        
        best_p = -100
        best_orientation = 0
        best_model = None

        models = list(set(self.remaining_models[category]))
            
        #print("Trying different models and orientation...")
        for (trial, modelId) in enumerate(models):
            inputs = []
            for orientation in range(num_orientations):
                r = math.pi*orientation/num_orientations*2
                new_obj = SynthNode(modelId, category, x, y, r, self).get_render()
                new_room = self.composite.add_and_get_composite \
                                (new_obj, category, math.sin(r), math.cos(r))
                new_obj[new_obj>0] = 1
                new_room[-1] = new_obj

                inputs.append(new_room)
            with torch.no_grad():
                inputs = torch.stack(inputs)
                inputs = inputs.cuda()
                outputs = synthesizer.model_rotation(inputs)
                outputs = synthesizer.softmax(outputs).cpu().data.numpy()

            for orientation in range(num_orientations):
                #print(orientation, outputs[orientation][1])
                if outputs[orientation][1] > best_p:
                    p = outputs[orientation][1]
                    r = math.pi*best_orientation/num_orientations*2
                    new_node = SynthNode(modelId, category, x, y, r, self)
                    collisions = self._get_collisions([new_node])
                    if (len(collisions) - len(self.existing_collisions)) > 0:
                        p -= 1
                    if p > best_p:
                        best_p = p
                        best_model = modelId
                        best_orientation = orientation

            print(f"Testing model {trial+1} of {len(models)}...", end = "\r")
        print()

        return best_model, math.pi*best_orientation/num_orientations*2, best_p