/**
 * @file   corridor_ekf.c
 * @brief  鲁棒 EKF 走廊相对定位滤波器实现
 *
 *         完整算法流程：
 *
 *         【预测步】Predict(vx, wz, dt)
 *           x_pred = f(x, u)  -- 非线性状态转移
 *           P_pred = F * P * F^T + Q  -- 协方差预测
 *
 *         【更新步】Update(obs)
 *           z     = h(x_pred) -- 观测预测
 *           y     = z_meas - z   -- 新息 (Innovation)
 *           S     = H * P_pred * H^T + R -- 新息协方差
 *           d²    = y^T * S^(-1) * y  -- 马氏距离平方
 *           if d² > χ²_threshold: 拒绝观测 (鲁棒)
 *           K     = P_pred * H^T * S^(-1) -- 卡尔曼增益
 *           x     = x_pred + K * y     -- 状态更新
 *           P     = (I - K * H) * P_pred -- 协方差更新
 */

#include "corridor_ekf.h"
#include <math.h>
#include <string.h>

/* =========================================================
 * 内部静态状态
 * ========================================================= */
static CorridorEKFConfig_t  s_cfg;
static CorridorEKFState_t  s_state;
static bool                s_initialized = false;
static uint32_t            s_last_update_ms = 0U;

/* 协方差上界保护阈值 */
#define P_MAX_DIAG    100.0f

/* =========================================================
 * 内部辅助函数
 * ========================================================= */

/** 限幅 */
static inline float clampf(float val, float lo, float hi)
{
    if (val < lo) return lo;
    if (val > hi) return hi;
    return val;
}

/** 对称矩阵拷贝 + 双向取平均 (减少舍入误差传播) */
static void symmetrize(float M[3][3])
{
    float avg;
    avg = (M[0][1] + M[1][0]) * 0.5f;
    M[0][1] = M[1][0] = avg;
    
    avg = (M[0][2] + M[2][0]) * 0.5f;
    M[0][2] = M[2][0] = avg;
    
    avg = (M[1][2] + M[2][1]) * 0.5f;
    M[1][2] = M[2][1] = avg;
}

/** 角度归一化到 [-π, π]
 * 防止 IMU yaw 累积角度跨越 ±π 时导致新息突变
 */
static float wrap_angle(float angle)
{
    const float PI = 3.14159265358979323846f;
    while (angle > PI)  angle -= 2.0f * PI;
    while (angle < -PI) angle += 2.0f * PI;
    return angle;
}

/** P 上界保护 */
static void protect_P(float P[3][3])
{
    for (int i = 0; i < 3; i++) {
        if (P[i][i] > P_MAX_DIAG) P[i][i] = P_MAX_DIAG;
        if (P[i][i] < 0.0f)      P[i][i] = 0.0f;
    }
}

/** Joseph 形式协方差更新 (1DOF 标量观测)
 * P_new = (I - K*H) * P * (I - K*H)^T + K * R * K^T
 * 
 * 参数:
 *   P[3][3]  - 先验协方差 (输入/输出)
 *   K[3]     - 卡尔曼增益向量
 *   H[3]     - 观测矩阵 (1x3 行向量)
 *   R        - 观测噪声方差 (标量)
 * 
 * 优点: 保证数值稳定性，即使有舍入误差也能保持 P 正定
 */
static void joseph_update_P(float P[3][3], const float K[3], const float H[3], float R)
{
    /* 计算 A = (I - K*H) */
    float A[3][3];
    for (int i = 0; i < 3; i++) {
        for (int j = 0; j < 3; j++) {
            A[i][j] = (i == j ? 1.0f : 0.0f) - K[i] * H[j];
        }
    }
    
    /* 计算 A * P */
    float AP[3][3];
    for (int i = 0; i < 3; i++) {
        for (int j = 0; j < 3; j++) {
            AP[i][j] = 0.0f;
            for (int k = 0; k < 3; k++) {
                AP[i][j] += A[i][k] * P[k][j];
            }
        }
    }
    
    /* 计算 A * P * A^T */
    float APAT[3][3];
    for (int i = 0; i < 3; i++) {
        for (int j = 0; j < 3; j++) {
            APAT[i][j] = 0.0f;
            for (int k = 0; k < 3; k++) {
                APAT[i][j] += AP[i][k] * A[j][k];  // A^T[k][j] = A[j][k]
            }
        }
    }
    
    /* 计算 K * R * K^T 并加到 APAT */
    for (int i = 0; i < 3; i++) {
        for (int j = 0; j < 3; j++) {
            P[i][j] = APAT[i][j] + K[i] * R * K[j];
        }
    }
}

