# ----------------------------------------------------------------------------
# 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 embodichain.lab.gym.envs import EmbodiedEnv, EmbodiedEnvCfg
from embodichain.lab.gym.utils.registration import register_env
from embodichain.utils import logger

__all__ = ["SimpleTaskEnv"]


@register_env("SimpleTask-v1", max_episode_steps=600)
class SimpleTaskEnv(EmbodiedEnv):
    """A demo environment with sinusoidal trajectory

    Args:
        EmbodiedEnv (_type_): _description_
    """

    def __init__(self, cfg: EmbodiedEnvCfg = None, **kwargs):
        super().__init__(cfg, **kwargs)

    def create_demo_action_list(self, *args, **kwargs):
        """
        Create a demonstration action list for the current task.

        This demo creates a simple sinusoidal trajectory for the robot joints.

        Returns:
            list: A list of demo actions generated by the task.
        """
        action_list = []
        num_steps = 100

        # Get initial pose
        init_pose = self.robot.get_qpos()  # shape: (num_envs, num_joints)

        # Create a sinusoidal trajectory
        for i in range(num_steps):
            # Calculate phase for sinusoidal motion
            t = i / num_steps  # 0 to 1
            phase = torch.full(
                (init_pose.shape[0],), t * 2 * torch.pi, device=self.device
            )  # repeat for num_envs

            # Create sinusoidal offsets for each joint
            # Joint 0: horizontal movement
            # Joint 1: vertical movement
            # Other joints: smaller oscillations
            offset = torch.zeros_like(
                init_pose, dtype=torch.float32, device=self.device
            )
            offset[:, 0] = torch.sin(phase) * 0.3  # ±0.3 rad
            offset[:, 1] = torch.cos(phase) * 0.2  # ±0.2 rad
            offset[:, 2] = torch.sin(phase * 2) * 0.1  # ±0.1 rad, double frequency

            # Add small random variation to make it more natural
            noise = (torch.rand_like(init_pose, device=self.device) - 0.5) * 0.02

            # Compute action
            action = init_pose + offset + noise

            # Clamp to joint limits if available
            if hasattr(self.robot.body_data, "qpos_limits"):
                qpos_limits = self.robot.body_data.qpos_limits[0]  # (num_joints, 2)
                action = torch.clamp(action, qpos_limits[:, 0], qpos_limits[:, 1])

            action_list.append(action)

        logger.log_info(
            f"Generated {len(action_list)} demo actions with sinusoidal trajectory"
        )
        return action_list
