"""
笛卡尔空间强化学习环境
"""
import gymnasium as gym
from gymnasium import spaces
import numpy as np
import logging

# 配置日志
logging.basicConfig(
    level=logging.INFO,
    format='%(asctime)s - %(name)s - %(levelname)s - %(message)s'
)
logger = logging.getLogger(__name__)

class CartesianSpaceEnv(gym.Env):
    """机械手臂在笛卡尔空间的运动规划环境"""
    metadata = {'render_modes': ['human', 'rgb_array'], 'render_fps': 50}

    def __init__(self, render_mode=None):
        # 定义动作空间：x, y, z轴的速度控制
        self.action_space = spaces.Box(
            low=np.array([-1.0, -1.0, -1.0]), 
            high=np.array([1.0, 1.0, 1.0]),
            shape=(3,),
            dtype=np.float32
        )

        # 定义观测空间：当前位置、目标位置、误差
        self.observation_space = spaces.Box(
            low=np.array([-np.inf]*9),
            high=np.array([np.inf]*9),
            dtype=np.float32
        )

        # 初始化环境参数
        self._target_pos = np.array([0.5, 0.5, 0.5])
        self._current_pos = np.array([0.0, 0.0, 0.0])
        self._dt = 0.01  # 时间步长
        self.render_mode = render_mode
        self.episode_step = 0
        self.episode_reward = 0.0
        
        # 测试场景支持
        self.test_scenarios = None
        self.current_scenario_idx = 0

    def _get_obs(self):
        # 获取当前观测值
        return np.concatenate([
            self._current_pos,
            self._target_pos,
            self._target_pos - self._current_pos
        ])

    # 如果使用旧版reset方法，需要更新为新的API格式
    def reset(self, seed=None, options=None):
        super().reset(seed=seed)
        
        # 重置环境参数
        self.episode_step = 0
        self.episode_reward = 0.0
        
        # 随机初始化当前位置（添加随机性）
        self._current_pos = self.np_random.uniform(low=-0.5, high=0.5, size=(3,)).astype(np.float32)
        
        # 如果有测试场景，使用场景中的目标位置
        if self.test_scenarios and self.current_scenario_idx < len(self.test_scenarios):
            self._target_pos = np.array(self.test_scenarios[self.current_scenario_idx]['target_position'], dtype=np.float32)
            self.current_scenario_idx += 1
        else:
            # 否则随机生成目标位置
            self._target_pos = self.np_random.uniform(low=0.0, high=1.0, size=(3,)).astype(np.float32)
        
        # 获取初始观测
        observation = self._get_obs()
        
        # 创建info字典
        info = {"reset": "environment reset to initial state"}
        
        return observation, info

    # 更新step方法以返回新的五元组格式
    def step(self, action):
        # 执行一步环境交互
        self._current_pos += action * self._dt
        observation = self._get_obs().astype(np.float32)

        # 计算奖励
        distance = np.linalg.norm(self._target_pos - self._current_pos)
        reward = -distance
        
        # 累计episode统计
        self.episode_step += 1
        self.episode_reward += reward

        # 判断是否终止
        terminated = distance < 0.01
        truncated = False
        info = {}
        
        # 记录每一步的信息
        logger.debug(f"Step {self.episode_step}: Position={self._current_pos}, Reward={reward:.4f}, Distance={distance:.4f}")
        
        # 如果episode结束，记录总结
        if terminated or truncated:
            logger.info(f"Episode end: Steps={self.episode_step}, Total reward={self.episode_reward:.4f}, Final distance={distance:.4f}")
            
        return observation, reward, terminated, truncated, info

    def render(self):
        # 渲染环境
        if self.render_mode == "human":
            print(f"Current position: {self._current_pos}, Target position: {self._target_pos}")

    def close(self):
        # 关闭环境
        pass

    def load_test_scenarios(self, scenarios):
        """
        加载测试场景
        
        参数:
            scenarios: 测试场景列表
        """
        self.test_scenarios = scenarios
        self.current_scenario_idx = 0