/** 计算 2x2 对称矩阵的逆 (原地) */
static bool invert_2x2_sym(float S[2][2])
{
    float det = S[0][0] * S[1][1] - S[0][1] * S[1][0];
    if (fabsf(det) < 1e-8f) {
        return false;  // 奇异矩阵
    }
    float inv_det = 1.0f / det;
    float S00 = S[0][0];
    S[0][0] =  inv_det * S[1][1];
    S[1][1] =  inv_det * S00;
    S[0][1] = -inv_det * S[0][1];
    S[1][0] =  S[0][1];
    return true;
}

/** 2x2 对称矩阵求逆 (原地) */
static bool invert_3x3_cholesky(float S[3][3])
{
    // 使用 Cholesky 分解求逆
    float L[3][3] = {0};
    
    for (int i = 0; i < 3; i++) {
        for (int j = 0; j <= i; j++) {
            float sum = S[i][j];
            for (int k = 0; k < j; k++) {
                sum -= L[i][k] * L[j][k];
            }
            if (i == j) {
                if (sum <= 0.0f) return false;
                L[i][j] = sqrtf(sum);
            } else {
                L[i][j] = sum / L[j][j];
            }
        }
    }

    // 求逆: S_inv = L^(-T) * L^(-1)
    float Linv[3][3] = {0};
    for (int i = 0; i < 3; i++) {
        Linv[i][i] = 1.0f / L[i][i];
        for (int j = i - 1; j >= 0; j--) {
            float sum = 0.0f;
            for (int k = j + 1; k <= i; k++) {
                sum += L[k][j] * Linv[k][i];
            }
            Linv[j][i] = -sum / L[j][j];
        }
    }

    // S_inv = Linv^T * Linv
    for (int i = 0; i < 3; i++) {
        for (int j = 0; j < 3; j++) {
            float sum = 0.0f;
            for (int k = 0; k < 3; k++) {
                sum += Linv[k][i] * Linv[k][j];
            }
            S[i][j] = sum;
        }
    }
    return true;
}

/** 计算马氏距离平方 (新息向量 y, 新息协方差 S_inv) */
static float mahalanobis_d2_3dof(const float y[3], const float S_inv[3][3])
{
    // d² = y^T * S_inv * y
    float tmp[3] = {0};
    for (int i = 0; i < 3; i++) {
        for (int j = 0; j < 3; j++) {
            tmp[i] += S_inv[i][j] * y[j];
        }
    }
    float d2 = 0.0f;
    for (int i = 0; i < 3; i++) {
        d2 += y[i] * tmp[i];
    }
    return d2;
}

/** 计算马氏距离平方 (1DOF: 只用 e_y) */
static float mahalanobis_d2_1dof(float y_ey, float S_inv_00)
{
    return y_ey * y_ey * S_inv_00;
}

/** 计算马氏距离平方 (2DOF: e_ey + e_th_avg) */
static float mahalanobis_d2_2dof(const float y[2], const float S_inv[2][2])
{
    float tmp[2] = {0};
    for (int i = 0; i < 2; i++) {
        for (int j = 0; j < 2; j++) {
            tmp[i] += S_inv[i][j] * y[j];
        }
    }
    float d2 = 0.0f;
    for (int i = 0; i < 2; i++) {
        d2 += y[i] * tmp[i];
    }
    return d2;
}

/** 零矩阵 */
static void zero_3x3(float M[3][3])
{
    memset(M, 0, sizeof(float) * 9);
}

