# rocket_sim.py
import numpy as np
from scipy.integrate import solve_ivp
import matplotlib.pyplot as plt

# 示踪剂：初始状态
# 火箭在 1000m 高度，垂直下降，初始速度 -50 m/s
initial_state = [0, 0, 1000, -50, 0, 0]  # [x, vx, y, vy, theta, omega]

def rocket_dynamics(t, state, u):
    """
    火箭 3-DOF 动力学
    state: [x, vx, y, vy, theta, omega]
    u: [thrust, gimbal_angle]
    """
    x, vx, y, vy, theta, omega = state
    thrust, delta = u

    # 参数
    m = 1000  # 质量 kg
    g = 9.81  # 重力
    I = 1000  # 转动惯量

    # 推力分量（考虑推力矢量控制）
    Fx = thrust * np.sin(delta)
    Fy = thrust * np.cos(delta) - m * g

    # 力矩
    M = thrust * delta * 1.5  # 假设推力中心到质心距离 1.5m

    # 动力学方程
    ax = Fx / m
    ay = Fy / m
    alpha = M / I

    return [vx, ax, vy, ay, omega, alpha]

# 示踪剂2：控制律（PD 控制）
def pd_controller(state, target):
    """
    PD 控制器
    target: [y_target, vy_target, theta_target, omega_target]
    """
    x, vx, y, vy, theta, omega = state
    y_t, vy_t, theta_t, omega_t = target

    # 误差
    e_y = y_t - y
    e_vy = vy_t - vy
    e_theta = theta_t - theta
    e_omega = omega_t - omega

    # PD 增益
    Kp_y = 5000
    Kd_y = 800
    Kp_theta = 2000
    Kd_theta = 500

    # 控制输出
    thrust = m * g + Kp_y * e_y + Kd_y * e_vy
    delta = Kp_theta * e_theta + Kd_theta * e_omega

    return [thrust, delta]

# 仿真
m = 1000
g = 9.81
dt = 0.01
T = 20  # 仿真时间 20s

t_eval = np.arange(0, T, dt)
states = [initial_state]
controls = []

target = [0, 0, 0, 0]  # 目标：着陆，速度为0，姿态垂直

for i in range(len(t_eval) - 1):
    # 计算控制
    u = pd_controller(states[-1], target)
    controls.append(u)

    # 积分一步
    sol = solve_ivp(
        rocket_dynamics,
        [t_eval[i], t_eval[i+1]],
        states[-1],
        args=(u,),
        dense_output=True
    )
    states.append(sol.y[:, -1])

# 绘图
states = np.array(states)
controls = np.array(controls)

plt.figure(figsize=(12, 8))

plt.subplot(3, 1, 1)
plt.plot(t_eval, states[:, 2], label='Altitude')
plt.plot(t_eval, np.zeros_like(t_eval), 'r--', label='Target')
plt.ylabel('Height (m)')
plt.legend()
plt.grid()

plt.subplot(3, 1, 2)
plt.plot(t_eval, states[:, 3], label='Vertical Velocity')
plt.plot(t_eval, np.zeros_like(t_eval), 'r--', label='Target')
plt.ylabel('Vy (m/s)')
plt.legend()
plt.grid()

plt.subplot(3, 1, 3)
plt.plot(t_eval[:-1], controls[:, 0], label='Thrust')
plt.plot(t_eval[:-1], controls[:, 1], label='Gimbal Angle')
plt.xlabel('Time (s)')
plt.ylabel('Control')
plt.legend()
plt.grid()

plt.tight_layout()
plt.savefig('rocket_landing.png')
print("✓ 仿真完成，结果保存到 rocket_landing.png")

python rocket_sim.py

# 添加传感器噪声
def simulate_imu(state, noise_std=0.1):
    """模拟 IMU 测量"""
    x, vx, y, vy, theta, omega = state

    # 真实加速度（需要知道外力）
    # 简化：假设已知推力
    a_true = thrust / m - g

    # 添加噪声
    a_measured = a_true + np.random.normal(0, noise_std)
    omega_measured = omega + np.random.normal(0, noise_std * 0.1)

    return a_measured, omega_measured

from filterpy.kalman import ExtendedKalmanFilter

# 创建 EKF
ekf = ExtendedKalmanFilter(dim_x=6, dim_z=2)

