import numpy as np
from cliport.tasks.task import Task
from cliport.utils import utils

import pybullet as p
import random

class StackBlockPyramidSeq(Task):
    """Stacking Block Pyramid Sequence base class."""

    def __init__(self):
        super().__init__()
        self.max_steps = 12
        self.lang_template = "put the {pick} block on {place}"
        self.task_completed_desc = "done stacking block pyramid."
        self.additional_reset()

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

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

        # Block colors.
        colors, color_names = utils.get_colors(self.mode)

        # Add blocks.
        # x, y, z dimensions for the asset size
        block_size = (0.04, 0.04, 0.04)
        block_urdf = 'stacking/block.urdf'

        objs = []
        for i in range(6):
            block_pose = self.get_random_pose(env, block_size)
            block_id = env.add_object(block_urdf, block_pose, color=colors[i])
            objs.append(block_id)

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

        # Goal: make bottom row.
        self.add_goal(objs=[objs[0]], matches=np.ones((1, 1)), targ_poses=[targs[0]], replace=False,
                rotations=True, metric='pose', params=None, step_max_reward=1 / 6, symmetries=[np.pi/2])

        self.lang_goals.append(self.lang_template.format(pick=color_names[0],
                                                         place="the lightest brown block"))

        self.add_goal(objs=[objs[1]], matches=np.ones((1, 1)), targ_poses=[targs[1]], replace=False,
                rotations=True, metric='pose', params=None, step_max_reward=1 / 6, symmetries=[np.pi/2])
        self.lang_goals.append(self.lang_template.format(pick=color_names[1],
                                                         place="the middle brown block"))

        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 / 6, symmetries=[np.pi/2])
        self.lang_goals.append(self.lang_template.format(pick=color_names[2],
                                                         place="the darkest brown block"))

        # Goal: make middle row.
        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 / 6, symmetries=[np.pi/2])
        self.lang_goals.append(self.lang_template.format(pick=color_names[3],
                                                         place=f"the {color_names[0]} and {color_names[1]} blocks"))

        self.add_goal(objs=[objs[4]], matches=np.ones((1, 1)), targ_poses=[targs[4]], replace=False,
                rotations=True, metric='pose', params=None, step_max_reward=1 / 6, symmetries=[np.pi/2])
        self.lang_goals.append(self.lang_template.format(pick=color_names[4],
                                                         place=f"the {color_names[1]} and {color_names[2]} blocks"))

        # Goal: make top row.
        self.add_goal(objs=[objs[5]], matches=np.ones((1, 1)), targ_poses=[targs[5]], replace=False,
                rotations=True, metric='pose', params=None, step_max_reward=1 / 6, symmetries=[np.pi/2])
        self.lang_goals.append(self.lang_template.format(pick=color_names[5],
                                                         place=f"the {color_names[3]} and {color_names[4]} blocks"))