/** 单位矩阵 */
static void eye_3x3(float M[3][3])
{
    zero_3x3(M);
    M[0][0] = M[1][1] = M[2][2] = 1.0f;
}

/* =========================================================
 * API 实现
 * ========================================================= */

void CorridorEKF_Init(const CorridorEKFConfig_t *config)
{
    s_cfg = *config;
    memset(&s_state, 0, sizeof(s_state));

    /* 初始化状态 */
    s_state.x[0] = 0.0f;  // e_y
    s_state.x[1] = 0.0f;  // e_th
    s_state.x[2] = 0.0f;  // s

    /* 初始化协方差 */
    eye_3x3(s_state.P);
    s_state.P[0][0] = config->P0_diag[0];
    s_state.P[1][1] = config->P0_diag[1];
    s_state.P[2][2] = config->P0_diag[2];

    s_initialized = true;
}

void CorridorEKF_Reset(void)
{
    if (!s_initialized) return;
    
    s_state.x[0] = 0.0f;
    s_state.x[1] = 0.0f;
    s_state.x[2] = 0.0f;
    
    eye_3x3(s_state.P);
    s_state.P[0][0] = s_cfg.P0_diag[0];
    s_state.P[1][1] = s_cfg.P0_diag[1];
    s_state.P[2][2] = s_cfg.P0_diag[2];
    
    s_last_update_ms = 0U;
}

void CorridorEKF_ResetHeading(void)
{
    if (!s_initialized) return;

    s_state.x[1] = 0.0f;

    /* 清理航向与其它状态的耦合，避免旧航向误差继续通过协方差传播。 */
    s_state.P[0][1] = 0.0f;
    s_state.P[1][0] = 0.0f;
    s_state.P[1][2] = 0.0f;
    s_state.P[2][1] = 0.0f;
    s_state.P[1][1] = s_cfg.P0_diag[1];
}

void CorridorEKF_RebaseAfterTurnaround(void)
{
    if (!s_initialized) return;

    /* 同一条走廊掉头后，新的前进方向相反，横向误差符号需要镜像。 */
    s_state.x[0] = -s_state.x[0];
    s_state.x[1] = 0.0f;

    /* e_y 与 e_th 的相关项在掉头后不再可直接沿用，清零重新收敛。 */
    s_state.P[0][1] = 0.0f;
    s_state.P[1][0] = 0.0f;
    s_state.P[1][2] = 0.0f;
    s_state.P[2][1] = 0.0f;
    s_state.P[1][1] = s_cfg.P0_diag[1];
}

void CorridorEKF_SetProcessNoise(float q_ey, float q_eth, float q_s)
{
    s_cfg.q_ey  = q_ey;
    s_cfg.q_eth = q_eth;
    s_cfg.q_s  = q_s;
}

void CorridorEKF_SetMeasurementNoise(float r_ey, float r_eth)
{
    s_cfg.r_ey  = r_ey;
    s_cfg.r_eth = r_eth;
}

/* =========================================================
 * 预测步 (Predict)
 * ========================================================= */
