import numpy as np

# 定义权重函数
def weight_function(zeta, k1=1.5, k2=5):
    abs_zeta = np.abs(zeta)
    psi = np.zeros_like(abs_zeta)
    psi[abs_zeta <= k1] = 1
    mask = (abs_zeta > k1) & (abs_zeta <= k2)
    psi[mask] = (1 - ((abs_zeta[mask] - k1) ** 2 / (k2 - k1) ** 2)) ** 2
    # 确保权重不为零
    psi[psi < 1e-6] = 1e-6
    # print(f"psi={psi}")
    return psi

# RESGQF滤波函数
def resgqf_filter(x_hat, P, z, f, h, Q, R, k1=1.5, k2=5):
    # 状态预测
    x_pred = f(x_hat)
    P_pred = P + Q

    # 计算量测残差
    z_hat = h(x_pred)
    r = z - z_hat
    # 标准化残差
    R_inv_sqrt = np.linalg.inv(np.linalg.cholesky(R))
    zeta = R_inv_sqrt @ r

    # 计算权重矩阵
    psi = weight_function(zeta)
    Psi = np.diag(psi)

    # 确保 Psi 是一个 2D 矩阵
    if Psi.ndim == 1:
        Psi = np.diag(Psi)

    # # 添加正则化项以避免矩阵不可逆
    # epsilon = 1e-6  # 小的正则化项
    # Psi += epsilon * np.eye(Psi.shape[0])

    # 更新量测噪声协方差
    R_bar = R_inv_sqrt @ np.linalg.inv(Psi) @ R_inv_sqrt.T

    # 计算跨协方差 P_xz 和自协方差 P_zz
    P_xz = P_pred[0, 0] * z_hat[0]
    P_zz = z_hat[0]**2 + R_bar[0, 0]


    # 确保 P_zz 不为零以避免除以零错误
    if P_zz < 1e-6:
        P_zz = 1e-6

    # 计算卡尔曼增益
    K = P_xz / P_zz

    # 状态更新
    x_hat = x_pred + K * (z - z_hat[0])
    P = P_pred - K**2 * P_xz

    # 打印调试信息
    print(f"量测: {z[0]}, 预测状态: {x_pred[0]}, 量测残差: {r[0]}, 标准化残差: {zeta[0]}, 权重: {psi[0]}, 量测噪声协方差: {R_bar[0, 0]}, 卡尔曼增益: {K}, 估计状态: {x_hat[0]}, 协方差: {P[0, 0]}")

    return x_hat, P

# 示例参数设置
# 状态转移函数
def f(x):
    return x  # 简单的恒等转移

# 量测函数
def h(x):
    return x  # 简单的恒等映射

# 初始状态估计
x_hat = np.array([0.1])
# 初始协方差
P = np.array([[1.0]])
# 过程噪声协方差
Q = np.array([[0.1]])
# 量测噪声协方差
R = np.array([[1.0]])

# 模拟量测值，包含异常值
measurements = [1.1, 1.2, 10.0, 1.3, 1.4, 1.4, 1.4, 1.4, 1.4, 1.4, 1.4, 1.4, 1.4, 1.4, 1.4, 1.4, 1.4, 1.4, 1.4, 1.4, 1.4, 1.4, 1.4, 1.4, 1.4, 1.4, 10, 1.4, 1.4, 1.4, 1.4, 1.4, 1.4, 1.4]  # 10.0 为异常值

# 运行滤波
for z in measurements:
    x_hat, P = resgqf_filter(x_hat, P, np.array([z]), f, h, Q, R)