#!/usr/bin/env python3

import argparse
import math
import os
import utils

import numpy as np


from data import House, DatasetToJSON
from math_utils import Transform
from priors.observations import ObjectCollection, ObservationCategory, RelativeObservation
from priors.pairwise import PairwiseArrangementPrior  # needed for pickle laod
from pyquaternion import Quaternion
from random import random


class ArrangementPriors:
    """
    A collection of object arrangement priors for estimating probability of an arrangement and sampling arrangements
    """
    def __init__(self, priors_file, w_dist=1.0, w_clearance=1.0, w_same_category=1.0, w_closest=1.0, w_orientation=1.0):
        self._priors = utils.pickle_load_compressed(priors_file)
        self._num_priors = 0.
        self._num_observations = 0.
        self._w_dist = w_dist
        self._w_clearance = w_clearance
        self._w_same_category = w_same_category
        self._w_closest = w_closest
        self._w_orientation = w_orientation
        for p_key, p in self._priors.items():
            self._priors[p_key] = p
            self._num_priors += 1
            self._num_observations += p.num_observations
            p.trim()

    @property
    def num_observations(self):
        """
        Return total number of observations in this ArrangementPriors set
        """
        return self._num_observations

    @property
    def pairwise_priors(self):
        """
        Return all pairwise priors
        """
        return self._priors

    def get(self, category_key):
        """
        Return pairwise prior for given category_key
        """
        return self._priors[category_key]

    def pairwise_occurrence_prob(self, category_key):
        """
        Return probability of pairwise observation with given category_key
        """
        norm = self._num_observations + self._num_priors  # add one smoothing
        if category_key in self._priors:
            return (self._priors[category_key].num_observations + 1) / norm
        else:
            return 1. / norm

    def log_prob(self, pairwise_observations_by_key):
        """
        Return log probability of given set of pairwise relative observations under arrangement priors
        """
        sum_lp = 0

        for key in pairwise_observations_by_key:
            # get observations and lp for this pairwise prior
            observations = pairwise_observations_by_key[key]

            # parameterize observations
            records = map(lambda o: o.parameterize(scheme='offsets_angles'), observations)
            x = np.stack(records, axis=0)
            # print(observations)
            # print(x)

            if key not in self._priors:
                lp = math.log(0.5)
            else:  # evaluate each observation against relevant prior
                prior = self._priors[key]
                if key.obj_category == 'room':  # distance and angle to closest point on wall
                    lp = self._w_clearance * prior.log_prob_closest(x)
                elif key.obj_category == key.ref_obj_category:  # same category instance
                    lp = self._w_same_category * self._w_closest * prior.log_prob_closest(x)
                else:  # all other categories
                    lp = self._w_closest * prior.log_prob_closest(x)
                lp += self._w_orientation * prior.log_prob_orientation(x)

            if self._w_dist > 0:
                dist = np.sum(np.square(x[:, 2:4]), axis=1, keepdims=True)
                lp = lp - (self._w_dist * dist/2)  # log(exp(-x^2/2))

            pairwise_occ_p = self.pairwise_occurrence_prob(key)
            lp = lp * pairwise_occ_p

            # sum up lps
            sum_lp += np.sum(lp, axis=0)

        if hasattr(sum_lp, 'shape'):
            sum_lp = np.sum(sum_lp)
        return sum_lp