void CorridorEKF_Predict(float odom_vx, float imu_wz, float dt)
{
    if (!s_initialized || dt <= 0.0f) return;

    float e_y  = s_state.x[0];
    float e_th = s_state.x[1];
    float s    = s_state.x[2];
    float vx   = odom_vx;
    float wz   = imu_wz;

    /* 状态预测: x_pred = f(x, u) */
    float cos_th = cosf(e_th);
    float sin_th = sinf(e_th);

    /* 安全检查: cos_th 不能太小 (防止数值爆炸) */
    if (fabsf(cos_th) < 0.01f) cos_th = (cos_th >= 0.0f) ? 0.01f : -0.01f;

    float e_y_pred  = e_y  + vx * sin_th * dt;
    float e_th_pred = e_th + wz * dt;
    float s_pred    = s    + vx * cos_th * dt;

    s_state.x[0] = e_y_pred;
    s_state.x[1] = e_th_pred;
    s_state.x[2] = s_pred;

    /* 雅可比矩阵 F (状态转移的 Jacobian) */
    float F[3][3] = {0};
    F[0][0] = 1.0f;
    F[0][1] = vx * cos_th * dt;  // de_y/de_th
    F[0][2] = 0.0f;

    F[1][0] = 0.0f;
    F[1][1] = 1.0f;
    F[1][2] = 0.0f;

    F[2][0] = 0.0f;
    F[2][1] = -vx * sin_th * dt;  // ds/de_th
    F[2][2] = 1.0f;

    /* 过程噪声协方差 Q (含耦合项)
     * [改进] 添加 e_y 和 e_th 的耦合噪声，反映横向-航向动力学耦合 */
    float Q[3][3] = {0};
    Q[0][0] = s_cfg.q_ey  * dt * dt;
    Q[1][1] = s_cfg.q_eth * dt * dt;
    Q[2][2] = s_cfg.q_s   * dt * dt;
    Q[0][1] = Q[1][0] = s_cfg.q_ey_eth * dt * dt;  // 横向-航向耦合

    /* 协方差预测: P_pred = F * P * F^T + Q */
    float F_P[3][3] = {0};
    for (int i = 0; i < 3; i++) {
        for (int j = 0; j < 3; j++) {
            for (int k = 0; k < 3; k++) {
                F_P[i][j] += F[i][k] * s_state.P[k][j];
            }
        }
    }

    float P_pred[3][3] = {0};
    for (int i = 0; i < 3; i++) {
        for (int j = 0; j < 3; j++) {
            for (int k = 0; k < 3; k++) {
                P_pred[i][j] += F_P[i][k] * F[j][k];  // F^T: F[j][k] = F[k][j]
            }
            P_pred[i][j] += Q[i][j];  // 加过程噪声 (含非对角项)
        }
    }

    symmetrize(P_pred);

    for (int i = 0; i < 3; i++) {
        for (int j = 0; j < 3; j++) {
            s_state.P[i][j] = P_pred[i][j];
        }
    }

    protect_P(s_state.P);
}

/* =========================================================
 * 观测步 (Update) - 鲁棒 EKF
 *
 * 设计决策 (方向 B — IMU 主导航向)：
 *   侧墙激光仅用于更新横向位置 e_y，不再构建航向观测 z_eth_L/z_eth_R。
 *   航向 e_th 完全由 IMU 主导：
 *     - 预测步: imu_wz 驱动 e_th 积分
 *     - 观测步: CorridorEKF_UpdateIMUYaw() 提供 yaw_continuous 标量约束
 *   侧墙前后差分 (d_lr-d_lf) 的噪声在 ±2cm 误差下过大，不适合做航向主观测。
 * ========================================================= */
