import numpy as np

from cliport.tasks import Task

from cliport.tasks.grippers import Spatula
import random

from cliport.utils import utils

import pybullet as p


class GeneratedTask(Task):
    """Build a car using blocks."""

    def __init__(self):
        super().__init__()
        self.max_steps = 20
        self.lang_template = "build a car using blocks"
        self.task_completed_desc = "done building car."
        self.additional_reset()

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

        # Add car base.
        base_size = (0.12, 0.12, 0.02)
        base_urdf = 'car-base/car-base.urdf'
        base_pose = self.get_random_pose(env, base_size)
        env.add_object(base_urdf, base_pose, category='fixed')

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

        # Add blocks.
        block_size = (0.04, 0.04, 0.04)
        block_urdf = 'block/block.urdf'

        objs = []
        for i in range(4):
            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)]
        targs = [(utils.apply(base_pose, i), base_pose[1]) for i in place_pos]

        # Goal: blocks are stacked vertically (bottom row: red, blue, green).
        for i in range(4):
            self.add_goal(objs=[objs[i]], matches=np.ones((1, 1)), targ_poses=[targs[i]], replace=False,
                             rotations=True, metric='pose', params=None, step_max_reward=1 / 4, symmetries=[np.pi/2])
            self.lang_goals.append(self.lang_template.format(blocks="the blocks", stacked="stacked vertically"))

        # Goal: blocks are stacked horizontally (top row: yellow).
        for i in range(2):
            self.add_goal(objs=[objs[i]], matches=np.ones((1, 1)), targ_poses=[targs[i]], replace=False,
                             rotations=True, metric='pose', params=None, step_max_reward=1 / 4)