from collections import Counter
from copy import deepcopy
from typing import Dict, List

import numpy as np
from numpy.random.mtrand import RandomState

from robogym.envs.rearrange.common.mesh import MeshRearrangeEnv
from robogym.envs.rearrange.datasets.envstates.base import Envstate, EnvstateDataset
from robogym.envs.rearrange.datasets.objects.base import ObjectDataset


class RandomObjectDataset(EnvstateDataset):
    """
    Class for randomly generating environment states.

    An environment state is generated by randomly picking a set of objects from object datasets
    and then applying default environment randomizations (initial position, orientation, scales ..)
    """

    def __init__(
        self,
        object_datasets: Dict[str, ObjectDataset],
        random_state: RandomState,
        object_sample_prob: Dict[str, float],
    ):
        """
        :param object_sample_prob: {object dataset name: probability for sampling from it}.
        """
        super().__init__(object_datasets, random_state)

        self.object_sample_prob = object_sample_prob
        self.object_dataset_names = sorted(list(object_sample_prob))
        self.object_dataset_prob = [
            object_sample_prob[i] for i in self.object_dataset_names
        ]
        assert np.sum(self.object_dataset_prob) == 1.0

    def _reset(self, env: MeshRearrangeEnv):
        object_groups = MeshRearrangeEnv._sample_random_object_groups(env)

        num_groups = len(object_groups)
        (
            meshes,
            mesh_object_dataset_names,
            mesh_object_ids,
            mesh_scales,
        ) = self._sample_meshes(num_groups)

        attrs_per_group = {
            "material_args": MeshRearrangeEnv._sample_object_materials(env, num_groups),
            "color": MeshRearrangeEnv._sample_object_colors(env, num_groups),
            "scale": MeshRearrangeEnv._sample_object_size_scales(env, num_groups),
            "mesh_files": meshes,
        }
        attrs_per_group["scale"] *= np.array(mesh_scales)

        object_groups = MeshRearrangeEnv._set_group_attributes(
            env, object_groups, attrs_per_group
        )

        # Be sure to recreate sim before randomizing position and rotation, because
        # randomization algorithm required bounding box for the current mesh objects.
        # In addition, object_groups should be deepcopied so that scale normalization inside
        # _recreate_sim is not applied to the current variable.
        MeshRearrangeEnv._set_object_groups(env, deepcopy(object_groups))
        MeshRearrangeEnv._recreate_sim(env)

        init_quats = self._sample_object_initial_rotations(env)
        init_quats = self._post_process_quat(
            init_quats, object_groups, mesh_object_dataset_names
        )

        # Note that environment object placement function should be called after object
        # orientation is set. Therefore we first set object orientation quats to environment,
        # and then call object placement function.
        MeshRearrangeEnv._set_object_initial_rotations(env, init_quats)
        init_pos, is_valid = MeshRearrangeEnv._generate_object_placements(env)

        self.envstate = Envstate(
            object_groups=object_groups,
            object_dataset_names=mesh_object_dataset_names,
            object_ids=mesh_object_ids,
            init_pos=init_pos,
            is_valid=is_valid,
            init_quats=init_quats,
        )

    def _sample_meshes(self, num_groups: int):
        """ Sample object meshes from object datasets """
        object_dataset_type_indices = self._random_state.choice(
            len(self.object_dataset_names), p=self.object_dataset_prob, size=num_groups
        )
        idx_to_num: Counter = Counter(object_dataset_type_indices)

        meshes: List[List[str]] = []
        mesh_object_dataset_names = []
        mesh_object_ids = []
        mesh_scales = []

        for object_dataset_idx, n_groups in idx_to_num.items():
            object_dataset_name = self.object_dataset_names[object_dataset_idx]
            object_dataset = self.object_datasets[object_dataset_name]

            object_ids = object_dataset.sample(n_groups)
            for object_id in object_ids:
                mesh_object_dataset_names.append(object_dataset_name)
                mesh_object_ids.append(object_id)
                mesh_scales.append(object_dataset.mesh_scale)

            meshes.extend(object_dataset.get_mesh_list(object_ids))

        assert len(meshes) == len(mesh_object_dataset_names)
        return meshes, mesh_object_dataset_names, mesh_object_ids, mesh_scales

    def _sample_object_initial_rotations(self, env: MeshRearrangeEnv):
        return MeshRearrangeEnv._sample_object_initial_rotations(env)


create = RandomObjectDataset