int CorridorEKF_Update(const CorridorObs_t *obs, CorridorState_t *out_state)
{
    if (!s_initialized) return 0;

    /* 维护最近一次 EKF 输出对应的观测时间戳，供 GetState() 返回一致结果。 */
    s_last_update_ms = obs->t_ms;

    int updated_obs_count = 0;

    /* 清除新息和拒绝掩码 */
    memset(s_state.K, 0, sizeof(s_state.K));
    uint8_t reject_mask = 0U;
    float max_maha_d2 = 0.0f;

    /* ----------------------------------------------------
     * 提取有效观测
     * ---------------------------------------------------- */
    bool left_ok  = ((obs->valid_mask & (1U << 0)) != 0U) &&
                    ((obs->valid_mask & (1U << 1)) != 0U);
    bool right_ok = ((obs->valid_mask & (1U << 2)) != 0U) &&
                    ((obs->valid_mask & (1U << 3)) != 0U);

    /* 左右侧横向平均距离 */
    float d_lf = obs->d_lf, d_lr = obs->d_lr;
    float d_rf = obs->d_rf, d_rr = obs->d_rr;
    float W    = s_cfg.corridor_width;
    float yoff = s_cfg.y_offset;
    float Rw   = s_cfg.robot_width;

    /* [改进A] 分侧传感器内缩补偿 — 消除左右安装不对称引起的系统性偏置
     * 左右各自使用独立的 inset 值计算期望居中读数 d_center */
    float d_center_left  = (W - Rw) / 2.0f + s_cfg.left_sensor_inset;
    float d_center_right = (W - Rw) / 2.0f + s_cfg.right_sensor_inset;

    /* 观测值 (测量) — 仅横向位置，不含航向 */
    float z_ey   = 0.0f;
    int valid_sides = 0;

    if (left_ok) {
        z_ey += d_center_left - ((d_lf + d_lr) / 2.0f) - yoff;
        valid_sides++;
    }

    if (right_ok) {
        z_ey += ((d_rf + d_rr) / 2.0f) - d_center_right - yoff;
        valid_sides++;
    }

    if (valid_sides == 0) {
        out_state->t_ms          = obs->t_ms;
        out_state->e_y           = s_state.x[0];
        out_state->e_th          = s_state.x[1];
        out_state->s             = s_state.x[2];
        out_state->conf          = clampf(1.0f - (s_state.P[0][0] + s_state.P[1][1]) * 0.1f, 0.0f, 1.0f);
        out_state->obs_reject_mask = 0xFF;
        out_state->mahalanobis_d2 = 0.0f;

        for (int i = 0; i < 3; i++) {
            for (int j = 0; j < 3; j++) {
                out_state->P[i][j] = s_state.P[i][j];
            }
        }

        /* 协方差膨胀 (无观测时的信任衰减) — 仅膨胀 e_y */
        s_state.P[0][0] += s_cfg.q_ey * 5.0f;
        protect_P(s_state.P);
        return 0;
    }

    if (valid_sides == 2) {
        z_ey /= 2.0f;
    }

    /* ----------------------------------------------------
     * 1DOF 标量 EKF 更新 — 仅 e_y
     * ---------------------------------------------------- */
    float e_y = s_state.x[0];
    float y_ey = z_ey - e_y;

    /* [改进F] 自适应观测噪声 R:
     *   双侧观测: R × 0.5 (噪声互相平均)
     *   单侧观测: R × 3.0 (缺少交叉验证, VL53 可信度低时尤需降低信任)
     */
    float R_ey = s_cfg.r_ey;
    if (valid_sides == 2) {
        R_ey *= 0.5f;
    } else if (valid_sides == 1) {
        R_ey *= 3.0f;
    }
    float S_ey = s_state.P[0][0] + R_ey;

    if (fabsf(S_ey) < 1e-8f) {
        goto output_result;
    }

    float d2_ey = y_ey * y_ey / S_ey;
    max_maha_d2 = d2_ey;

    if (d2_ey > s_cfg.chi2_1dof) {
        reject_mask |= (1U << 0);
        goto output_result;
    }

    float S_inv_ey = 1.0f / S_ey;
    float K_ey[3];
    K_ey[0] = s_state.P[0][0] * S_inv_ey;
    K_ey[1] = s_state.P[1][0] * S_inv_ey;
    K_ey[2] = s_state.P[2][0] * S_inv_ey;

    s_state.x[0] += K_ey[0] * y_ey;
    s_state.x[1] += K_ey[1] * y_ey;
    s_state.x[2] += K_ey[2] * y_ey;

    /* Joseph 形式协方差更新: P = (I-KH)*P*(I-KH)^T + K*R*K^T
     * H = [1, 0, 0] (仅观测 e_y) */
    float H_ey[3] = {1.0f, 0.0f, 0.0f};
    joseph_update_P(s_state.P, K_ey, H_ey, R_ey);

    symmetrize(s_state.P);
    protect_P(s_state.P);

    updated_obs_count = 1;

output_result:
    out_state->t_ms = obs->t_ms;
    out_state->e_y  = s_state.x[0];
    out_state->e_th = s_state.x[1];
    out_state->s    = s_state.x[2];

    for (int i = 0; i < 3; i++) {
        for (int j = 0; j < 3; j++) {
            out_state->P[i][j] = s_state.P[i][j];
        }
    }

    out_state->mahalanobis_d2 = max_maha_d2;
    out_state->obs_reject_mask = reject_mask;

    float P_trace = s_state.P[0][0] + s_state.P[1][1] + s_state.P[2][2];
    float conf_from_P = clampf(1.0f - P_trace * 0.5f, 0.0f, 1.0f);
    float side_factor = (valid_sides == 2) ? 1.0f : 0.7f;
    float reject_penalty = (reject_mask & (1U << 0)) ? 0.5f : 1.0f;

    out_state->conf = clampf(conf_from_P * side_factor * reject_penalty, 0.0f, 1.0f);

    return updated_obs_count;
}

