# ----------------------------------------------------------------------------
# Copyright (c) 2021-2026 DexForce Technology Co., Ltd.
#
# 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
#
#     http://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, express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# ----------------------------------------------------------------------------

import torch
import numpy as np

from embodichain.lab.gym.envs import EmbodiedEnv, EmbodiedEnvCfg
from embodichain.lab.gym.utils.registration import register_env
from embodichain.lab.sim.planners import (
    MotionGenerator,
    MotionGenCfg,
    MotionGenOptions,
    ToppraPlannerCfg,
    ToppraPlanOptions,
    PlanState,
    MoveType,
    MovePart,
)
from embodichain.lab.sim.planners.utils import TrajectorySampleMethod
from embodichain.utils import logger

__all__ = ["StackBlocksTwoEnv"]


@register_env("StackBlocksTwo-v1", max_episode_steps=600)
class StackBlocksTwoEnv(EmbodiedEnv):
    def __init__(self, cfg: EmbodiedEnvCfg = None, **kwargs):
        super().__init__(cfg, **kwargs)

        action_config = kwargs.get("action_config", None)
        if action_config is not None:
            self.action_config = action_config

    def create_demo_action_list(self, *args, **kwargs):
        """
        Create a demonstration action list for stacking block_2 on top of block_1.

        This expert trajectory:
        1. Moves to block_2 position
        2. Grasps block_2
        3. Lifts block_2
        4. Moves to block_1 position
        5. Places block_2 on top of block_1
        6. Releases and retracts

        Returns:
            list: A list of demo actions generated by the task.
        """
        try:
            block1 = self.sim.get_rigid_object("block_1")
            block2 = self.sim.get_rigid_object("block_2")
        except Exception as e:
            logger.log_warning(f"Blocks not found: {e}, returning empty action list.")
            return []

        # Get block positions
        block1_pose = block1.get_local_pose(to_matrix=True)
        block2_pose = block2.get_local_pose(to_matrix=True)

        block1_pos = block1_pose[:, :3, 3]  # (num_envs, 3)
        block2_pos = block2_pose[:, :3, 3]  # (num_envs, 3)

        right_arm_ids = self.robot.get_joint_ids(name="right_arm")
        right_eef_ids = self.robot.get_joint_ids(name="right_eef")

        init_qpos = self.robot.get_qpos()
        init_right_arm_qpos = init_qpos[:, right_arm_ids]

        init_right_arm_xpos = self.robot.compute_fk(
            qpos=init_right_arm_qpos, name="right_arm", to_matrix=True
        )

        # 1. Approach block_2
        approach_block2_xpos = init_right_arm_xpos.clone()
        approach_block2_xpos[:, :3, 3] = block2_pos + torch.tensor(
            [0.02, 0.0, 0.1], dtype=torch.float32, device=self.device
        ).unsqueeze(0)

        # 2. Pick block_2
        pick_block2_xpos = init_right_arm_xpos.clone()
        pick_block2_xpos[:, :3, 3] = block2_pos + torch.tensor(
            [0.02, 0.0, -0.025], dtype=torch.float32, device=self.device
        ).unsqueeze(0)

        # 3. Lift block_2
        lift_block2_xpos = pick_block2_xpos.clone()
        lift_block2_xpos[:, 2, 3] += 0.15

        # 4. Approach block_1
        approach_block1_xpos = init_right_arm_xpos.clone()
        target_place_pos = block1_pos.clone()
        target_place_pos[:, 2] += 0.05  # block_1 height + block_2 height
        approach_block1_xpos[:, :3, 3] = target_place_pos + torch.tensor(
            [0.025, 0.0, 0.05], dtype=torch.float32, device=self.device
        ).unsqueeze(0)

        # 5. Place block_2 on block_1
        place_block2_xpos = approach_block1_xpos.clone()
        place_block2_xpos[:, :3, 3] = target_place_pos + torch.tensor(
            [0.025, 0.0, 0.02], dtype=torch.float32, device=self.device
        ).unsqueeze(0)

        # 6. Retract
        retract_xpos = approach_block1_xpos.clone()

        retract_xpos[:, 2, 3] += 0.1

        # Compute IK for each waypoint
        waypoints = [
            ("approach_block2", approach_block2_xpos),
            ("pick_block2", pick_block2_xpos),
            ("lift_block2", lift_block2_xpos),
            ("approach_block1", approach_block1_xpos),
            ("place_block2", place_block2_xpos),
            ("retract", retract_xpos),
        ]

        qpos_waypoints = []
        current_qpos = init_right_arm_qpos

        for waypoint_name, waypoint_xpos in waypoints:
            is_success, qpos = self.robot.compute_ik(
                pose=waypoint_xpos, joint_seed=current_qpos, name="right_arm"
            )

            # Check IK success
            success_flag = (
                is_success.all() if isinstance(is_success, torch.Tensor) else is_success
            )
            if not success_flag:
                logger.log_warning(f"IK failed for {waypoint_name}")
                qpos = current_qpos

            qpos_waypoints.append(qpos)
            current_qpos = qpos

        # Create motion generator for smooth trajectory
        motion_cfg = MotionGenCfg(
            planner_cfg=ToppraPlannerCfg(
                robot_uid=self.robot.uid,
            )
        )
        self.motion_generator = MotionGenerator(cfg=motion_cfg)

        action_list = []

        gripper_open = torch.tensor(
            [0.05, 0.05], dtype=torch.float32, device=self.device
        )
        gripper_close = torch.tensor(
            [0.0, 0.0], dtype=torch.float32, device=self.device
        )

        # Convert waypoints to numpy for motion generator
        qpos_waypoints_np = [q[0].cpu().numpy() for q in qpos_waypoints]

        # Generate trajectory segments
        segments = [
            (0, 1, 30, gripper_open),
            (1, 1, 10, gripper_close),
            (1, 2, 30, gripper_close),
            (2, 3, 40, gripper_close),
            (3, 4, 30, gripper_close),
            (4, 4, 10, gripper_open),
            (4, 5, 30, gripper_open),
        ]

        for start_idx, end_idx, num_steps, gripper_state in segments:
            if start_idx == end_idx:
                # Hold position (for gripper action)
                qpos = qpos_waypoints[start_idx]
                for _ in range(num_steps):
                    action = init_qpos.clone()
                    action[:, right_arm_ids] = qpos
                    action[:, right_eef_ids] = gripper_state.unsqueeze(0).expand(
                        self.num_envs, -1
                    )
                    action_list.append(action)
            else:
                # Generate smooth trajectory between waypoints
                qpos_list = [qpos_waypoints_np[start_idx], qpos_waypoints_np[end_idx]]

                options = MotionGenOptions(
                    control_part="right_arm",
                    plan_opts=ToppraPlanOptions(
                        constraints={
                            "velocity": 0.2,
                            "acceleration": 0.5,
                        },
                        sample_method=TrajectorySampleMethod.QUANTITY,
                        sample_interval=num_steps,
                    ),
                )
                # Joint space trajectory
                target_states = [
                    PlanState(
                        qpos=torch.as_tensor(qpos_waypoints_np[start_idx]),
                        move_type=MoveType.JOINT_MOVE,
                    ),
                    PlanState(
                        qpos=torch.as_tensor(qpos_waypoints_np[end_idx]),
                        move_type=MoveType.JOINT_MOVE,
                    ),
                ]
                plan_result = self.motion_generator.generate(
                    target_states=target_states, options=options
                )

                # Convert to torch and add to action list
                for qpos_item in plan_result.positions:
                    qpos = torch.as_tensor(
                        qpos_item, dtype=torch.float32, device=self.device
                    )

                    qpos = qpos.flatten()

                    if qpos.shape[0] != len(right_arm_ids):
                        logger.log_warning(
                            f"Qpos shape mismatch: got {qpos.shape[0]}, expected {len(right_arm_ids)}"
                        )
                        continue

                    qpos = qpos.unsqueeze(0).expand(self.num_envs, -1)

                    action = init_qpos.clone()
                    action[:, right_arm_ids] = qpos
                    action[:, right_eef_ids] = gripper_state.unsqueeze(0).expand(
                        self.num_envs, -1
                    )
                    action_list.append(action)

        logger.log_info(
            f"Generated {len(action_list)} demo actions for stacking blocks"
        )
        return action_list

    def is_task_success(self, **kwargs) -> torch.Tensor:
        """Determine if the task is successfully completed. This is mainly used in the data generation process
        of the imitation learning.

        The task is successful if:
        1. Block2 is stacked on top of Block1
        2. Both blocks haven't fallen over

        Args:
            **kwargs: Additional arguments for task-specific success criteria.

        Returns:
            torch.Tensor: A boolean tensor indicating success for each environment in the batch.
        """
        try:
            block1 = self.sim.get_rigid_object("block_1")
            block2 = self.sim.get_rigid_object("block_2")
        except Exception as e:
            logger.log_warning(f"Block1 or Block2 not found: {e}, returning False.")
            return torch.zeros(self.num_envs, dtype=torch.bool, device=self.device)

        # Get poses
        block1_pose = block1.get_local_pose(to_matrix=True)
        block2_pose = block2.get_local_pose(to_matrix=True)

        # Extract positions
        block1_pos = block1_pose[:, :3, 3]  # (num_envs, 3)
        block2_pos = block2_pose[:, :3, 3]

        # Check if blocks haven't fallen
        block1_fallen = self._is_fall(block1_pose)
        block2_fallen = self._is_fall(block2_pose)

        # Block2 should be on top of block1
        expected_block2_pos = torch.stack(
            [
                block1_pos[:, 0],
                block1_pos[:, 1],
                block1_pos[:, 2] + 0.05,  # block1 z + block height
            ],
            dim=1,
        )

        # Tolerance
        eps = torch.tensor(
            [0.025, 0.025, 0.012], dtype=torch.float32, device=self.device
        )

        # Check if block2 is within tolerance of expected position
        position_diff = torch.abs(block2_pos - expected_block2_pos)  # (num_envs, 3)
        within_tolerance = torch.all(
            position_diff < eps.unsqueeze(0), dim=1
        )  # (num_envs,)

        # Task succeeds if blocks are stacked correctly and haven't fallen
        success = within_tolerance & ~block1_fallen & ~block2_fallen

        return success

    def _is_fall(self, pose: torch.Tensor) -> torch.Tensor:
        # Extract z-axis from rotation matrix (last column, first 3 elements)
        pose_rz = pose[:, :3, 2]
        world_z_axis = torch.tensor([0, 0, 1], dtype=pose.dtype, device=pose.device)

        # Compute dot product for each batch element
        dot_product = torch.sum(pose_rz * world_z_axis, dim=-1)  # Shape: (batch_size,)

        # Clamp to avoid numerical issues with arccos
        dot_product = torch.clamp(dot_product, -1.0, 1.0)

        # Compute angle and check if fallen
        angle = torch.arccos(dot_product)
        return angle >= torch.pi / 4
