# ----------------------------------------------------------------------------
# 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, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# ----------------------------------------------------------------------------

import torch
from typing import Dict, Optional

from embodichain.lab.gym.envs import EmbodiedEnv, EmbodiedEnvCfg
from embodichain.lab.gym.utils.registration import register_env
from embodichain.utils import logger

from embodichain.lab.gym.envs.tasks.tableware.base_agent_env import BaseAgentEnv
from embodichain.lab.gym.envs.tasks.tableware.pour_water.action_bank import (
    PourWaterActionBank,
)

__all__ = ["PourWaterEnv", "PourWaterAgentEnv"]


@register_env("PourWater-v3", max_episode_steps=600)
class PourWaterEnv(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 the current task.

        Returns:
            list: A list of demo actions generated by the task.
        """
        logger.log_info("Create demo action list for PourWaterTask.")

        if getattr(self, "action_config") is not None:
            self._init_action_bank(PourWaterActionBank, self.action_config)
            action_list = self.create_expert_demo_action_list(*args, **kwargs)
        else:
            logger.log_error("No action_config found in env, please check again.")

        if action_list is None:
            return action_list

        logger.log_info(
            f"Demo action list created with {len(action_list)} steps.", color="green"
        )
        return action_list

    def create_expert_demo_action_list(self, **kwargs):
        """
        Create an expert demonstration action list using the action bank.

        This function generates a trajectory based on expert knowledge, mapping joint and end-effector
        states to the required action format for the environment and robot type.

        Args:
            **kwargs: Additional keyword arguments.

        Returns:
            list: A list of actions, each containing joint positions ("qpos").
        """

        if hasattr(self, "action_bank") is False or self.action_bank is None:
            logger.log_error(
                "Action bank is not initialized. Cannot create expert demo action list."
            )

        ret = self.action_bank.create_action_list(
            self, self.graph_compose, self.packages
        )

        if ret is None:
            logger.log_warning("Failed to generate expert demo action list.")
            return None

        # TODO: to be removed, need a unified interface in robot class
        left_arm_joints = self.robot.get_joint_ids(name="left_arm")
        right_arm_joints = self.robot.get_joint_ids(name="right_arm")
        left_eef_joints = self.robot.get_joint_ids(name="left_eef", remove_mimic=True)
        right_eef_joints = self.robot.get_joint_ids(name="right_eef", remove_mimic=True)

        total_traj_num = ret[list(ret.keys())[0]].shape[-1]
        actions = torch.zeros(
            (total_traj_num, self.num_envs, self.robot.dof), dtype=torch.float32
        )

        for key, joints in [
            ("left_arm", left_arm_joints),
            ("left_eef", left_eef_joints),
            ("right_arm", right_arm_joints),
            ("right_eef", right_eef_joints),
        ]:
            if key in ret:
                # TODO: only 1 env supported now
                actions[:, 0, joints] = torch.as_tensor(ret[key].T, dtype=torch.float32)

        return actions

    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.

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

        Returns:
            torch.Tensor: A boolean tensor indicating success for each environment in the batch.
        """

        bottle = self.sim.get_rigid_object("bottle")
        cup = self.sim.get_rigid_object("cup")

        bottle_final_xpos = bottle.get_local_pose(to_matrix=True)
        cup_final_xpos = cup.get_local_pose(to_matrix=True)

        bottle_ret = self._is_fall(bottle_final_xpos)
        cup_ret = self._is_fall(cup_final_xpos)

        return ~(bottle_ret | cup_ret)

    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


@register_env("PourWaterAgent-v3", max_episode_steps=600)
class PourWaterAgentEnv(BaseAgentEnv, PourWaterEnv):
    def __init__(self, cfg: EmbodiedEnvCfg = None, **kwargs):
        super().__init__(cfg, **kwargs)
        super()._init_agents(**kwargs)

    def reset(self, seed: Optional[int] = None, options: Optional[Dict] = None):
        obs, info = super().reset(seed=seed, options=options)
        super().get_states()
        return obs, info
