from enum import Enum
import time
import krpc
from pid import PID
conn = krpc.connect(name='SpaceX Landing')
if not conn.space_center:
    print("No active vessel found.")
    exit()

# control = vessel.control
vessel = conn.space_center.active_vessel
celestial = vessel.orbit.body # 当前天体
# 创建 PID 控制器
# 外环：高度 -> 目标垂直速度（允许负值）
altitude_pid = PID(kp=0.3, ki=0.03, kd=0.05, integral_limit=2.0, alpha=0.8, output_limit=(-10.0, 10.0))

# 内环：垂直速度 -> 节流阀（0~1）
speed_pid = PID(kp=0.06, ki=0.05, kd=0.05, integral_limit=10.0, alpha=0.5, output_limit=(0.0, 1.0))

target_altitude = 30  # 目标高度，单位 m

# 状态机定义：高于 100m 为滑翔（GLIDE），100m 及以下为减速着陆（LAND）
class STATE(str, Enum):
    GLIDE = "GLIDE"
    LAND = "LAND"

current_state = STATE.GLIDE

last_time = time.monotonic()
loop_dt = 0.1

# 进入减速着陆时使用的目标垂直速度（保持原有行为）
target_vertical_speed = -3

try:
    while True:
        now = time.monotonic()
        dt = now - last_time
        last_time = now

        flight = vessel.flight(vessel.orbit.body.reference_frame)
        current_altitude = flight.surface_altitude
        current_vertical_speed = flight.vertical_speed

        # 计算转换距离
        available_thrust = vessel.available_thrust # 可用推力
        gav = celestial.surface_gravity # 重力加速度
        twr = available_thrust / (vessel.mass * gav) # 推重比
        if twr > 1.0:
            transition_distance = (current_vertical_speed**2) / (2 * gav * (twr - 1)) # 米
        else:
            transition_distance = float('inf')  # 无法减速，设为无穷大

        # 状态判断：大于 200m 为滑翔（无需控制，油门 0），否则为减速着陆（使用现有内环 PID）
        if current_state == STATE.GLIDE: # 禁止反向切换
            new_state = STATE.GLIDE if current_altitude > transition_distance else STATE.LAND

        if new_state != current_state and new_state is not None:
            print(f"状态切换: {current_state} -> {new_state}")
            current_state = new_state
            # 在状态切换时重置 PID，避免积分风up
            try:
                altitude_pid.reset()
            except Exception:
                pass
            try:
                speed_pid.reset()
            except Exception:
                pass

        if current_state == STATE.GLIDE:
            # 滑翔：不做具体控制，油门设为 0
            vessel.control.throttle = 0.0
            print(f"当前高度: {current_altitude:.2f} m, 预计点火距离: {transition_distance:.2f} m, 状态: GLIDE, 油门: 0.00"," "*10, end='\r')
        else:
            # 内环：垂直速度误差 -> 节流阀
            speed_error = target_vertical_speed - current_vertical_speed
            throttle = speed_pid.update(speed_error, dt)

            vessel.control.throttle = throttle
            print(f"当前高度: {current_altitude:.2f} m, 预计点火距离: {transition_distance:.2f} m, 目标速: {target_vertical_speed:.2f} m/s, 速: {current_vertical_speed:.2f} m/s, 节流阀: {throttle:.2f}"," "*10, end='\r')

        # 保持大致的控制周期
        sleep_time = loop_dt - (time.monotonic() - now)
        if sleep_time > 0:
            time.sleep(sleep_time)

except KeyboardInterrupt:
    vessel.control.throttle = 0.0
    print("已停止控制，节流设为 0。")