# 初始协方差
ekf.P = np.diag([1.0, 1.0, 10.0, 10.0, 0.1, 0.1])

# 过程噪声（物理不确定性）
ekf.Q = np.diag([0.01, 0.01, 0.1, 0.1, 0.01, 0.01])

# 测量噪声（传感器噪声）
ekf.R = np.diag([0.1, 0.1])

def predict_update(ekf, u, z):
    """预测-更新步骤"""
    ekf.predict(u=u)
    ekf.update(z, HJacobian=H_jacobian, h=measurement_function)
    return ekf.x

# 仿真
for i in range(len(t_eval) - 1):
    # 预测
    ekf.predict(u=controls[i])

    # 模拟测量
    z = simulate_imu(states[i+1])

    # 更新
    ekf.update(z, HJacobian=H_jacobian, h=measurement_function)

    # 记录估计状态
    estimated_states.append(ekf.x.copy())

def generate_landing_trajectory(current_state, landing_point):
    """
    生成着陆轨迹
    使用多项式轨迹
    """
    x0, v0, y0, vy0, theta0, omega0 = current_state

    # 目标：在地面，速度为0
    tf = 10.0  # 预计着陆时间

    # 三次多项式: y(t) = a0 + a1*t + a2*t^2 + a3*t^3
    # 约束:
    # y(0) = y0, y'(0) = vy0
    # y(tf) = 0, y'(tf) = 0

    A = np.array([
        [1, 0, 0, 0],
        [0, 1, 0, 0],
        [1, tf, tf**2, tf**3],
        [0, 1, 2*tf, 3*tf**2]
    ])
    b = np.array([y0, vy0, 0, 0])

    a = np.linalg.solve(A, b)

    # 生成需求轨迹
    t_traj = np.linspace(0, tf, 100)
    y_d = a[0] + a[1]*t_traj + a[2]*t_traj**2 + a[3]*t_traj**3
    vy_d = a[1] + 2*a[2]*t_traj + 3*a[3]*t_traj**2

    return y_d, vy_d

class PIDController:
    def __init__(self, Kp, Ki, Kd):
        self.Kp = Kp
        self.Ki = Ki
        self.Kd = Kd
        self.integral = 0
        self.prev_error = 0

    def update(self, error, dt):
        # P 项
        P = self.Kp * error

        # I 项
        self.integral += error * dt
        I = self.Ki * self.integral

        # D 项
        D = self.Kd * (error - self.prev_error) / dt
        self.prev_error = error

        return P + I + D

# 使用
pid_alt = PIDController(Kp=5000, Ki=10, Kd=800)
pid_vel = PIDController(Kp=2000, Ki=5, Kd=500)

for i in range(len(t_eval) - 1):
    # 高度控制
    error_alt = y_target - states[i][2]
    thrust_command = pid_alt.update(error_alt, dt)

    # 速度控制
    error_vel = vy_target - states[i][3]
    thrust_adjust = pid_vel.update(error_vel, dt)

    thrust = m * g + thrust_command + thrust_adjust

from control import lqr

# 线性化系统矩阵
A = np.array([
    [0, 1, 0, 0],
    [0, 0, 0, 0],
    [0, 0, 0, 1],
    [0, 0, 0, 0]
])

B = np.array([
    [0, 0],
    [1/m, 0],
    [0, 0],
    [0, 1/I]
])

# 权衡矩阵
Q = np.diag([100, 10, 1000, 100])  # 状态权重
R = np.diag([1, 1])  # 控制权重

# 计算 LQR 增益
K, S, E = lqr(A, B, Q, R)

# 控制律: u = -K * (x - x_desired)
def lqr_controller(state, target):
    error = np.array(state) - np.array(target)
    u = -K @ error
    return u

def tvc_allocation(control_force, control_moment):
    """
    将控制力和力矩分配到 TVC 角
    """
    # 假设单发动机 TVC
    max_gimbal = np.radians(15)  # 最大摆角 15 度

    # 分配算法
    # 1. 优先满足力矩需求
    desired_delta = np.clip(control_moment / (thrust * 1.5), -max_gimbal, max_gimbal)

    # 2. 推力 = 垂直分量
    thrust_vertical = control_force / np.cos(desired_delta)

    # 3. 推力限制
    max_thrust = 10000  # 最大推力 10kN
    thrust = np.clip(thrust_vertical, 0, max_thrust)

    return thrust, desired_delta

