#include "kalman_filter.h"
#include <math.h>

// 初始化卡尔曼滤波器
void kalman_init(KalmanFilter* kf, float dt, float initial_x, float initial_y) {
    // ==================== 状态初始化 ====================
    // 设置初始状态 [x, vx, y, vy]
    kf->x[0] = initial_x; // x位置
    kf->x[1] = 0.0f;      // x速度 (初始为0)
    kf->x[2] = initial_y; // y位置
    kf->x[3] = 0.0f;      // y速度 (初始为0)
    
    // 初始化协方差矩阵 - 表示状态估计的不确定性
    for (int i = 0; i < 4; i++) {
        for (int j = 0; j < 4; j++) {
            kf->P[i][j] = (i == j) ? 1.0f : 0.0f; // 对角线元素为1，其余为0
        }
    }
    
    kf->dt = dt; // 设置采样时间间隔
    
    // ==================== 状态转移矩阵 ====================
    // 恒定速度模型的状态转移矩阵
    // 对于x方向: x_{k+1} = x_k + vx_k * dt
    //           vx_{k+1} = vx_k
    // 对于y方向: y_{k+1} = y_k + vy_k * dt
    //           vy_{k+1} = vy_k
    kf->F[0][0] = 1.0f; kf->F[0][1] = dt;   kf->F[0][2] = 0.0f; kf->F[0][3] = 0.0f;
    kf->F[1][0] = 0.0f; kf->F[1][1] = 1.0f; kf->F[1][2] = 0.0f; kf->F[1][3] = 0.0f;
    kf->F[2][0] = 0.0f; kf->F[2][1] = 0.0f; kf->F[2][2] = 1.0f; kf->F[2][3] = dt;
    kf->F[3][0] = 0.0f; kf->F[3][1] = 0.0f; kf->F[3][2] = 0.0f; kf->F[3][3] = 1.0f;
    
    // ==================== 过程噪声协方差矩阵 Q ====================
    // 过程噪声反映了模型不确定性，通常与加速度扰动相关
    // 修改位置1: 过程噪声参数 sigma_a (加速度扰动标准差)
    // 增大此值 → 更信任测量值（响应更快，但可能震荡）
    // 减小此值 → 更信任预测值（响应平滑，但可能滞后）
    float sigma_a = 0.01f;          // 加速度扰动标准差 (m/s²)
    
    // 计算Q矩阵元素
    float dt2 = dt * dt;
    float dt3 = dt2 * dt;
    float dt4 = dt3 * dt;
    float q11 = dt4 / 4.0f * sigma_a * sigma_a;
    float q12 = dt3 / 2.0f * sigma_a * sigma_a;
    float q22 = dt2 * sigma_a * sigma_a;
    
    // 构建过程噪声协方差矩阵
    kf->Q[0][0] = q11; kf->Q[0][1] = q12; kf->Q[0][2] = 0.0f;  kf->Q[0][3] = 0.0f;
    kf->Q[1][0] = q12; kf->Q[1][1] = q22; kf->Q[1][2] = 0.0f;  kf->Q[1][3] = 0.0f;
    kf->Q[2][0] = 0.0f; kf->Q[2][1] = 0.0f; kf->Q[2][2] = q11; kf->Q[2][3] = q12;
    kf->Q[3][0] = 0.0f; kf->Q[3][1] = 0.0f; kf->Q[3][2] = q12; kf->Q[3][3] = q22;
    
    // ==================== 测量噪声协方差矩阵 R ====================
    // 测量噪声反映了传感器的不确定性
    // 修改位置2: 测量噪声参数 measurement_error
    // 题目中给出测量误差为20cm，即0.2m
    // 增大此值 → 更信任预测模型
    // 减小此值 → 更信任传感器数据
    float measurement_error = 0.15f;  // 测量误差(米)
    
    // 构建测量噪声协方差矩阵
    kf->R[0][0] = measurement_error * measurement_error;
    kf->R[0][1] = 0.0f;
    kf->R[1][0] = 0.0f;
    kf->R[1][1] = measurement_error * measurement_error;
}

