#!/usr/bin/env python3

import argparse
import csv
import glob
import os
from random import random
import Utils
from collections import namedtuple
from dataset.preprocess.Object import ObjectCategories, ObjectData
from dataset.preprocess.House import House
from dataset.preprocess.math_utils import *
from dataset.preprocess.math_utils.OBB import OBB
from dataset.preprocess.math_utils.Simulator import Simulator


class ObjectCollection:
    """Provides observation information for a collection of objects"""
    def __init__(self, categorization_type='final', sim_mode='direct'):
        self._object_data = ObjectData()
        self._object_categories = ObjectCategories()
        self._objects = {}
        self._room = None
        self._categorization_type = categorization_type
        self._sim = Simulator(mode=sim_mode)

    @property
    def simulator(self):
        return self._sim

    @property
    def room(self):
        return self._room

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

    def add_object(self, o):
        if o.id in self._objects and self._objects[o.id] != o:
            print(f'Warning: ignoring node with duplicate node id={o.id}')
            return None
        if hasattr(o, 'type') and o.type == 'Room':  # room node
            self._room = o
            self._sim.add_room(o, wall=True, floor=True, ceiling=False)
        else:  # other nodes
            self._sim.add_object(o)
        self.update(o, update_sim=True)
        return o.id

    def _remove(self, obj_id):
        if obj_id not in self._objects:
            print(f'Warning: tried to remove not present object with id={obj_id}')
        else:
            del self._objects[obj_id]
            self._sim.remove(obj_id)

    def update(self, o, xform=None, update_sim=False):
        if not hasattr(o, 'category'):
            o.category = self.category(o.modelId, scheme=self._categorization_type)
        model_id = o.modelId if hasattr(o, 'modelId') else None
        o.model2world = self.semantic_frame_matrix(model_id)
        o.xform = xform if xform else Transform.from_node(o)
        if hasattr(o, 'transform'):
            o.transform = o.xform.as_mat4_flat_row_major()
        o.obb = OBB.from_local2world_transform(np.matmul(o.xform.as_mat4(), o.model2world))
        o.frame = self.node_to_semantic_frame(o, o.obb)
        self._objects[o.id] = o
        # room geometries pre-transformed, so after obb computation above is done, set back to identity transform
        if hasattr(o, 'type') and o.type == 'Room':
            o.xform = Transform()
            o.transform = o.xform.as_mat4_flat_row_major()
        if update_sim:
            self._sim.set_state(obj_id=o.id, position=o.xform.translation, rotation_q=o.xform.rotation)

    def randomize_object_transforms(self):
        for o in self._objects.values():
            if o is self._room:
                continue
            t = self._room.obb.sample()
            t[1] = o.xform.translation[1]
            o.xform.set_translation(t)
            r = random() * 2 * math.pi
            o.xform.set_rotation(radians=r)
            self.update(o, o.xform)

    def init_from_room(self, house, room_id, only_architecture=False, update_sim=True):
        self.reset()
        room = next(r for r in house.rooms if r.id == room_id)
        self.add_object(room)
        if not only_architecture:
            for o in room.nodes:
                self.add_object(o)
        # if update_sim: self._sim.add_house_room_only(house, room, only_architecture=only_architecture,
        # no_ceil=True, no_floor=True)

    def init_from_house(self, house, update_sim=True):
        self.reset()
        for o in house.nodes:
            self.add_object(o)
        # if update_sim:
        #     self._sim.add_house(house, no_ceil=True, no_floor=True)

    def as_house(self):
        room_nodes = [dict(n.__dict__) for n in self._objects.values() if n.type != 'Room']
        for i, n in enumerate(room_nodes):
            n['originalId'] = n.originalId if hasattr(n, 'originalId') else n['id']
            n['id'] = f'0_{str(i + 1)}'  # overwrite id with linearized id
        room = {
            'id': '0_0',
            'originalId': self.room.originalId if hasattr(self.room, 'originalId') else self.room.id,
            'type': 'Room',
            'valid': 1,
            'modelId': self.room.modelId,
            'nodeIndices': list(range(1, len(room_nodes) + 1)),
            'roomTypes': self.room.roomTypes,
            'bbox': self.room.bbox,
        }
        house_dict = {
            'version': 'suncg@1.0.2',
            'id': self.room.house_id,
            'up': [0, 1, 0],
            'front': [0, 0, 1],
            'scaleToMeters': 1,
            'levels': [{'id': '0', 'nodes': [room] + room_nodes}]
        }
        return House(house_json=house_dict)

    def reset(self):
        self._objects = {}
        self._sim.reset()

    def reinit_simulator(self, wall=True, floor=True, ceiling=False):
        self._sim.reset()
        for o in self._objects.values():
            if hasattr(o, 'type') and o.type == 'Room':  # room node:
                self._sim.add_room(self.room, wall=wall, floor=floor, ceiling=ceiling)
            else:  # other nodes
                self._sim.add_object(o)
            self._sim.set_state(obj_id=o.id, position=o.xform.translation, rotation_q=o.xform.rotation)

    def get_relative_observations(self, room_id, filter_ref_obj, ignore_categories):
        out = {}
        ref_objects = [filter_ref_obj] if filter_ref_obj else self._objects.values()
        for o_r in ref_objects:
            if o_r.category in ignore_categories:
                continue
            for o_i in self._objects.values():
                if o_i is o_r:
                    continue
                if o_i.category in ignore_categories:
                    continue
                out[(o_i.id, o_r.id)] = self.object_frames_to_relative_observation(o_i.frame, o_r.frame, room_id)
        return out

    def get_collisions(self, include_collision_with_static=True, obj_id_a=None):
        # update sim state to match state of this ObjectCollection
        for o_id, o in self._objects.items():
            self._sim.set_state(obj_id=o.id, position=o.xform.translation, rotation_q=o.xform.rotation)
        self._sim.step()  # sim step needed to create contacts
        return self._sim.get_contacts(obj_id_a=obj_id_a, include_collision_with_static=include_collision_with_static)

    def get_observation_key(self, observation):
        room_node = self._objects[observation.room_id]
        room_types = '-'.join(room_node.roomTypes) if hasattr(room_node, 'roomTypes') else ''
        obj_node = self._objects[observation.obj_id]
        obj_cat = self.category(obj_node.modelId, scheme=self._categorization_type)
        ref_node = self._objects[observation.ref_id]
        ref_cat = self.category(ref_node.modelId, scheme=self._categorization_type)
        key = ObservationCategory(room_types=room_types, obj_category=obj_cat, ref_obj_category=ref_cat)
        return key

    def category(self, model_id, scheme):
        if 'rm' in model_id:
            return 'room'
        if scheme == 'coarse':
            return self._object_categories.get_coarse_category(model_id)
        elif scheme == 'fine':
            return self._object_categories.get_fine_category(model_id)
        elif scheme == 'final':
            return self._object_categories.get_final_category(model_id)
        else:
            raise RuntimeError(f'Unknown categorization type: {scheme}')

    def semantic_frame_matrix(self, model_id):
        if model_id in self._object_data.model_to_data:
            return self._object_data.get_model_semantic_frame_matrix(model_id)
        else:  # not a model, so assume identity semantic frame
            return np.identity(4)

    def object_frames_to_relative_observation(self, frame, ref_frame, room_id):
        ref_dims = ref_frame['obb'].half_dimensions
        rel_centroid = ref_frame['obb'].transform_point(frame['obb'].centroid)
        rel_min = ref_frame['obb'].transform_point(frame['aabb']['min'])
        rel_max = ref_frame['obb'].transform_point(frame['aabb']['max'])
        rel_up = ref_frame['obb'].transform_direction(frame['obb'].rotation_matrix[:3, 1])
        rel_front = ref_frame['obb'].transform_direction(-frame['obb'].rotation_matrix[:3, 2])  # note: +z = back
        cp = self._sim.get_closest_point(obj_id_a=frame['obj_id'], obj_id_b=ref_frame['obj_id'])
        rel_cp = ref_frame['obb'].transform_point(cp.positionOnAInWS)
        # NOTE: below approximate closest point calls are for removing pybullet call and debugging memory leak
        # cp = frame['obb'].closest_point(ref_frame['obb'].centroid)
        # rel_cp = ref_frame['obb'].transform_point(cp)
        out = RelativeObservation(room_id=room_id, obj_id=frame['obj_id'], ref_id=ref_frame['obj_id'],
                                  ref_dims=ref_dims, centroid=rel_centroid, min=rel_min, max=rel_max, closest=rel_cp,
                                  front=rel_front, up=rel_up)
        return out

    @staticmethod
    def node_to_semantic_frame(node, obb):
        aabb_min, aabb_max = obb.to_aabb()
        out = {
            'obj_id': node.id,
            'obb': obb,
            'aabb': {'min': aabb_min, 'max': aabb_max}
        }
        return out


