import numpy as np
import matplotlib.pyplot as plt

# Define constants
dt = 0.05  # time interval
g = 9.81  # gravitational acceleration

# Initial state [x, vx, y, vy, ay]
initial_state = np.array([0, 3, 0, 0, -g])  # initial position (0,0) with vx0=3, vy0=0

# State transition matrix
F = np.array([[1, dt, 0, 0, 0],
              [0, 1, 0, 0, 0],
              [0, 0, 1, dt, 0.5 * dt**2],
              [0, 0, 0, 1, dt],
              [0, 0, 0, 0, 1]])

# Observation matrix
H = np.array([[1, 0, 0, 0, 0], 
              [0, 0, 1, 0, 0]])

# Process and measurement noise matrices
Q = np.eye(5) * 0.02    # Process noise covariance
Q[4,4] = 0              # 1g - gravity acceleration

R = np.eye(2) * 1.5     # Measurement noise covariance

# Initial estimation covariance matrix
P = np.eye(5) * 50

# Lists to store positions
posterior_state_estimate, ideal_positions, measurements = [], [], []
prior_state_estimate = []

# Time steps for simulation
n_steps = 51

# Initialize state variables
state = initial_state.copy()
ideal_state = initial_state.copy()

# 生成过程噪声，可以使用 Q 的对角线元素生成
process_noise_variance = np.diag(Q)  # 假设是一个对角矩阵

# 用于存储方差的列表
position_X_variances = []
position_X_variances.append(P[0, 0])
position_Y_variances = []
position_Y_variances.append(P[2, 2])

process_state = initial_state.copy()
z_state_measure = initial_state[[0, 2]]

for _ in range(n_steps):
    # 状态空间方程得到的模拟数据
    process_noise = np.random.normal(0, process_noise_variance, size=state.shape)
    process_state = F @ process_state + process_noise

    # 模拟测量得到的状态：measurement state

    z_state_measure = H @ process_state + np.random.normal(0, np.sqrt(R.diagonal()))

    # Prediction step
    process_noise = np.random.normal(0, process_noise_variance, size=state.shape)
    state = F @ state + process_noise
    P = F @ P @ F.T + Q

    prior_state_estimate.append((state[0], state[2]))

    # Assume we get some measurements (with noise)
    ideal_state = F @ ideal_state

    # Kalman Gain
    S = H @ P @ H.T + R
    K = P @ H.T @ np.linalg.inv(S)

    # Update step
    y = z_state_measure - H @ state
    state += K @ y
    P = (np.eye(len(K)) - K @ H) @ P

    # 存储 x，y 位置的方差
    position_X_variances.append(P[0, 0])
    position_Y_variances.append(P[2, 2])

    # Store the results
    posterior_state_estimate.append((state[0], state[2]))
    ideal_positions.append((ideal_state[0], ideal_state[2]))
    measurements.append((z_state_measure[0], z_state_measure[1]))

# Convert lists to numpy arrays
prior_state_estimate = np.array(prior_state_estimate)
posterior_state_estimate = np.array(posterior_state_estimate)
ideal_positions = np.array(ideal_positions)
measurements = np.array(measurements)

# Plot the results
plt.figure(figsize=(10, 6))
plt.plot(prior_state_estimate[:,0],prior_state_estimate[:,1], label = 'Prior State Estimate',linestyle='--')
plt.plot(ideal_positions[:, 0], ideal_positions[:, 1], label='Ideal Trajectory', linestyle='--')

#plt.scatter(measurements[:, 0], measurements[:, 1], c='orange', label='Measurements', marker='o')
plt.plot(measurements[:, 0], measurements[:, 1], c='red', label='Measurements', linestyle='-')
plt.plot(posterior_state_estimate[:, 0], posterior_state_estimate[:, 1], c='blue', label='Posterior Estimated Trajectory', linestyle='-')
plt.xlabel('x position (m)')
plt.ylabel('y position (m)')
plt.title('Projectile Motion with Kalman Filter - Ideal vs Estimated vs Measured')
plt.legend()
plt.grid()
plt.show()

