import time
import krpc

conn = krpc.connect(name='SpaceX Landing')
if not conn.space_center:
    print("No active vessel found.")
    exit()
vessel = conn.space_center.active_vessel

# 复用之前的 PID 类
class PID:
    def __init__(self, kp, ki, kd, integral_limit=1.0):
        self.kp, self.ki, self.kd = kp, ki, kd
        self.prev_error = 0
        self.integral = 0
        self.integral_limit = integral_limit

    def update(self, error, dt):
        self.integral += error * dt
        self.integral = max(-self.integral_limit, min(self.integral_limit, self.integral))
        derivative = (error - self.prev_error) / dt
        self.prev_error = error
        return (self.kp * error) + (self.ki * self.integral) + (self.kd * derivative)

# 1. 初始设置
vessel.auto_pilot.engage()
vessel.auto_pilot.reference_frame = vessel.surface_velocity_reference_frame
vessel.auto_pilot.target_direction = (0, -1, 0) # 始终指向速度反方向（Retrograde）

speed_pid = PID(kp=0.15, ki=0.05, kd=0.1, integral_limit=0.5)

# 节流阀平滑处理相关变量
last_throttle = 0.0
alpha = 0.3  # 平滑系数。越大响应越快，越小越平滑

print("着陆程序启动...")

vessel.control.throttle = 0.99
time.sleep(0.1)
vessel.control.throttle = 0  # 自检


while True:
    flight = vessel.flight(vessel.orbit.body.reference_frame)
    # 获取距地高度（考虑到雷达高度计）
    alt = max(0, flight.surface_altitude - 9.50) 
    vel = flight.vertical_speed
    
    # --- 姿态控制逻辑优化 ---
    # 防止低速时由于速度方向频繁改变导致的姿态抖动
    if alt < 30 or abs(vel) < 0.5:
        # 离地很近或速度很慢时，强制垂直向上
        vessel.auto_pilot.reference_frame = vessel.surface_reference_frame
        vessel.auto_pilot.target_direction = (1, 0, 0) # 向天
    else:
        # 正常减速阶段，锁定速度反方向（Retrograde）
        vessel.auto_pilot.reference_frame = vessel.surface_velocity_reference_frame
        vessel.auto_pilot.target_direction = (0, -1, 0)

    # 打开刹车
    if alt < 10000 and vessel.available_thrust > 0 and vessel.control.brakes == False:
        vessel.control.brakes = True


    # 2. 动态速度规划 (关键逻辑)
    # 这是一个简化的着陆曲线：目标速度 = -sqrt(2 * 理想加速度 * 高度) 
    # 我们希望在靠近地面时速度减小到 -2 m/s
    target_vel = -10000.0 # 高度较高时先全速下降
    if alt < 5000:
        if alt > 500:
            target_vel = -min(300.0, 0.5 * alt) # 高空较快下降
        elif alt > 200:
            target_vel = -min(100.0, 0.5 * alt) # 高空较快下降
        elif alt > 30:
            # 打开刹车
            if vessel.control.brakes == True:
                vessel.control.brakes = False
            target_vel = -20.0 # 中高度恒定下降速度
        else:
            target_vel = -3.0 # 接近地面时恒定慢速
        
    
    # 3. 触地检测
    if alt < 1.0 and abs(vel) < 0.5:
        vessel.control.throttle = 0
        vessel.control.sas = True 
        print("已着陆！")
        break

    # 4. 自动部署脚架
    if alt < 500:
        vessel.control.gear = True

    # 5. PID 计算
    dt = 0.1
    error = target_vel - vel
    raw_throttle = speed_pid.update(error, dt)
    
    # --- 简易低延迟滤波 (EMA) ---
    # 使用 0.4 的 alpha 值兼顾了响应速度和稳定性
    throttle = (alpha * raw_throttle) + ((1.0 - alpha) * last_throttle)
    
    # 限制并在执行前更新 last_throttle
    throttle = max(0.0, min(1.0, throttle))
    last_throttle = throttle
    
    vessel.control.throttle = throttle
    
    print(f"高度: {alt:.2f}m | 目标速: {target_vel:.1f}m/s | 实际速: {vel:.1f}m/s | 节流阀: {throttle:.2f}")
    time.sleep(dt)