def rigid_body_dynamics(t, state, u):
    """
    6 自由度刚体动力学
    """
    # 状态: [x, y, z, vx, vy, vz, phi, theta, psi, p, q, r]
    # 输入: [Fx, Fy, Fz, Mx, My, Mz]

    # 提取状态
    r = state[0:3]      # 位置
    v = state[3:6]      # 速度
    euler = state[6:9]  # 欧拉角
    omega = state[9:12] # 角速度

    # 提取输入
    F = u[0:3]
    M = u[3:6]

    # 参数
    m = 1000
    I = np.diag([1000, 1000, 500])  # 惯性张量

    # 姿态矩阵（从体坐标系到惯性坐标系）
    R = rotation_matrix(euler)

    # 重力
    F_gravity = np.array([0, 0, -m * g])

    # 合力（惯性系）
    F_total = R @ F + F_gravity

    # 合力矩（体系）
    M_total = M - np.cross(omega, I @ omega)  # 陀螺力矩

    # 加速度
    a = F_total / m
    alpha = np.linalg.inv(I) @ M_total

    # 角速度与欧拉角关系
    phi, theta, psi = euler
    W = np.array([
        [1, np.sin(phi)*np.tan(theta), np.cos(phi)*np.tan(theta)],
        [0, np.cos(phi), -np.sin(phi)],
        [0, np.sin(phi)/np.cos(theta), np.cos(phi)/np.cos(theta)]
    ])
    euler_dot = W @ omega

    # 导数
    dr_dt = v
    dv_dt = a
    deuler_dt = euler_dot
    domega_dt = alpha

    return np.concatenate([dr_dt, dv_dt, deuler_dt, domega_dt])

import numpy as np
import matplotlib.pyplot as plt

# 测试不同 PID 参数
kp_values = [1000, 5000, 10000]
ki_values = [0, 10, 100]
kd_values = [0, 500, 1000]

plt.figure(figsize=(15, 5))

# 测试 Kp
for i, kp in enumerate(kp_values):
    states = simulate_landing(kp=kp, ki=0, kd=500)
    plt.subplot(1, 3, 1)
    plt.plot(t_eval, states[:, 2], label=f'Kp={kp}')

# 测试 Ki
for i, ki in enumerate(ki_values):
    states = simulate_landing(kp=5000, ki=ki, kd=500)
    plt.subplot(1, 3, 2)
    plt.plot(t_eval, states[:, 2], label=f'Ki={ki}')

# 测试 Kd
for i, kd in enumerate(kd_values):
    states = simulate_landing(kp=5000, ki=10, kd=kd)
    plt.subplot(1, 3, 3)
    plt.plot(t_eval, states[:, 2], label=f'Kd={kd}')

plt.legend()
plt.show()

# 无滤波
states_no_filter = []
x_est = np.zeros(6)
for i in range(len(t_eval) - 1):
    # 直接积分
    a_measured, omega_measured = simulate_imu(states[i])
    x_est[2] += x_est[3] * dt  # y += vy * dt
    x_est[3] += a_measured * dt  # vy += a * dt
    states_no_filter.append(x_est.copy())

# 有卡尔曼滤波
states_with_filter = []
ekf = ExtendedKalmanFilter(...)
# ... 滤波代码 ...

# 对比
plt.figure(figsize=(12, 6))
plt.plot(t_eval, [s[2] for s in states_no_filter], label='No Filter')
plt.plot(t_eval, [s[2] for s in states_with_filter], label='With EKF')
plt.plot(t_eval, [s[2] for s in states_true], 'r--', label='True')
plt.legend()
plt.show()

def falcon9_landing():
    # 初始状态（再入大气层）
    state = [0, 0, 50000, -1000, 0, 0]  # 50km 高度，-1km/s 速度

    # 着陆算法
    for i in range(len(t_eval)):
        altitude = state[2]

        if altitude > 10000:
            # 气动减速
            thrust = 0
        elif altitude > 2000:
            # 单发减速
            thrust = max_thrust / 9
        else:
            # hoverslam（PD 控制）
            e_h = target_altitude - altitude
            e_v = target_velocity - state[3]
            thrust = m * g + Kp * e_h + Kd * e_v

        # 仿真
        state = rocket_dynamics(state, thrust, gimbal)