# 绘制X位置的方差随时间的变化
"""
1倍标准差（±1σ）：大约68.27%的数据落在一个标准差范围内，即从分布平均值向左右各一个标准差。
2倍标准差（±2σ）：大约95.45%的数据落在两个标准差范围内，这一个常用的置信区域，许多统计分析中常被用于定义数据的正常范围。
3倍标准差（±3σ）：大约99.73%的数据落在三个标准差范围内，通常用于识别异常值或者极端值。
"""
time_steps = range(n_steps+1)
plt.figure(figsize=(10, 5))
plt.plot(time_steps, position_X_variances, label='Variance of X Position', color='blue')

# Calculate twice the standard deviation
twice_std_dev = 2 * np.sqrt(position_X_variances)
plt.fill_between(time_steps, 
                 0 - np.sqrt(twice_std_dev), 
                 0 + np.sqrt(twice_std_dev), 
                 color='gray', alpha=0.5, step='mid', label='2 Std Dev Range')
plt.xlabel('Time Step')
plt.ylabel('Variance')
plt.title('Convergence of X Position Variance')
plt.legend()
plt.grid(True)
plt.show()

# 绘制Y位置的方差随时间的变化
plt.figure(figsize=(10, 5))
plt.plot(time_steps, position_Y_variances, label='Variance of Y Position', color='blue')
# Calculate twice the standard deviation
twice_std_dev = 2 * np.sqrt(position_Y_variances)
plt.fill_between(time_steps, 
                 0 - np.sqrt(twice_std_dev), 
                 0 + np.sqrt(twice_std_dev), 
                 color='gray', alpha=0.5, step='mid', label='2 Std Dev Range')
plt.xlabel('Time Step')
plt.ylabel('Variance')
plt.title('Convergence of Y Position Variance')
plt.legend()
plt.grid(True)
plt.show()

import numpy as np
import matplotlib.pyplot as plt

# A state space model for a differential drive mobile robot

# A matrix - 3x3 identity matrix
A_t_minus_1 = np.array([[1.0,  0,  0],
                        [  0,1.0,  0],
                        [  0,  0, 1.0]])

# Function to get the B matrix
def getB(yaw, dt):
    """
    Calculates and returns the B matrix
    3x2 matrix -> number of states x number of control inputs
    Expresses how the state of the system [x,y,yaw] changes
    due to the control commands
    :param yaw: The yaw angle in radians
    :param dt: The time change in seconds
    """
    B = np.array([[np.cos(yaw) * dt, 0],
                  [np.sin(yaw) * dt, 0],
                  [0, dt]])
    return B

# Define Jacobian for the transition function
def jacobian_of_motion(state, control, dt):
    _, _, yaw = state
    v, _ = control
    J_f = np.array([
        [1, 0, -v * np.sin(yaw) * dt],
        [0, 1, v * np.cos(yaw) * dt],
        [0, 0, 1]
    ])
    return J_f