/* =========================================================
 * 辅助 API
 * ========================================================= */

/* =========================================================
 * IMU 航向观测更新 (独立 1DOF 标量 EKF 更新)
 *
 * 观测方程: z_eth_imu = imu_yaw_rad - imu_yaw_ref_rad
 * 对应状态: e_th (x[1])
 * H = [0, 1, 0]  (只观测 e_th)
 *
 * 设计意图：
 *   - IMU 内部维护的 yaw 经过模块校准，比外部积分 wz 更稳定
 *   - 在侧墙观测丢失 (转弯/单侧退化) 时提供航向约束
 *   - 使用较大 R 值，让侧墙观测在有效时主导
 * ========================================================= */
void CorridorEKF_UpdateIMUYaw(float imu_yaw_rad, float imu_yaw_ref_rad, bool valid)
{
    if (!s_initialized || !valid) return;

    /* 观测值: IMU 相对于走廊参考方向的航向偏差
     * [改进] 角度归一化防止 ±π 跨越时新息突变 */
    float z_eth_imu = wrap_angle(imu_yaw_rad - imu_yaw_ref_rad);

    /* 新息: y = z - h(x), h(x) = e_th = x[1]
     * [改进] 再次归一化，防止 e_th 和 z_eth_imu 符号不同时差值超出 [-π, π] */
    float y_imu = wrap_angle(z_eth_imu - s_state.x[1]);

    /* H = [0, 1, 0] → S = P[1][1] + R_imu */
    float R_imu = s_cfg.r_eth_imu;
    float S_imu = s_state.P[1][1] + R_imu;

    if (fabsf(S_imu) < 1e-8f) return;

    /* χ² 1DOF 检验: d² = y² / S */
    float d2_imu = y_imu * y_imu / S_imu;
    if (d2_imu > s_cfg.chi2_1dof) return;  /* 拒绝异常观测 */

    /* 卡尔曼增益: K = P * H^T / S = P[:][1] / S */
    float K[3];
    float S_inv = 1.0f / S_imu;
    K[0] = s_state.P[0][1] * S_inv;
    K[1] = s_state.P[1][1] * S_inv;
    K[2] = s_state.P[2][1] * S_inv;

    /* 状态更新: x += K * y */
    s_state.x[0] += K[0] * y_imu;
    s_state.x[1] += K[1] * y_imu;
    s_state.x[2] += K[2] * y_imu;

    /* Joseph 形式协方差更新: P = (I-KH)*P*(I-KH)^T + K*R*K^T
     * H = [0, 1, 0] (仅观测 e_th) */
    float H_imu[3] = {0.0f, 1.0f, 0.0f};
    joseph_update_P(s_state.P, K, H_imu, R_imu);

    symmetrize(s_state.P);
    protect_P(s_state.P);
}

void CorridorEKF_GetState(CorridorState_t *out)
{
    if (!s_initialized || out == NULL) return;

    out->t_ms = s_last_update_ms;
    out->e_y  = s_state.x[0];
    out->e_th = s_state.x[1];
    out->s    = s_state.x[2];

    for (int i = 0; i < 3; i++) {
        for (int j = 0; j < 3; j++) {
            out->P[i][j] = s_state.P[i][j];
        }
    }

    out->mahalanobis_d2 = 0.0f;
    out->obs_reject_mask = 0U;

    float P_trace = s_state.P[0][0] + s_state.P[1][1] + s_state.P[2][2];
    out->conf = clampf(1.0f - P_trace * 0.5f, 0.0f, 1.0f);
}