# an observation of object id relative to a reference object ref_id
class RelativeObservation(namedtuple('RelativeObservation',
                                     ['room_id', 'obj_id', 'ref_id', 'ref_dims', 'centroid', 'min', 'max', 'closest',
                                      'front', 'up'])):
    __slots__ = ()  # prevent per-instance dict

    def to_str_list(self):
        return [self.room_id, self.obj_id, self.ref_id, nparr2str_compact(self.ref_dims),
                nparr2str_compact(self.centroid), nparr2str_compact(self.min), nparr2str_compact(self.max),
                nparr2str_compact(self.closest), nparr2str_compact(self.front), nparr2str_compact(self.up)]

    @classmethod
    def fromstring(cls, s):
        return RelativeObservation(room_id=s['room_id'], obj_id=s['obj_id'], ref_id=s['ref_id'], ref_dims=s['ref_dims'],
                                   centroid=str2nparr(s['centroid']), min=str2nparr(s['min']), max=str2nparr(s['max']),
                                   closest=str2nparr(s['closest']), front=str2nparr(s['front']), up=str2nparr(s['up']))

    def parameterize(self, scheme):
        if scheme == 'dist_angles':
            return self._to_distance_angles()
        elif scheme == 'dist_cos_sin_angles':
            return self._to_distance_cos_sin_angles()
        elif scheme == 'offsets_cos_sin_angles':
            return self._to_offsets_cos_sin_angles()
        elif scheme == 'offsets_angles':
            return self._to_offsets_angles()
        else:
            raise RuntimeError(f'Unknown RelativeObservation parameterization scheme {scheme}')

    def _to_distance_angles(self):
        d_center, a_center = relative_pos_to_xz_distance_angle(self.centroid)
        d_closest, a_closest = relative_pos_to_xz_distance_angle(self.closest)
        a_front = relative_dir_to_xz_angle(self.front)
        return [d_center, d_closest, a_center, a_closest, a_front]

    def _to_distance_cos_sin_angles(self):
        d_center, a_center = relative_pos_to_xz_distance_angle(self.centroid)
        d_closest, a_closest = relative_pos_to_xz_distance_angle(self.closest)
        a_front = relative_dir_to_xz_angle(self.front)
        return [d_center, d_closest, math.cos(a_center), math.sin(a_center), math.cos(a_closest), math.sin(a_closest),
                math.cos(a_front), math.sin(a_front)]

    def _to_offsets_angles(self):
        a_front = relative_dir_to_xz_angle(self.front)
        return [self.centroid[0], self.centroid[2], self.closest[0], self.closest[2], a_front]

    def _to_offsets_cos_sin_angles(self):
        a_front = relative_dir_to_xz_angle(self.front)
        return [self.centroid[0], self.centroid[2], self.closest[0], self.closest[2],
                math.cos(a_front), math.sin(a_front)]


