import gymnasium as gym
from XPlaneGym import XPlaneEnv

# 尝试创建环境实例
env = gym.make("XPlane-v0")
print("环境创建成功！")
env.close()

import gymnasium as gym
from XPlaneGym import XPlaneEnv
import numpy as np

# 创建环境，配置为连续动作空间
env = gym.make(
    "XPlane-v0",
    ip='127.0.0.1',           # X-Plane IP地址
    port=49000,               # UDP端口
    timeout=1.0,              # 通信超时
    pause_delay=0.05,         # 动作执行延迟
    starting_latitude=37.558, # 初始纬度（首尔金浦国际机场附近）
    starting_longitude=126.790,
    starting_altitude=3000.0, # 初始高度3000米
    starting_velocity=100.0,  # 初始速度100节
    starting_pitch_range=10.0,
    starting_roll_range=20.0,
    random_desired_state=True,
    desired_pitch_range=5.0,
    desired_roll_range=10.0,
    continuous_actions=True   # 使用连续动作空间
)

# 运行几个随机步骤来验证环境工作正常
obs, info = env.reset()
print(f"初始观察: {obs}")

for episode in range(3):
    obs, info = env.reset()
    total_reward = 0
    done = False
    step = 0
    
    while not done and step < 100:
        # 随机选择动作
        action = env.action_space.sample()
        obs, reward, terminated, truncated, info = env.step(action)
        total_reward += reward
        done = terminated or truncated
        step += 1
    
    print(f"Episode {episode + 1}: 奖励 = {total_reward:.2f}, 步数 = {step}")

env.close()
print("测试完成！")

import gymnasium as gym
import numpy as np
from XPlaneGym import XPlaneEnv

class MyCustomXPlaneEnv(XPlaneEnv):
    """自定义飞行控制环境"""
    
    def __init__(self, **kwargs):
        super().__init__(**kwargs)
        
    def _compute_reward(self, state, desired_state):
        """
        自定义奖励函数
        可以根据任务需求设计不同的奖励机制
        """
        # 计算状态差异
        pitch_diff = abs(state[2] - desired_state[0])
        roll_diff = abs(state[3] - desired_state[1])
        
        # 高度保持奖励
        altitude_reward = -abs(state[1] - 3000.0) * 0.01
        
        # 姿态控制奖励（越接近目标姿态，奖励越高）
        attitude_reward = -0.5 * (pitch_diff + roll_diff)
        
        # 达到目标的额外奖励
        if pitch_diff < 2.0 and roll_diff < 2.0:
            attitude_reward += 10.0
            
        return altitude_reward + attitude_reward
    
    def _get_observation(self):
        """
        自定义观察空间
        可以选择包含或排除特定的状态分量
        """
        # 使用自定义的状态表示
        obs = np.array([
            self.current_state.altitude / 10000.0,      # 归一化高度
            self.current_state.velocity / 500.0,       # 归一化速度
            self.current_state.pitch / 90.0,           # 归一化俯仰角
            self.current_state.roll / 180.0,           # 归一化滚转角
            self.current_state.heading / 360.0,       # 归一化航向
        ], dtype=np.float32)
        
        return obs

# 注册自定义环境
gym.register(
    id='XPlane-custom-v0',
    entry_point='__main__:MyCustomXPlaneEnv',
    kwargs={...}  # 传递给构造函数的参数
)

class XPlaneEnv(gym.Env):
    """X-Plane强化学习环境主类"""
    
    metadata = {'render_modes': ['human']}
    
    def __init__(self, ...):
        super().__init__()
        
        # 初始化UDP通信
        self.socket = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
        self.socket.settimeout(timeout)
        
        # 初始化环境参数
        self.ip = ip
        self.port = port
        # ... 其他参数初始化
        
        # 定义观察空间和动作空间
        self._define_spaces()
    
    def reset(self, seed=None, options=None):
        """重置环境到初始状态"""
        super().reset(seed=seed)
        
        # 随机设置初始状态
        self._set_initial_conditions()
        
        # 等待环境稳定
        self._wait_for_stable_state()
        
        return self._get_observation(), {}
    
    def step(self, action):
        """执行一个时间步"""
        # 发送控制指令到X-Plane
        self._send_control(action)
        
        # 等待状态更新
        time.sleep(self.pause_delay)
        
        # 获取新状态
        state = self._get_current_state()
        
        # 计算奖励
        reward = self._compute_reward(state, self.desired_state)
        
        # 检查是否终止
        terminated = self._check_termination(state)
        
        # 获取观察
        obs = self._get_observation(state)
        
        info = {'state': state, 'desired': self.desired_state}
        
        return obs, reward, terminated, False, info