// 预测步骤
void kalman_predict(KalmanFilter* kf) {
    // 临时变量
    float temp[4] = {0};
    float temp_P[4][4] = {{0}};
    
    // ==================== 状态预测 ====================
    // 预测状态: x = F * x
    for (int i = 0; i < 4; i++) {
        for (int j = 0; j < 4; j++) {
            temp[i] += kf->F[i][j] * kf->x[j];
        }
    }
    // 更新状态估计
    for (int i = 0; i < 4; i++) kf->x[i] = temp[i];
    
    // ==================== 协方差预测 ====================
    // 预测协方差: P = F * P * F^T + Q
    // 计算 F * P
    for (int i = 0; i < 4; i++) {
        for (int j = 0; j < 4; j++) {
            temp_P[i][j] = 0;
            for (int k = 0; k < 4; k++) {
                temp_P[i][j] += kf->F[i][k] * kf->P[k][j];
            }
        }
    }
    
    // 计算 (F * P) * F^T
    for (int i = 0; i < 4; i++) {
        for (int j = 0; j < 4; j++) {
            kf->P[i][j] = 0;
            for (int k = 0; k < 4; k++) {
                kf->P[i][j] += temp_P[i][k] * kf->F[j][k]; // F^T[j][k] = F[k][j]
            }
            // 加上过程噪声Q
            kf->P[i][j] += kf->Q[i][j];
        }
    }
}

// 更新步骤
void kalman_update(KalmanFilter* kf, float z_x, float z_y) {
    float z[2] = {z_x, z_y}; // 测量值 [x, y]
    
    // ==================== 计算残差 ====================
    // 计算残差: y = z - H * x
    // H = [1 0 0 0; 0 0 1 0] (只测量位置，不测量速度)
    float H_x[2] = {kf->x[0], kf->x[2]}; // H * x
    float y[2] = {z[0] - H_x[0], z[1] - H_x[1]}; // 残差
    
    // ==================== 计算残差协方差 ====================
    // 计算残差协方差: S = H * P * H^T + R
    float S[2][2] = {{0}};
    
    // 由于H矩阵的特殊形式，可以简化计算
    // H * P * H^T 实际上只需要P矩阵的(0,0), (0,2), (2,0), (2,2)元素
    S[0][0] = kf->P[0][0] + kf->R[0][0]; // 加上测量噪声
    S[0][1] = kf->P[0][2];
    S[1][0] = kf->P[2][0];
    S[1][1] = kf->P[2][2] + kf->R[1][1]; // 加上测量噪声
    
    // ==================== 计算卡尔曼增益 ====================
    // 计算卡尔曼增益: K = P * H^T * inv(S)
    // 首先计算S的逆
    float det = S[0][0] * S[1][1] - S[0][1] * S[1][0];
    float S_inv[2][2] = {
        {S[1][1] / det, -S[0][1] / det},
        {-S[1][0] / det, S[0][0] / det}
    };
    
    // 计算K = P * H^T * S_inv
    // H^T = [1 0; 0 0; 0 1; 0 0]
    float K[4][2] = {{0}};
    for (int i = 0; i < 4; i++) {
        // K[i][0] = P[i][0] * S_inv[0][0] + P[i][2] * S_inv[1][0]
        K[i][0] = (kf->P[i][0] * S_inv[0][0] + kf->P[i][2] * S_inv[1][0]);
        // K[i][1] = P[i][0] * S_inv[0][1] + P[i][2] * S_inv[1][1]
        K[i][1] = (kf->P[i][0] * S_inv[0][1] + kf->P[i][2] * S_inv[1][1]);
    }
    
    // ==================== 更新状态估计 ====================
    // 更新状态估计: x = x + K * y
    for (int i = 0; i < 4; i++) {
        kf->x[i] += K[i][0] * y[0] + K[i][1] * y[1];
    }
    
    // ==================== 更新协方差估计 ====================
    // 更新协方差估计: P = (I - K * H) * P
    // 计算K * H
    float KH[4][4] = {{0}};
    for (int i = 0; i < 4; i++) {
        KH[i][0] = K[i][0]; // H[0][0] = 1, H[0][1] = 0, H[0][2] = 0, H[0][3] = 0
        KH[i][2] = K[i][1]; // H[1][0] = 0, H[1][1] = 0, H[1][2] = 1, H[1][3] = 0
    }
    
    // 计算I - K*H
    float I_KH[4][4] = {{0}};
    for (int i = 0; i < 4; i++) {
        for (int j = 0; j < 4; j++) {
            I_KH[i][j] = (i == j ? 1.0f : 0.0f) - KH[i][j];
        }
    }
    
    // 更新协方差: P = (I - K*H) * P
    float temp_P[4][4] = {{0}};
    for (int i = 0; i < 4; i++) {
        for (int j = 0; j < 4; j++) {
            temp_P[i][j] = 0;
            for (int k = 0; k < 4; k++) {
                temp_P[i][j] += I_KH[i][k] * kf->P[k][j];
            }
        }
    }
    
    // 复制回原矩阵
    for (int i = 0; i < 4; i++) {
        for (int j = 0; j < 4; j++) {
            kf->P[i][j] = temp_P[i][j];
        }
    }
}

// 获取当前位置
void kalman_get_position(KalmanFilter* kf, float* x, float* y) {
    *x = kf->x[0]; // x位置
    *y = kf->x[2]; // y位置
}