class ArrangementGreedySampler:
    """
    Iterative optimization of object arrangements using greedy sampling of ArrangementPriors
    """
    def __init__(self, arrangement_priors, num_angle_divisions=8, num_pairwise_priors=-1, sim_mode='direct'):
        self._objects = ObjectCollection(sim_mode=sim_mode)
        self.room_id = None
        self._priors = arrangement_priors
        self._num_angle_divisions = num_angle_divisions
        self._num_pairwise_priors = num_pairwise_priors

    @property
    def objects(self):
        return self._objects

    def init(self, house, only_architecture=True, room_id=None):
        if not room_id:
            room_id = house.rooms[0].id
        self._objects.init_from_room(house, room_id, only_architecture=only_architecture)
        self.room_id = room_id

    def log_prob(self, filter_ref_obj=None, ignore_categories=list()):
        observations = self._objects.get_relative_observations(self.room_id, filter_ref_obj=filter_ref_obj,
                                                               ignore_categories=ignore_categories)
        observations_by_key = {}

        # top_k_prior_categories = None
        # if filter_ref_obj and num_pairwise_priors specified, filter observations to only those in top k priors
        # if filter_ref_obj and self._num_pairwise_priors > 0:
        #     category = self._objects.category(filter_ref_obj.modelId, scheme='final')
        #     priors = list(filter(lambda p: p.ref_obj_category == category, self._priors.pairwise_priors))
        #     k = min(self._num_pairwise_priors, len(priors))
        #     priors = list(sorted(priors, key=lambda p: self._priors.pairwise_occurrence_log_prob(p)))[-k:]
        #     top_k_prior_categories = set(map(lambda p: p.obj_category, priors))

        for o in observations.values():
            # only pairwise observations in which filter_ref_obj is the reference
            if filter_ref_obj and o.ref_id != filter_ref_obj.id:
                continue
            key = self._objects.get_observation_key(o)
            os_key = observations_by_key.get(key, [])
            os_key.append(o)
            observations_by_key[key] = os_key
        return self._priors.log_prob(observations_by_key)

    def get_candidate_transform(self, node, max_iterations=100):
        num_checks = 0
        zmin = self._objects.room.zmin
        while True:
            num_checks += 1
            p = self._objects.room.obb.sample()
            ray_from = [p[0], zmin - .5, p[2]]
            ray_to = [p[0], zmin + .5, p[2]]
            intersection = ags._objects.simulator.ray_test(ray_from, ray_to)
            if intersection.id == self.room_id + 'f' or num_checks > max_iterations:
                break
        xform = Transform()
        xform.set_translation([p[0], zmin + .1, p[2]])
        angle = random() * 2 * math.pi
        angular_resolution = 2 * math.pi / self._num_angle_divisions
        angle = round(angle / angular_resolution) * angular_resolution
        xform.set_rotation(radians=angle)
        return xform

    def sample_placement(self, node, n_samples, houses_log=None, max_attempts_per_sample=10, ignore_categories=list(),
                         collision_threshold=0):
        """
        Sample placement for given node
        """
        self._objects.add_object(node)
        max_lp = -np.inf
        max_xform = None
        max_house = None
        num_noncolliding_samples = 0
        for i in range(max_attempts_per_sample*n_samples):
            xform = self.get_candidate_transform(node)
            self._objects.update(node, xform=xform, update_sim=True)
            collisions = self._objects.get_collisions(obj_id_a=node.id)
            # print(f'i={i}, samples_so_far={num_noncolliding_samples}, n_samples={n_samples},'
            #       f'max_attempts_per_sample={max_attempts_per_sample}')
            if collision_threshold > 0:
                if min(collisions.values(), key=lambda c: c.distance).distance < -collision_threshold:
                    continue
            elif len(collisions) > 0:
                continue
            lp = self.log_prob(filter_ref_obj=node, ignore_categories=ignore_categories)
            print(f'lp={lp}')
            if lp > max_lp:
                max_xform = xform
                max_lp = lp
                if houses_log is not None:
                    max_house = self._objects.as_house()
            num_noncolliding_samples += 1
            if num_noncolliding_samples == n_samples:
                break
        if houses_log is not None:
            houses_log.append(max_house)
        self._objects.update(node, xform=max_xform, update_sim=True)

    def placeable_objects_sorted_by_size(self, house):
        objects = []
        fixed_objects = []
        for n in house.levels[0].nodes:
            if n.type != 'Object':
                continue
            category = self._objects.category(n.modelId, scheme='final')
            if category == 'door' or category == 'window':
                fixed_objects.append(n)
                continue
            objects.append(n)

        for o in objects:
            # to_delete = []
            # for k in o.__dict__:
            #     if k not in ['id', 'modelId', 'transform', 'type', 'valid', 'bbox']:
            #         to_delete.append(k)
            # for k in to_delete:
            #     delattr(o, k)
            dims = self._objects._object_data.get_aligned_dims(o.modelId)
            o.volume = dims[0] * dims[1] * dims[2]
        objects = list(sorted(objects, key=lambda x: x.volume, reverse=True))
        return objects, fixed_objects


