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 BuildCylinderStructure(Task):
    """Construct a structure using four colored cylinders (red, blue, green, yellow) on a square base."""

    def __init__(self):
        super().__init__()
        self.max_steps = 5
        self.lang_template = "construct a structure using four colored cylinders on a square base"
        self.task_completed_desc = "done building the cylinder structure."
        self.additional_reset()

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

        # Add square base.
        # x, y, z dimensions for the asset size
        base_size = (0.15, 0.15, 0.005)
        base_urdf = 'square/square-template.urdf'
        base_pose = self.get_random_pose(env, base_size)
        env.add_object(base_urdf, base_pose, category='fixed')

        # Cylinder colors.
        colors = [
            utils.COLORS['red'], utils.COLORS['blue'], utils.COLORS['green'], utils.COLORS['yellow']
        ]

        # Add cylinders.
        # x, y, z dimensions for the asset size
        cylinder_size = (0.04, 0.04, 0.08)
        cylinder_urdf = 'cylinder/cylinder-template.urdf'

        objs = []
        for i in range(4):
            cylinder_pose = self.get_random_pose(env, cylinder_size)
            cylinder_id = env.add_object(cylinder_urdf, cylinder_pose, color=colors[i])
            objs.append(cylinder_id)

        # Associate placement locations for goals.
        place_pos = [(0, -0.05, 0.04), (0, 0.05, 0.04),
                     (0, 0.05, 0.12), (0, -0.05, 0.12)]
        targs = [(utils.apply(base_pose, i), base_pose[1]) for i in place_pos]

        # Goal: red and blue cylinders are placed side by side on the base.
        self.add_goal(objs=objs[:2], matches=np.ones((2, 2)), targ_poses=targs[:2], replace=False,
                rotations=True, metric='pose', params=None, step_max_reward=1 / 2, symmetries=[np.pi/2]*2)
        self.lang_goals.append("place the red and blue cylinders side by side on the base")

        # Goal: green cylinder is placed on top of the blue cylinder.
        self.add_goal(objs=[objs[2]], matches=np.ones((1, 1)), targ_poses=[targs[2]], replace=False,
                rotations=True, metric='pose', params=None, step_max_reward=1 / 2, symmetries=[np.pi/2])
        self.lang_goals.append("place the green cylinder on top of the blue cylinder")

        # Goal: yellow cylinder is placed on top of the red cylinder.
        self.add_goal(objs=[objs[3]], matches=np.ones((1, 1)), targ_poses=[targs[3]], replace=False,
                rotations=True, metric='pose', params=None, step_max_reward=1 / 2, symmetries=[np.pi/2])
        self.lang_goals.append("place the yellow cylinder on top of the red cylinder")