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


class BlockInsertion(Task):
    """pick up the L-shaped red block and place it into the L-shaped fixture."""

    def __init__(self):
        super().__init__()
        self.max_steps = 3
        self.lang_template = "put the L shape block in the L shape hole"
        self.task_completed_desc = "done with insertion."
        self.additional_reset()

    def get_random_pose(self, env, obj_size):
        pose = super().get_random_pose(env, obj_size)
        pos, rot = pose
        rot = utils.eulerXYZ_to_quatXYZW((0, 0, np.pi / 2))
        return pos, rot

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

        """Add L-shaped block."""
        size = (0.1, 0.1, 0.04)
        urdf = 'insertion/ell.urdf'
        pose = self.get_random_pose(env, size)
        block_id = env.add_object(urdf, pose)

        """Add L-shaped fixture to place block."""
        size = (0.1, 0.1, 0.04)
        urdf = 'insertion/fixture.urdf'
        targ_pose = self.get_random_pose(env, size)
        env.add_object(urdf, targ_pose, 'fixed')

        self.add_goal(objs=[block_id], matches=np.int32([[1]]), targ_poses=[targ_pose], replace=False,
                rotations=False, metric='pose', params=None, step_max_reward=1, symmetries=[2 * np.pi])
        self.lang_goals.append(self.lang_template)