from egogym.utils import make_cabinet_xml, make_drawer_xml
from egogym.components import Object, ArticulableObject

class ObjectsManager:
    def __init__(self, objects_dict, np_random, shuffle=True):
        self.np_random = np_random
        sorted_objects = sorted(objects_dict, key=lambda x: x[0])
        if shuffle:
            np_random.shuffle(sorted_objects)
        self.objects_dict = sorted_objects
        self.index = -1

    def sample(self, random=False):
        if random:
            index = self.np_random.integers(len(self.objects_dict))
        else:
            self.index = (self.index + 1) % len(self.objects_dict)
            index = self.index
        if self.objects_dict[index][1] == "procedurally_generated":
            return ArticulableObject(
                name=self.objects_dict[index][0],
                func=globals()[f"make_{self.objects_dict[index][0]}_xml"],
                np_random=self.np_random,
            )
        else:
            return Object(
                name=self.objects_dict[index][0],
                path=self.objects_dict[index][1],
                np_random=self.np_random,
            )
