import numpy as np
from core.skills.base_skill import BaseSkill, register_skill
from omegaconf import DictConfig
from omni.isaac.core.controllers import BaseController
from omni.isaac.core.robots.robot import Robot
from omni.isaac.core.tasks import BaseTask


# pylint: disable=unused-argument
@register_skill
class Home(BaseSkill):
    def __init__(self, robot: Robot, controller: BaseController, task: BaseTask, cfg: DictConfig, *args, **kwargs):
        super().__init__()
        self.robot = robot
        self.controller = controller
        self.task = task
        self.skill_cfg = cfg

        self.lr_hand = "right" if "right" in self.controller.robot_file else "left"
        if self.lr_hand == "left":
            self._joint_indices = self.robot.left_joint_indices
            self._joint_home = self.robot.left_joint_home
            if self.skill_cfg.get("gripper_state", None):
                self._gripper_state = self.skill_cfg["gripper_state"]
            else:
                self._gripper_state = self.robot.left_gripper_state
        elif self.lr_hand == "right":
            self._joint_indices = self.robot.right_joint_indices
            self._joint_home = self.robot.right_joint_home
            if self.skill_cfg.get("gripper_state", None):
                self._gripper_state = self.skill_cfg["gripper_state"]
            else:
                self._gripper_state = self.robot.right_gripper_state

        # !!! keyposes should be generated after previous skill is done
        self.manip_list = []

    def simple_generate_manip_cmds(self):
        manip_list = []
        curr_ee_trans, curr_ee_ori = self.controller.get_ee_pose()
        curr_joints = self.robot.get_joint_positions()[self._joint_indices]
        home_joints = self._joint_home

        for k in range(0, 50):
            arm_action = np.array(home_joints) * ((k + 1) / 40) + np.array(curr_joints) * (1 - (k + 1) / 40)
            cmd = (
                curr_ee_trans,
                curr_ee_ori,
                "dummy_forward",
                {"arm_action": arm_action, "gripper_state": self._gripper_state},
            )
            manip_list.append(cmd)

        self.manip_list = manip_list

    def is_feasible(self, th=5):
        return self.controller.num_plan_failed <= th

    def is_subtask_done(self, t_eps=0.088):
        assert len(self.manip_list) != 0
        curr_joints = self.robot.get_joint_positions()[self._joint_indices]
        target_joints = self.manip_list[0][3]["arm_action"]
        diff_trans = np.linalg.norm(curr_joints - target_joints)
        pose_flag = (diff_trans < t_eps,)
        self.plan_flag = self.controller.num_last_cmd > 10
        return np.logical_or(pose_flag, self.plan_flag)

    def is_done(self):
        if len(self.manip_list) == 0:
            return True
        if self.is_subtask_done():
            self.manip_list.pop(0)
        if self.is_success():
            self.manip_list.clear()
            print("Home Done")
        return len(self.manip_list) == 0

    def is_success(self, t_eps=0.088):
        curr_joints = self.robot.get_joint_positions()[self._joint_indices]
        diff_trans = np.linalg.norm(curr_joints - self._joint_home)
        # print(diff_trans)
        return diff_trans < t_eps