# a record categorizing a relative observation by originating room type, object category and reference object category
ObservationCategory = namedtuple('ObservationCategory', ['room_types', 'obj_category', 'ref_obj_category'])


class RelativeObservationsDatabase:
    """Encapsulates a set of object arrangement priors"""

    def __init__(self, name, priors_dir, verbose=False, categorization_type='final'):
        self._name = name
        self._priors_dir = priors_dir
        Utils.ensuredir(self._priors_dir)
        self._verbose = verbose
        self._objects = ObjectCollection(categorization_type=categorization_type)
        self._semantic_frames = {}  # house_id -> {obj_id: SemanticFrame}
        self._observations = {}  # house_id -> {(obj_id,ref_obj_id): RelativeObservation}
        self._grouped_observations = {}  # {ObservationCategory: [RelativeObservation]}

    def clear(self):
        self._semantic_frames = {}
        self._observations = {}
        self._grouped_observations = {}

    @property
    def num_houses(self):
        return len(self._observations.keys())

    @property
    def grouped_observations(self):
        return self._grouped_observations

    def collect_observations(self, houses):
        for h in houses:
            self._observations[h.id] = self._get_observations_from_house(h)

    def save_observations_by_house(self, prefix='', save_frames=False):
        for (house_id, house_observations) in self._observations.items():
            basename = prefix + '_' + house_id
            csv_file = os.path.join(self._priors_dir, basename + '.relpos.csv')
            f = csv.writer(open(csv_file, 'w'))
            f.writerow(RelativeObservation._fields)
            for o in house_observations.values():
                f.writerow(o.to_str_list())

        if save_frames:
            for (house_id, house_frames) in self._semantic_frames.items():
                basename = prefix + '_' + house_id
                csv_file = os.path.join(self._priors_dir, basename + '.semframes.csv')
                f = csv.writer(open(csv_file, 'w'))
                f.writerow(['obj_id', 'local2world', 'world_aabb_min', 'world_aabb_max'])
                for frame in house_frames.values():
                    local2world_flat = nparr2str_compact(np.squeeze(np.asarray(frame['obb'].local2world.flatten())))
                    f.writerow([frame['obj_id'], local2world_flat,
                                nparr2str_compact(frame['aabb']['min']), nparr2str_compact(frame['aabb']['max'])])

    def load_observations(self, filenames, house_dir, load_frames=False):
        counts_house_ids = list(map(lambda f: f.split('_'), filenames))
        num_observations = len(counts_house_ids)
        for i, (obs_id, house_id) in enumerate(counts_house_ids):
            print(f'Loading observation obs_id={obs_id} house id={house_id}, {i}/{num_observations}')
            house_observations = self._observations.get(house_id, {})
            num_rows = 0
            for row in csv.DictReader(open(os.path.join(self._priors_dir, obs_id+'_'+house_id + '.relpos.csv'))):
                if len(row) == 0:
                    continue
                obs = RelativeObservation.fromstring(row)
                house_observations[(obs.obj_id, obs.ref_id)] = obs
                num_rows += 1
            if num_rows > 0:
                self._observations[obs_id] = house_observations

            if load_frames:
                house_frames = self._semantic_frames.get(house_id, {})
                for row in csv.DictReader(open(os.path.join(self._priors_dir, obs_id+'_'+house_id + '.semframes.csv'))):
                    local2world = np.matrix(str2nparr(row['local2world'])).reshape(4, 4)
                    frame = {'obj_id': row['obj_id'], 'obb': OBB.from_local2world_transform(local2world),
                             'aabb': {'min': row['world_aabb_min'], 'max': row['world_aabb_max']}}
                    house_frames[frame['obj_id']] = frame
                self._semantic_frames[house_id] = house_frames

        print('Grouping observations by categories...')
        groups = {}
        num_obs = len(self._observations)
        for i, (obs_id, house_id) in enumerate(counts_house_ids):
            print(f'Observation in house id={house_id} {i}/{num_obs}')
            house = House(file_dir=os.path.join(house_dir, obs_id + '.json'), include_support_information=False)
            self._objects.init_from_house(house, update_sim=False)
            for observation in self._observations[obs_id].values():
                key = self._objects.get_observation_key(observation)
                key_bin = groups.get(key, [])
                key_bin.append(observation)
                groups[key] = key_bin
        print('Done grouping')

        self._grouped_observations = groups

    def save(self):
        pkl_file = os.path.join(self._priors_dir, self._name + '.priors.pkl.gz')
        Utils.pickle_dump_compressed(self._grouped_observations, pkl_file)

    def load(self):
        pkl_file = os.path.join(self._priors_dir, self._name + '.priors.pkl.gz')
        self._grouped_observations = Utils.pickle_load_compressed(pkl_file)

    def _get_observations_from_house(self, house):
        if self._verbose:
            print(f'Observations for house id={house.id}...')

        self._objects.reset()
        self._objects.simulator.add_house(house, no_ceil=True, no_floor=True)
        relative_observations = {}  # (node_i.id,node_j.id) -> observation of node_i relative to node_j

        if len(house.rooms) == 0:
            print(f'House id={house.id} has no rooms, skipping.')
            return relative_observations

        for room in house.rooms:
            if self._verbose:
                print(f'Room id={room.id}, roomTypes={room.roomTypes}')
            # prepare node and room metadata
            self._objects.add_object(room)
            nodes = [n for n in room.nodes if n.type == 'Object' and n.valid]
            for o in nodes:
                self._objects.add_object(o)
            room_relative_observations = self._objects.get_relative_observations(room.id, 
                filter_ref_obj=None,
                ignore_categories=list())
            relative_observations.update(room_relative_observations)

            return relative_observations


