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
import pybullet as p

class BallTowerInZones(Task):
    """Pick up each sphere and stack it on top of the stand of the same color while maintaining the balance. 
    However, the stacking should be done in a specific color sequence - blue at the bottom, followed by green, 
    and finally red at the top. The spheres and stands are each placed within separate zones."""

    def __init__(self):
        super().__init__()
        self.max_steps = 10
        self.lang_template = "stack the {color} sphere on the {color} stand"
        self.task_completed_desc = "done stacking spheres."
        self.additional_reset()

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

        # Add zones.
        zone_size = (0.12, 0.12, 0)
        zone_urdf = 'zone/zone.urdf'
        zone_poses = []
        for _ in range(3):
            zone_pose = self.get_random_pose(env, zone_size)
            env.add_object(zone_urdf, zone_pose, 'fixed')
            zone_poses.append(zone_pose)

        # Sphere and stand colors.
        colors = [
            utils.COLORS['blue'], utils.COLORS['green'], utils.COLORS['red']
        ]

        # Add spheres and stands.
        sphere_size = (0.04, 0.04, 0.04)
        sphere_urdf = 'sphere/sphere.urdf'
        stand_size = (0.05, 0.05, 0.05)
        stand_urdf = 'stacking/stand.urdf'

        spheres = []
        stands = []
        for i in range(3):
            sphere_pose = self.get_random_pose(env, sphere_size)
            sphere_id = env.add_object(sphere_urdf, sphere_pose, color=colors[i])
            spheres.append(sphere_id)

            stand_pose = self.get_random_pose(env, stand_size)
            stand_id = env.add_object(stand_urdf, stand_pose, color=colors[i])
            stands.append(stand_id)

        # Goal: each sphere is stacked on the stand of the same color in the correct zone.
        for i in range(3):
            self.add_goal(objs=[spheres[i]], matches=np.ones((1, 1)), targ_poses=[zone_poses[i]], replace=False,
                    rotations=True, metric='pose', params=None, step_max_reward=1 / 3)
            self.lang_goals.append(self.lang_template.format(color=colors[i]))

        # Add distractors.
        n_distractors = 0
        while n_distractors < 6:
            is_sphere = np.random.rand() > 0.5
            urdf = sphere_urdf if is_sphere else stand_urdf
            size = sphere_size if is_sphere else stand_size
            pose = self.get_random_pose(env, obj_size=size)
            color = colors[n_distractors % len(colors)]

            obj_id = env.add_object(urdf, pose, color=color)
            n_distractors += 1