def _send_control(self, action):
    """发送控制指令到X-Plane"""
    if self.continuous_actions:
        # 连续动作：直接使用浮点值
        pitch_input = action[0]
        roll_input = action[1]
        yaw_input = action[2] if len(action) > 2 else 0.0
    else:
        # 离散动作：映射到具体控制量
        action_mapping = {
            0: [0, 0, 0],      # 无操作
            1: [0.1, 0, 0],    # 俯仰向上
            2: [-0.1, 0, 0],   # 俯仰向下
            3: [0, 0.1, 0],    # 向左滚转
            4: [0, -0.1, 0],   # 向右滚转
            # ... 其他动作映射
        }
        pitch_input, roll_input, yaw_input = action_mapping[action]
    
    # 构造并发送UDP数据包
    control_data = self._build_control_packet(pitch_input, roll_input, yaw_input)
    self.socket.sendto(control_data, (self.ip, self.port))

def _receive_state(self):
    """从X-Plane接收状态数据"""
    try:
        data, addr = self.socket.recvfrom(4096)
        state = self._parse_state_packet(data)
        return state
    except socket.timeout:
        # 超时处理：返回最后已知状态或默认值
        return self._get_default_state()

def _parse_state_packet(self, data):
    """解析X-Plane状态数据包"""
    # X-Plane UDP数据包格式：
    # 前5个字节是头部信息，之后是多个数据索引-值对
    
    # 跳过头部
    offset = 5
    
    state = XPlaneState()
    
    while offset < len(data) - 8:
        # 读取数据索引
        index = int.from_bytes(data[offset:offset+4], byteorder='big')
        offset += 4
        
        # 读取数据值
        value = struct.unpack('f', data[offset:offset+4])[0]
        offset += 4
        
        # 根据索引提取所需数据
        # X-Plane数据索引定义请参考官方文档
        if index == 0:  # 海平面气压高度
            state.altitude = value
        elif index == 17:  # 真空速
            state.velocity = value
        elif index == 11:  # 俯仰角
            state.pitch = value
        elif index == 12:  # 滚转角
            state.roll = value
        elif index == 13:  # 航向角
            state.heading = value
        # ... 其他状态字段
    
    return state

import gymnasium as gym
from stable_baselines3 import PPO
from XPlaneGym import XPlaneEnv

# 创建姿态控制环境
env = gym.make(
    "XPlane-v0",
    continuous_actions=True,
    starting_altitude=3000.0,
    starting_velocity=150.0,
    random_desired_state=True,
    desired_pitch_range=15.0,
    desired_roll_range=30.0,
)

# 使用PPO算法进行训练
model = PPO(
    "MlpPolicy",
    env,
    learning_rate=3e-4,
    n_steps=2048,
    batch_size=64,
    n_epochs=10,
    gamma=0.99,
    verbose=1
)

# 训练100000步
model.learn(total_timesteps=100000)

# 保存模型
model.save("ppo_xplane_attitude")

# 评估模型
eval_env = gym.make("XPlane-v0", continuous_actions=True)
obs, _ = eval_env.reset()
rewards = []

for _ in range(1000):
    action, _ = model.predict(obs, deterministic=True)
    obs, reward, terminated, truncated, _ = eval_env.step(action)
    rewards.append(reward)
    
    if terminated or truncated:
        obs, _ = eval_env.reset()

print(f"平均奖励: {sum(rewards)/len(rewards):.2f}")
eval_env.close()

class PathFollowingEnv(XPlaneEnv):
    """路径跟踪环境"""
    
    def __init__(self, waypoints=None, **kwargs):
        super().__init__(**kwargs)
        self.waypoints = waypoints or self._default_waypoints()
        self.current_waypoint_index = 0
    
    def _default_waypoints(self):
        """默认航点序列（经度，纬度，高度）"""
        return [
            (126.790, 37.558, 3000),  # 起点
            (126.850, 37.600, 3500),
            (126.900, 37.550, 3000),
            (126.850, 37.500, 2500),
        ]
    
    def _compute_reward(self, state, desired_state):
        """路径跟踪奖励函数"""
        # 当前航点
        target = self.waypoints[self.current_waypoint_index]
        
        # 计算位置误差
        lat_diff = abs(state.latitude - target[1])
        lon_diff = abs(state.longitude - target[0])
        alt_diff = abs(state.altitude - target[2])
        
        position_error = -(lat_diff + lon_diff) * 100 - alt_diff * 0.1
        
        # 检查是否到达当前航点
        if lat_diff < 0.001 and lon_diff < 0.001 and alt_diff < 50:
            self.current_waypoint_index = (self.current_waypoint_index + 1) % len(self.waypoints)
            position_error += 100  # 到达奖励
        
        # 姿态保持奖励
        attitude_penalty = -abs(state.pitch - desired_state.pitch) * 0.5
        attitude_penalty -= abs(state.roll - desired_state.roll) * 0.5
        
        return position_error + attitude_penalty