import numpy as np
import os
import pybullet as p
import random
from cliport.tasks import primitives
from cliport.tasks.grippers import Spatula
from cliport.tasks.task import Task
from cliport.utils import utils
import numpy as np
from cliport.tasks.task import Task
from cliport.utils import utils

class ContainerPyramidConstruction(Task):
    """Construct a pyramid of containers with specific color and shape arrangement."""

    def __init__(self):
        super().__init__()
        self.max_steps = 10
        self.lang_template = "stack the {color} {shape} container on the {level} level"
        self.task_completed_desc = "done constructing container pyramid."
        self.additional_reset()

    def reset(self, env):
        super().reset(env)

        # Add pallet.
        pallet_size = (0.35, 0.35, 0.01)
        pallet_pose = self.get_random_pose(env, pallet_size)
        pallet_urdf = 'pallet/pallet.urdf'
        env.add_object(pallet_urdf, pallet_pose, 'fixed')

        # Define container shapes and colors.
        shapes = ['square', 'circle', 'triangle']
        colors = ['red', 'blue', 'green']
        levels = ['bottom', 'middle', 'top']

        # Add containers.
        container_urdf = 'container/container-template.urdf'
        containers = []
        for i in range(3):
            for j in range(3 - i):
                # Define container size and color.
                container_size = (0.1, 0.1, 0.1)
                container_color = colors[i]

                # Define container pose.
                x_offset = 0.05 * (j - (3 - i - 1) / 2)
                z_offset = 0.1 * i
                container_pose = (pallet_pose[0] + x_offset, pallet_pose[1], pallet_pose[2] + z_offset)

                # Add container to the environment.
                replace = {'DIM': container_size, 'HALF': (container_size[0] / 2, container_size[1] / 2, container_size[2] / 2)}
                container_urdf_filled = self.fill_template(container_urdf, replace)
                container_id = env.add_object(container_urdf_filled, container_pose, color=container_color)
                containers.append(container_id)

        # Add goals.
        for i in range(3):
            for j in range(3 - i):
                idx = int(i * (i + 1) / 2 + j)
                language_goal = self.lang_template.format(color=colors[i], shape=shapes[j], level=levels[i])
                self.add_goal(objs=[containers[idx]], matches=np.ones((1, 1)), targ_poses=[pallet_pose], replace=False,
                              rotations=True, metric='pose', params=None, step_max_reward=1 / 6, language_goal=language_goal)