def ekf_simulation(initial_state, control_vector, total_time, time_step):
    # Initialize state for EKF
    state_estimate = initial_state
    states_over_time = [state_estimate]

    # Simulated and ideal states tracking
    simulated_states = [initial_state]
    ideal_states = [initial_state]
    prior_states_estimate = [initial_state]
    
    # Initial posterior covariance of new estimate
    Pk = np.diag([1.0, 1.0, 1.0])

    # Measurement noise covariance matrix R_k
    # Measurement Noise init
    std_dev_x_mes = 0.5    # Standard deviation for x measurement
    std_dev_y_mes = 0.5    # Standard deviation for y measurement
    std_dev_yaw_mes = 0.15  # Standard deviation for yaw measurement
    
    R_k = np.diag([std_dev_x_mes**2, std_dev_y_mes**2, std_dev_yaw_mes**2])
    
    H = np.array([[1, 0, 0],
                  [0, 1, 0],
                  [0, 0, 1]])

    # Jacobian matrix J_h, which is the same as H for a linear model
    J_h = H
    
    # Process Noise init
    std_dev_x = 0.15
    std_dev_y = 0.15
    std_dev_yaw = 0.15

    # Simulation loop
    for _ in np.arange(0, total_time, time_step):
        # Ideal state (no noise)
        ideal_state = A_t_minus_1 @ ideal_states[-1] + getB(ideal_states[-1][2], time_step) @ control_vector
        ideal_states.append(ideal_state)

        # 1-1. State Prediction step
        # 生成过程噪声
        process_noise_sim = np.random.normal(0, [std_dev_x, std_dev_y, std_dev_yaw], size=ideal_state.shape)

        Prior_state_Estimate = (A_t_minus_1 @ states_over_time[-1] + getB(states_over_time[-1][2], time_step) @ control_vector
                                + process_noise_sim)
        prior_states_estimate.append(Prior_state_Estimate)
        # 1-2. Process Error covariance
        Jf = jacobian_of_motion(states_over_time[-1], control_vector, dt=time_step)
        
        # Generate process noise for simulation as additive noise
        process_noise = np.random.normal(0, [std_dev_x, std_dev_y, std_dev_yaw])
        # Convert to a diagonal covariance matrix if needed for use in the EKF
        process_noise_cov = np.diag(process_noise**2)   
        
        Pk = Jf @ Pk @ Jf.T +  process_noise_cov #np.random.uniform(-process_noise, process_noise, 3)     
        
        # 2-1. Compute the Kalman gain
        K_k = Pk @  J_h.T @ np.linalg.inv(J_h @ Pk @ J_h.T + R_k)

        # Simulate sensor reading with measurement noise (using ideal_state)
        #tolerance_noise = np.random.uniform(-measurement_noise, measurement_noise, 3)
        measurement_noise = np.random.normal(0, [std_dev_x_mes, std_dev_y_mes, std_dev_yaw_mes])
        simulated_state = ideal_state + measurement_noise
        
        simulated_states.append(simulated_state)
        
        # 2-2. Update estimate with measurement z_k
        state_estimate = Prior_state_Estimate + K_k @ (simulated_state - J_h @ Prior_state_Estimate)
        
        # 2-3. Update the error covariance
        Pk = (np.eye(3) - K_k @ J_h) @ Pk

        # Record the updated state
        states_over_time.append(state_estimate)

    return np.array(prior_states_estimate), np.array(states_over_time), np.array(simulated_states), np.array(ideal_states)

def plot_states(prior_states_estimate, states_over_time, simulated_states, ideal_states):
    plt.figure(figsize=(12, 8))

    # Plot position on XY plane
    plt.subplot(2, 1, 1)
    plt.plot(ideal_states[:, 0], ideal_states[:, 1], 'g-', label='Ideal Path')
    plt.plot(prior_states_estimate[:,0],prior_states_estimate[:,1], label='prior Path estimate')
    plt.plot(simulated_states[:, 0], simulated_states[:, 1], 'b--', label='Simulated Measure Path')
    plt.plot(states_over_time[:, 0], states_over_time[:, 1], 'r-', label='EKF posterior state Path')

    plt.xlabel('X Position (meters)')
    plt.ylabel('Y Position (meters)')
    plt.title('Trajectory Over Time')
    plt.legend()
    plt.grid(True)

    # Plot yaw angle over time
    plt.subplot(2, 1, 2)
    time_points = np.arange(0, len(states_over_time))
    plt.plot(time_points, ideal_states[:, 2], 'g-', label='Ideal Yaw')
    plt.plot(prior_states_estimate[:,2], label='prior Yaw estimate')
    plt.plot(time_points, simulated_states[:, 2], 'b--', label='Simulated Measure Yaw')
    plt.plot(time_points, states_over_time[:, 2], 'r-', label='EKF posterior state Yaw')

    plt.xlabel('Time Step')
    plt.ylabel('Yaw (radians)')
    plt.title('Yaw Angle Over Time')
    plt.legend()
    plt.grid(True)

    plt.tight_layout()
    plt.show()

def main():
    # Initial conditions
    initial_state = np.array([0.0, 0.0, 0.0]) # [x_init, y_init, yaw_init]
    control_vector = np.array([4.5, 0.15])    # [v, yaw_rate]

    # Driving parameters
    total_time = 10.0 # Total simulation time in seconds
    time_step = 0.05   # Time step for each iteration in seconds

    # Run simulation
    prior_states_estimate, ekf_states, simulated_states, ideal_states = ekf_simulation(
        initial_state, control_vector, total_time=total_time, time_step=time_step)

    # Plot the simulation results
    plot_states(prior_states_estimate, ekf_states, simulated_states, ideal_states)

main()