import time
import krpc
import math


G_EARTH = 9.80665

def get_surface_speed():
    conn = krpc.connect(name='Vessel speed')
    vessel = conn.space_center.active_vessel
    obt_frame = vessel.orbit.body.non_rotating_reference_frame
    srf_frame = vessel.orbit.body.reference_frame

    while True:
        obt_speed = vessel.flight(obt_frame).speed
        srf_speed = vessel.flight(srf_frame).speed
        print('Orbital speed = %.1f m/s, Surface speed = %.1f m/s' %
              (obt_speed, srf_speed))
        time.sleep(1)

def get_surface_velocity():
    conn = krpc.connect(name='Orbital speed')
    vessel = conn.space_center.active_vessel
    ref_frame = conn.space_center.ReferenceFrame.create_hybrid(
        position=vessel.orbit.body.reference_frame,
        rotation=vessel.surface_reference_frame)

    while True:
        velocity = vessel.flight(ref_frame).velocity
        print('Surface velocity = (%.1f, %.1f, %.1f)' % velocity)
        time.sleep(1)


def surface_prograde():
    conn = krpc.connect(name='Surface prograde')
    vessel = conn.space_center.active_vessel
    ap = vessel.auto_pilot

    ap.reference_frame = vessel.surface_velocity_reference_frame
    ap.target_direction = (0, 1, 0)
    ap.engage()
    ap.wait()
    print('Auto-pilot now pointing prograde in surface reference frame')
    ap.disengage()

def surface_retrograde():
    conn = krpc.connect(name='Surface retrograde')
    vessel = conn.space_center.active_vessel
    ap = vessel.auto_pilot

    ap.reference_frame = vessel.surface_velocity_reference_frame
    ap.target_direction = (0, -1, 0)
    ap.engage()
    ap.wait()
    print('Auto-pilot now pointing retrograde in surface reference frame')
    ap.disengage()

def vessel_available_thrust():
    conn = krpc.connect(name='Vessel available thrust')
    vessel = conn.space_center.active_vessel

    while True:
        thrust = vessel.available_thrust
        print('Vessel available thrust = %.1f kN' % (thrust / 1000))
        time.sleep(1)

def get_body_g(local_vessel):
    try:
        body = local_vessel.orbit.body
        g = float(getattr(body, 'surface_gravity', G_EARTH))
        if g > 0:
            return g
    except Exception:
        pass
    try:
        g = float(getattr(local_vessel.flight(), 'surface_gravity', G_EARTH))
        if g > 0:
            return g
    except Exception:
        pass
    return G_EARTH

