import numpy as np
import matplotlib.pyplot as plt

def particle_filter(y_obs, N_p, F, H, Q, R, x0_mean, x0_var):
    """粒子フィルタの実装"""
    T = len(y_obs)
    particles = np.random.normal(x0_mean, np.sqrt(x0_var), N_p)
    weights = np.ones(N_p) / N_p

    x_est = np.zeros(T)
    x_particles_history = np.zeros((T, N_p))

    for t in range(T):
        # 予測ステップ: 状態遷移モデルに従って粒子を伝播
        particles = F * particles + np.random.normal(0, np.sqrt(Q), N_p)

        # 重み更新: 観測尤度で重みを計算
        likelihood = (1.0 / np.sqrt(2 * np.pi * R)) * np.exp(
            -0.5 * (y_obs[t] - H * particles)**2 / R
        )
        weights = likelihood
        weights += 1e-300  # アンダーフロー防止
        weights /= weights.sum()  # 正規化

        # 状態推定（重み付き平均）
        x_est[t] = np.sum(weights * particles)
        x_particles_history[t] = particles

        # リサンプリング（系統リサンプリング）
        cumsum = np.cumsum(weights)
        u = (np.arange(N_p) + np.random.uniform()) / N_p
        indices = np.searchsorted(cumsum, u)
        particles = particles[indices]
        weights = np.ones(N_p) / N_p

    return x_est, x_particles_history

# シミュレーション設定
np.random.seed(42)
T = 80
F = 1.0    # 状態遷移
H = 1.0    # 観測
Q = 0.5    # プロセスノイズ
R = 2.0    # 観測ノイズ
N_p = 500  # 粒子数

# 非線形の真の状態を生成
x_true = np.zeros(T)
y_obs = np.zeros(T)
x_true[0] = 0.0
for t in range(1, T):
    x_true[t] = F * x_true[t-1] + 0.5 * np.sin(0.1 * t) + np.random.normal(0, np.sqrt(Q))
y_obs = H * x_true + np.random.normal(0, np.sqrt(R), T)

# 粒子フィルタの実行
x_est_pf, particles_hist = particle_filter(y_obs, N_p, F, H, Q, R, 0, 1.0)

# 可視化
fig, axes = plt.subplots(2, 1, figsize=(12, 8))

# 上段: 状態推定結果
ax1 = axes[0]
ax1.plot(x_true, 'k-', linewidth=1.5, label='True state')
ax1.scatter(range(T), y_obs, c='gray', s=15, alpha=0.5, label='Observations')
ax1.plot(x_est_pf, 'r-', linewidth=1.5, label='Particle filter estimate')
ax1.set_xlabel('Time step', fontsize=12)
ax1.set_ylabel('State', fontsize=12)
ax1.set_title(f'Particle Filter (N_p={N_p})', fontsize=13)
ax1.legend(fontsize=10)
ax1.grid(True, alpha=0.3)

# 下段: 粒子の分布の推移（一部の時刻で表示）
ax2 = axes[1]
show_times = [5, 20, 40, 60]
colors = ['blue', 'green', 'orange', 'red']
for i, t_show in enumerate(show_times):
    ax2.scatter(np.full(N_p, t_show), particles_hist[t_show],
                alpha=0.05, s=5, color=colors[i])
    ax2.scatter(t_show, x_est_pf[t_show], color=colors[i], s=80,
                marker='D', edgecolors='black', zorder=5, label=f't={t_show}')
ax2.plot(x_true, 'k-', linewidth=1, alpha=0.5)
ax2.set_xlabel('Time step', fontsize=12)
ax2.set_ylabel('State', fontsize=12)
ax2.set_title('Particle Distribution at Selected Time Steps', fontsize=13)
ax2.legend(fontsize=10)
ax2.grid(True, alpha=0.3)

plt.tight_layout()
plt.show()

# 推定精度の評価
rmse = np.sqrt(np.mean((x_true - x_est_pf)**2))
print(f"RMSE (Particle Filter): {rmse:.4f}")