if __name__ == '__main__':
    module_dir = os.path.dirname(os.path.abspath(__file__))
    parser = argparse.ArgumentParser(description='Compute arrangement priors')
    parser.add_argument('--task', type=str, required=True, help='<Required> task [collect|save_pkl]')
    parser.add_argument('--input', type=str, help='house json or directory with house json files')
    parser.add_argument('--priors_dir', type=str, required=True, help='priors data directory')
    parser.add_argument('--house_dir', type=str, help='house data directory')
    args = parser.parse_args()

    rod = RelativeObservationsDatabase(name='suncg_priors', priors_dir=args.priors_dir, verbose=True)

    if os.path.isdir(args.input):
        house_files = sorted(glob.glob(os.path.join(args.input, '**/*.json'), recursive=True))
    else:
        house_files = [args.input]

    if args.task == 'collect':
        for f in house_files:
            house = House(file_dir=f, include_support_information=False)
            prefix = os.path.splitext(os.path.basename(f))[0]
            priors_file = os.path.join(args.priors_dir, prefix + '_' + house.id + '.relpos.csv')
            if os.path.exists(priors_file):
                print(f'Priors already exist at {priors_file}')
            else:
                rod.collect_observations([house])
                rod.save_observations_by_house(prefix)
                rod.clear()

    if args.task == 'save_pkl':
        relpos_files = filter(lambda p: '.relpos.csv' in p, os.listdir(args.input))
        relpos_files = list(map(lambda p: os.path.basename(p).split('.')[0], relpos_files))
        print(relpos_files)
        rod.load_observations(relpos_files, args.house_dir, load_frames=False)
        # reduced = {k: rod.grouped_observations[k] for k in [
        #     ObservationCategory('Bedroom', 'chair', 'desk'), ObservationCategory('Bedroom', 'desk', 'chair'),
        # ]}
        # rod._grouped_observations = reduced
        rod.save()

    print('DONE')
