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

class PID:
    def __init__(self, kp, ki, kd, integral_limit=10.0, alpha=1.0, output_limit=(0.0, 1.0)):
        self.kp = kp
        self.ki = ki
        self.kd = kd
        self.prev_error = 0.0
        self.integral = 0.0
        self.integral_limit = integral_limit
        self.alpha = alpha
        self.prev_output = 0.0
        self.output_limit = output_limit

    def update(self, error, dt):
        if dt <= 0:
            dt = 1e-6
        self.integral += error * dt
        self.integral = max(-self.integral_limit, min(self.integral_limit, self.integral))
        derivative = (error - self.prev_error) / dt
        output = self.kp * error + self.ki * self.integral + self.kd * derivative
        self.prev_error = error
        # 一阶滤波
        output = self.alpha * output + (1 - self.alpha) * self.prev_output
        self.prev_output = output

        # 输出限制
        output = max(self.output_limit[0], min(self.output_limit[1], output))

        return output


# control = vessel.control
vessel = conn.space_center.active_vessel

# 创建 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.12, ki=0.05, kd=0.05, integral_limit=50.0, alpha=0.5, output_limit=(0.0, 1.0))

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

last_time = time.monotonic()
loop_dt = 0.1

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

        # 外环：高度误差 -> 目标垂直速度
        altitude_error = target_altitude - current_altitude
        target_vertical_speed = altitude_pid.update(altitude_error, dt)
        target_vertical_speed = max(-10.0, min(10.0, target_vertical_speed))

        # 内环：垂直速度误差 -> 节流阀
        speed_error = target_vertical_speed - current_vertical_speed
        throttle = speed_pid.update(speed_error, dt)
        throttle = max(0.0, min(1.0, throttle))

        vessel.control.throttle = throttle
        print(f"当前高度: {current_altitude:.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。")