# Copyright 2025 The RLinf Authors.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     https://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

import copy
import time
from dataclasses import dataclass, field

import numpy as np

from ..franka_env import FrankaEnv, FrankaRobotConfig


@dataclass
class DexpnpConfig(FrankaRobotConfig):
    target_ee_pose: np.ndarray = field(default_factory=lambda: np.zeros(6))
    reward_threshold: np.ndarray = field(
        default_factory=lambda: np.array([0.01, 0.01, 0.01, 0.2, 0.2, 0.2])
    )
    enable_random_reset: bool = True
    enable_gripper_penalty: bool = False
    step_frequency: float = 5.0

    def __post_init__(self):
        self.compliance_param = {
            "translational_stiffness": 1000,
            "translational_damping": 89,
            "rotational_stiffness": 150,
            "rotational_damping": 7,
            "translational_Ki": 0,
            "translational_clip_x": 0.015,
            "translational_clip_y": 0.015,
            "translational_clip_z": 0.015,
            "translational_clip_neg_x": 0.015,
            "translational_clip_neg_y": 0.015,
            "translational_clip_neg_z": 0.015,
            "rotational_clip_x": 0.02,
            "rotational_clip_y": 0.02,
            "rotational_clip_z": 0.02,
            "rotational_clip_neg_x": 0.02,
            "rotational_clip_neg_y": 0.02,
            "rotational_clip_neg_z": 0.02,
            "rotational_Ki": 0,
        }
        self.precision_param = {
            "translational_stiffness": 3000,
            "translational_damping": 89,
            "rotational_stiffness": 300,
            "rotational_damping": 9,
            "translational_Ki": 0.1,
            "translational_clip_x": 0.01,
            "translational_clip_y": 0.01,
            "translational_clip_z": 0.01,
            "translational_clip_neg_x": 0.01,
            "translational_clip_neg_y": 0.01,
            "translational_clip_neg_z": 0.01,
            "rotational_clip_x": 0.05,
            "rotational_clip_y": 0.05,
            "rotational_clip_z": 0.05,
            "rotational_clip_neg_x": 0.05,
            "rotational_clip_neg_y": 0.05,
            "rotational_clip_neg_z": 0.05,
            "rotational_Ki": 0.1,
        }
        self.target_ee_pose = np.array(self.target_ee_pose)
        self.reset_ee_pose = self.target_ee_pose + np.array(
            [0.0, 0.0, 0.05, 0.0, 0.0, 0.0]
        )
        self.reward_threshold = np.array(self.reward_threshold)
        self.action_scale = np.array([0.03, 0.5, 1])
        self.ee_pose_limit_min = self.target_ee_pose - np.array(
            [0.02, 0.02, 0.02, 0.003, 0.003, 0.003]
        )
        self.ee_pose_limit_max = self.target_ee_pose + np.array(
            [0.02, 0.02, 0.1, 0.003, 0.003, 0.003]
        )
        self.hand_target_state = np.array(self.hand_target_state)
        self.hand_reset_state = np.array(self.hand_reset_state)


class DexpnpEnv(FrankaEnv):
    CONFIG_CLS = DexpnpConfig

    @property
    def task_description(self):
        return "pick up the toy and place it onto the plate"

    def go_to_rest(self, joint_reset=False):
        """
        Move to the rest position defined in base class.
        Add a small z offset before going to rest to avoid collision with object.
        """
        if self._is_hand:
            self._end_effector_action(self.config.hand_reset_state)
        else:
            self._end_effector_action(np.array([1.0]))
        self._franka_state = self._controller.get_state().wait()[0]
        self._move_action(self._franka_state.tcp_pose)

        self._franka_state = self._controller.get_state().wait()[0]
        # Move up to clear the slot
        reset_pose = copy.deepcopy(self._franka_state.tcp_pose)
        reset_pose[2] += 0.03
        time.sleep(5)
        self._interpolate_move(reset_pose, timeout=1)
        time.sleep(2)
        reset_pose[2] += 0.02
        self._interpolate_move(reset_pose, timeout=1)

        super().go_to_rest(joint_reset)