def auto_land_predictive(
        bottom_clearance_margin=12,
        thrust_discount_wait=0.85,
        host='localhost',
        port=50000,
        target_altitude=1.0,
        target_vspeed=-1.0,
        safety_margin=3.0,
        verbose=True):
    """预测式自动着陆（suicide burn 思路）。

    关键点：
    1) 速度从接口获取：surface speed 用 flight(surface_frame).speed（标量），垂直速度用 flight(surface_frame).vertical_speed。
       不再用“高度差/时间差”估算。
    2) 不再脚本激活即点火；先计算 braking distance，根据可用推力决定何时点火（countdown）。

    近似模型（忽略空气阻力、推力变化、TWR变化导致的积分误差）：
      a_max_up = max(0, available_thrust/mass - g)
      braking_distance ≈ (v^2 - vf^2) / (2*a_max_up) + safety_margin
    当 remaining_alt <= braking_distance 时开始点火。

    燃烧阶段按剩余距离反推所需净向上加速度：
      a_req_up = max(0, (v^2 - vf^2) / (2*remaining_alt))
      throttle = (a_req_up + g) * mass / available_thrust

    注：这里 v 取“相对地表速度”的标量（surface speed），更贴近你提到的 get_surface_speed。
        若需要更精细，可改成取 surface_velocity 向量的径向分量/或用垂直速度替代。

        bottom_clearance_margin = 12  # m：按载具尺寸调整（越大越早开始末端减速）

    参数说明：
      bottom_clearance_margin (float):
        - 单位：米(m)
        - 说明：从雷达（bedrock_altitude）读取的高度中减去该几何余量以近似载具最低点到地面的距离。
          用于避免脚本在实际着陆时因底部结构接触地面而过早判断已着陆。值越大会更早开始末端减速。
        - 典型值：10~25，根据载具底部高度调整。

      thrust_discount_wait (float):
        - 单位：无（比例系数）
        - 说明：在 WAIT 阶段估算可用推力时对 vessel.available_thrust 的折减因子。
          会用 avail_thrust * thrust_discount_wait 来保守估算用于决定何时点火（补偿推力建立/响应延迟）。
        - 取值范围：0.0~1.0；越小越保守（更早点火）。默认 0.85。

      host (str) / port (int):
        - 说明：连接 kRPC 服务的地址和 RPC 端口（用于 krpc.connect）。一般本地运行时用 'localhost' 和 50000。

      target_altitude (float):
        - 单位：米(m)
        - 说明：目标的“触地点高度基准”，脚本在计算剩余高度时会把 radar_alt - bottom_clearance_margin 与此值比较。
          通常设为想要在接触地面时的质心高度（或 0/地面高度上方的偏移）。
        - 典型值：0~5（例如 3 表示目标在底部余量之上 3m 处被视为触地）。

      target_vspeed (float):
        - 单位：m/s（竖直速度，负值表示向下）
        - 说明：期望的着陆竖直速度（向下为负）。脚本会把该值取绝对值作为目标的向下速度（vf_down）。
          在计算刹车距离和 PID/前馈控制里以该目标速度为参考。
        - 典型值：-0.5 ~ -1.5（更接近 0 的值更“柔和”但更难实现）。默认为 -1.0。

      safety_margin (float):
        - 单位：米(m)
        - 说明：在刹车距离上再加的安全余量，用于补偿模型误差（如空气阻力、推力随时间变化等）。
        - 典型值：几米（例如 3~10）。默认 5.

      verbose (bool):
        - 说明：是否打印并记录运行时信息到返回的 status['logs']，用于调试或观察倒计时/燃烧过程。

    返回值：
      返回一个 dict 格式的状态描述，包含 'result' 字段（'failed' / 'landed' / 'touchdown_pending'）和 'logs'（运行时消息列表）。
    """

    conn = krpc.connect(name='auto_land_predictive', address=host, rpc_port=port)
    sc = conn.space_center
    vessel = sc.active_vessel
    if vessel is None:
        conn.close()
        raise RuntimeError('未找到活动载具')

    status = {'result': 'failed', 'logs': []}

    # 低空/接地参数（用于“最后几米慢下来”）
    touchdown_vspeed_tol = 0.6  # m/s：认为“软着陆”的竖直速度容差
    low_altitude_pid_switch = 10.0  # m：低于该高度切换到竖直速度控制
    vspeed_p_gain = 1.2  # 比例增益：越大越积极（建议 0.6~2.0）
    low_altitude_ff_gain = 0.3  # 低空前馈增益：按“需要的减速距离/剩余距离”做额外加速度
    # 由于 surface_altitude 近似是“载具质心高度”，末端会偏乐观。
    # 用 bedrock_altitude（雷达高度）并减去一个底部几何余量，近似得到“最低点离地高度”。

    # 使用 body 的 reference_frame 读取 surface speed / vertical_speed（对应 get_surface_speed 思路）
    try:
        surface_frame = vessel.orbit.body.reference_frame
    except Exception:
        surface_frame = None

    # autopilot 指向 surface retrograde
    try:
        ap = vessel.auto_pilot
        ap.reference_frame = vessel.surface_velocity_reference_frame
        ap.target_direction = (0, -1, 0)
        ap.engage()
    except Exception:
        ap = None

    def _read_state():
        """从接口读取状态：alt(surface_altitude), radar_alt(bedrock_altitude), surface_speed, vertical_speed, mass, available_thrust, g, situation_name"""
        flight = vessel.flight(surface_frame) if surface_frame is not None else vessel.flight()
        alt = float(getattr(flight, 'surface_altitude', 0.0))
        radar_alt = float(getattr(flight, 'bedrock_altitude', alt))
        srf_speed = float(getattr(flight, 'speed', 0.0))
        v_speed = float(getattr(flight, 'vertical_speed', 0.0))
        mass = float(getattr(vessel, 'mass', 0.0))
        avail_thrust = float(getattr(vessel, 'available_thrust', 0.0))
        g_local = get_body_g(vessel)
        situation_name = str(getattr(getattr(vessel, 'situation', None), 'name', ''))
        return alt, radar_alt, srf_speed, v_speed, mass, avail_thrust, g_local, situation_name

    try:
        # 初始确保不点火
        vessel.control.throttle = 0.0

        # 倒计时/等待点火
        last_log_t = 0.0
        while True:
            alt, radar_alt, srf_speed, v_speed, mass, avail_thrust, g_local, situation_name = _read_state()
            # 用“雷达高度-底部余量”近似最低点离地高度
            alt_eff = max(0.0, float(radar_alt) - float(bottom_clearance_margin))
            remaining_alt = max(0.0, alt_eff - float(target_altitude))

            # v：使用 surface speed（标量）作为“需要刹掉的速度”
            v = max(0.0, float(srf_speed))
            vf = max(0.0, abs(float(target_vspeed)))

            # 若推力不可用，直接点火尝试（避免永远等不到）
            avail_thrust_wait = float(avail_thrust) * float(thrust_discount_wait)
            if mass <= 0 or avail_thrust_wait <= 1e-6:
                a_max_up = 0.0
            else:
                a_max_up = max(0.0, avail_thrust_wait / mass - g_local)

            if a_max_up <= 1e-6:
                braking_dist = float('inf')
            else:
                braking_dist = max(0.0, (v * v - vf * vf) / (2.0 * a_max_up)) + float(safety_margin)

            countdown = None
            if v > 1e-6 and braking_dist != float('inf'):
                # 近似把“还剩多少距离”除以当前速度当作倒计时（粗略）
                countdown = (remaining_alt - braking_dist) / v

            now = time.time()
            if verbose and now - last_log_t > 0.2:
                last_log_t = now
                msg = (
                    f'WAIT alt={alt:.1f} radar={radar_alt:.1f} alt_eff={alt_eff:.1f} rem={remaining_alt:.1f} '
                    f'srf_speed={srf_speed:.2f} vs={v_speed:.2f} '
                    f'a_max_up={a_max_up:.2f} braking={braking_dist:.1f} thr_disc={thrust_discount_wait:.2f} '
                    f'countdown={(countdown if countdown is not None else float("nan")):.2f}'
                )
                status['logs'].append(msg)
                print(msg)

            # 点火条件：进入刹车距离 or 推力未知（inf）但高度很低
            should_ignite = (remaining_alt <= braking_dist + v_speed) or (braking_dist == float('inf') and remaining_alt < 200.0)
            if should_ignite:
                break

            # 确保仍在滑行
            time.sleep(0.05)

        # 燃烧闭环：根据 remaining_alt 反推需要的减速度 -> 油门
        while True:
            alt, radar_alt, srf_speed, v_speed, mass, avail_thrust, g_local, situation_name = _read_state()
            alt_eff = max(0.0, float(radar_alt) - float(bottom_clearance_margin))
            remaining_alt = max(0.0, alt_eff - float(target_altitude))

            # BURN 阶段以“竖直下降速度”为控制对象：下降为负，统一转成向下为正的标量 v_down
            v_down = max(0.0, -float(v_speed))
            vf_down = max(0.0, abs(float(target_vspeed)))  # 目标接地竖直速度（向下为正）

            # 终止条件：LANDED 可能会在仍有较大竖直速度时触发（弹跳/碰撞瞬间）
            # 因此要求：LANDED 且竖直速度足够小才结束
            is_landed = str(situation_name).strip().lower() == 'landed'
            if is_landed and v_down <= (vf_down + float(touchdown_vspeed_tol)):
                vessel.control.throttle = 0.0
                status['result'] = 'landed'
                status['logs'].append('检测到 vessel.situation=LANDED 且竖直速度满足软着陆阈值，着陆完成')
                if verbose:
                    print('检测到 vessel.situation=LANDED 且竖直速度满足软着陆阈值，着陆完成')
                break

            # 兜底：极低高度 + 速度足够小（防止 situation 迟迟不切换导致卡死）
            if remaining_alt <= max(0.2, float(target_altitude)) and v_down <= (vf_down + float(touchdown_vspeed_tol)):
                vessel.control.throttle = 0.0
                status['result'] = 'touchdown_pending'
                status['logs'].append('达到高度/速度阈值，但 vessel.situation 仍未变为 LANDED')
                if verbose:
                    print('达到高度/速度阈值，但 vessel.situation 仍未变为 LANDED')
                break

            # 低空策略：切换到“竖直速度跟踪”以确保最后几米把速度压到目标值
            if alt_eff <= float(low_altitude_pid_switch):
                # 低空：比例 + 前馈（用当前速度估算“需要多少距离才能刹到目标速度”）
                # a_p：把竖直速度拉向目标
                a_p = float(vspeed_p_gain) * max(0.0, (v_down - vf_down))
                # a_ff：基于能量/距离的前馈，防止在最后几米才突然补救
                if remaining_alt <= 1e-3:
                    a_ff = 0.0
                else:
                    a_ff = float(low_altitude_ff_gain) * max(0.0, (v_down * v_down - vf_down * vf_down) / (2.0 * remaining_alt))
                a_req_up = max(0.0, a_p + a_ff)
            else:
                # 中高空：用距离反推所需净向上加速度（使用竖直下降速度）
                if remaining_alt <= 1e-3:
                    a_req_up = 0.0
                else:
                    a_req_up = max(0.0, (v_down * v_down - vf_down * vf_down) / (2.0 * remaining_alt))

            # 推力不足时给满油门
            if mass <= 0 or avail_thrust <= 1e-6:
                throttle = 1.0
            else:
                throttle = (a_req_up + g_local) * mass / avail_thrust

            throttle = max(0.0, min(1.0, float(throttle)))

            # 低空兜底：高度<5m 且仍很快 -> 满油门
            if alt_eff < 5.0 and v_down > 10.0:
                throttle = 1.0

            # 若已经触发 LANDED 但竖直速度仍偏大，避免油门过低导致“硬砸/反弹”
            if is_landed and v_down > (vf_down + float(touchdown_vspeed_tol)):
                throttle = max(throttle, 0.8)

            vessel.control.throttle = throttle

            if verbose:
                msg = (
                    f'BURN alt={alt:.1f} radar={radar_alt:.1f} alt_eff={alt_eff:.1f} rem={remaining_alt:.1f} '
                    f'srf_speed={srf_speed:.2f} vs={v_speed:.2f} '
                    f'v_down={v_down:.2f} vf_down={vf_down:.2f} '
                    f'a_req_up={a_req_up:.2f} throttle={throttle:.2f} '
                    f'landed={int(is_landed)} '
                    f'current_thrust={vessel.thrust:.1f}kN'
                )
                status['logs'].append(msg)
                print(msg)

            time.sleep(0.05)

    finally:
        try:
            if ap is not None and hasattr(ap, 'disengage'):
                try:
                    ap.disengage()
                except Exception:
                    pass
        except Exception:
            pass
        try:
            conn.close()
        except Exception:
            pass

    return status

if __name__ == "__main__":
    xxx = auto_land_predictive(
        bottom_clearance_margin=30,
    )
    print(xxx)
    print('done')