if __name__ == '__main__':
    module_dir = os.path.dirname(os.path.abspath(__file__))
    parser = argparse.ArgumentParser(description='Arrangement model')
    parser.add_argument('--task', type=str, default='arrange', help='task [arrange]')
    parser.add_argument('--input', type=str, required=True, help='input scene to arrange')
    parser.add_argument('--output_dir', type=str, required=True, help='output directory to save arranged house states')
    parser.add_argument('--priors_dir', type=str, required=True, help='base directory to load priors')
    parser.add_argument('--only_final_state', action='store_true', help='save only final')
    parser.add_argument('--sim_mode', dest='sim_mode', default='direct', help='sim server [gui|direct|shared_memory]')
    parser.add_argument('--restart_sim_every_round', action='store_true', help='restart sim after every object placed')
    parser.add_argument('--max_objects', type=int, default=np.inf, help='maximum number of objects to place')
    parser.add_argument('--num_samples_per_object', type=int, default=100, help='placement samples per object')
    parser.add_argument('--ignore_doors_windows', action='store_true', help='ignore door and window priors')
    parser.add_argument('--collision_threshold', type=float, default=0, help='how much collision penetration to allow')
    args = parser.parse_args()

    if args.task == 'arrange':
        # initialize arrangement priors and sampler
        filename = os.path.join(args.priors_dir, f'priors.pkl.gz')
        ap = ArrangementPriors(priors_file=filename,
                               w_dist=0.0,
                               w_clearance=1.0,
                               w_closest=1.0,
                               w_orientation=1.0,
                               w_same_category=1.0)
        ags = ArrangementGreedySampler(arrangement_priors=ap, sim_mode=args.sim_mode)

        # load house and set architecture
        original_house = House(file_dir=args.input, include_support_information=False)
        ags.init(original_house, only_architecture=True)

        # split into placeable objects and fixed portals (doors, windows)
        placeables, fixed_objects = ags.placeable_objects_sorted_by_size(original_house)

        # add fixed objects (not to be arranged)
        for n in fixed_objects:
            ags.objects.add_object(n)

        # place objects one-by-one
        houses_states = []
        num_placeables = len(placeables)
        for (i, n) in enumerate(placeables):
            print(f'Placing {i+1}/{num_placeables}')
            if args.sim_mode == 'gui' and i < 2:  # center view on first couple of placements if gui_mode
                ags.objects.simulator.set_gui_rendering(enabled=True)
            houses_log = None if args.only_final_state else houses_states
            ags.sample_placement(node=n, n_samples=args.num_samples_per_object, houses_log=houses_log,
                                 ignore_categories=['window', 'door'] if args.ignore_doors_windows else [])
            if args.restart_sim_every_round:
                ags.objects.simulator.connect()
                ags.objects.reinit_simulator()
            if i+1 >= args.max_objects:
                break

        # append final house state
        house_final = ags._objects.as_house()
        houses_states.append(house_final)

        # save out states
        input_id = os.path.splitext(os.path.basename(args.input))[0]
        output_dir = os.path.join(args.output_dir, input_id)
        print(f'Saving results to {output_dir}...')
        dj = DatasetToJSON(dest=output_dir)
        empty_is = []
        for i, h in enumerate(houses_states):
            if h:
                h.trim()
            else:
                print(f'Warning: empty house state for step {i}, ignoring...')
        houses_states = list(filter(lambda h: h is not None, houses_states))
        for (i, s) in enumerate(dj.step(houses_states)):
            print(f'Saved {i}')

    print('DONE')