###################################################################################
# PROGRAM: mh_Gaussian_mixture.py
# DATE: 2026-8-30
# NOTICE: This program accompanies the book "マルコフ連鎖モンテカルロ法入門"
#                     by Koji Hukushima and Yoshihiko Nishikawa
###################################################################################

import numpy as np
import matplotlib.pyplot as plt
from scipy.stats import multivariate_normal

# ==========================================
# 遊べるパラメータ設定
# ==========================================
# ★ここの数字や座標を変更して、サンプル列の動きを観察してみよう

n_samples = 10000  # ステップ数

# サンプリングのスタート地点（初期値）
# 例1: [6.0, -2.0] （分布から遠く離れた場所）
# 例2: [0.0, 0.0]  （左下の山のど真ん中）
# 例3: [8.0, 8.0]  （右上のさらに奥深く）
initial_state = np.array([6.0, -2.0]) 


# ==========================================
# Step 1: ターゲット分布（混合ガウス分布）の定義
# ==========================================
mu1 = np.array([0, 0])
mu2 = np.array([4, 4])
cov1 = np.array([[1.0, 0.9], [0.9, 1.5]])  # 左下の山
cov2 = np.array([[1.0, -0.8], [-0.8, 1.0]]) # 右上の山

def target_density(x):
    """ターゲットとなる2D確率密度関数"""
    return 0.5 * multivariate_normal.pdf(x, mean=mu1, cov=cov1) + \
           0.5 * multivariate_normal.pdf(x, mean=mu2, cov=cov2)

# ==========================================
# Step 2: メトロポリス・ヘイスティングス (MH) 法
# ==========================================
def metropolis_hastings(delta, n_steps, x0):
    samples = []
    x = x0.copy()  # 初期値を設定
    accepted = 0
    
    for _ in range(n_steps):
        # 一様分布から提案（現在地 x の周辺 -delta/2 から delta/2 の範囲）
        x_proposal = x + np.random.uniform(-delta/2, delta/2, size=2)
        
        # 確率密度の比（受容確率）を計算
        p_current = target_density(x)
        p_proposal = target_density(x_proposal)
        acceptance_ratio = min(1.0, p_proposal / p_current)
        
        # 受容判定
        if np.random.rand() < acceptance_ratio:
            x = x_proposal
            accepted += 1
            
        samples.append(x.copy())
        
    acceptance_rate = accepted / n_steps
    return np.array(samples), acceptance_rate

# ==========================================
# Step 3: シミュレーションとグラフ描画
# ==========================================
deltas = [0.2, 1.0, 3.0]

plt.rcParams['mathtext.fontset'] = 'cm'
plt.rcParams["font.size"] = 14

fig, axes = plt.subplots(1, 3, figsize=(15, 5))

# 等高線描画用のグリッド作成
x_grid, y_grid = np.mgrid[-3:9:.05, -3:9:.05]
pos = np.dstack((x_grid, y_grid))
z = target_density(pos)

print(f"=== MH法サンプリング ===")
print(f"ステップ数: {n_samples}")
print(f"初期値: {initial_state}")

for ax, delta in zip(axes, deltas):
    # サンプリング実行
    samples, acc_rate = metropolis_hastings(delta, n_samples, initial_state)
    
    # ターゲット分布の等高線を描画
    ax.contour(x_grid, y_grid, z, levels=15, cmap='viridis', alpha=0.5)
    
    # サンプル数が多すぎる場合、描画点を間引く
    skip = max(1, n_samples // 5000) 
    thinned_samples = samples[::skip]
    
    # サンプルの軌跡をプロット
    ax.plot(thinned_samples[:, 0], thinned_samples[:, 1], 
            'o-', markersize=2, linewidth=0.5, alpha=0.4, color='magenta')
    
    # スタート地点を赤い星マークで強調
    ax.plot(initial_state[0], initial_state[1], 
            marker='*', color='red', markersize=15, markeredgecolor='black', label='Start')
    
    # グラフの装飾
    ax.set_title(f"$\delta$ = {delta}\nAcceptance Rate: {acc_rate:.2f}")
    ax.set_xlim(-3, 9)
    ax.set_ylim(-3, 9)
    ax.set_aspect('equal')
    if delta == deltas[0]:
        ax.legend(loc='upper left')

plt.tight_layout()
plt.show()