import numpy as np
import matplotlib.pyplot as plt

np.random.seed(42)

fig, axes = plt.subplots(2, 2, figsize=(14, 10))

# --- (I) 離散時間・離散状態: マルコフ連鎖 ---
# 天気モデル（晴れ=0, 曇り=1, 雨=2）
P_mc = np.array([
    [0.7, 0.2, 0.1],
    [0.3, 0.4, 0.3],
    [0.2, 0.3, 0.5]
])
states_label = ['晴れ', '曇り', '雨']
n_steps = 50
n_paths = 3

for path in range(n_paths):
    state = 0
    trajectory = [state]
    for _ in range(n_steps - 1):
        state = np.random.choice(3, p=P_mc[state])
        trajectory.append(state)
    axes[0, 0].step(range(n_steps), trajectory, alpha=0.7, linewidth=1.5,
                    label=f'パス {path+1}')

axes[0, 0].set_yticks([0, 1, 2])
axes[0, 0].set_yticklabels(states_label)
axes[0, 0].set_xlabel('Time step $n$')
axes[0, 0].set_title('(I) Markov Chain (Discrete Time, Discrete State)')
axes[0, 0].legend(fontsize=8)
axes[0, 0].grid(True, alpha=0.3)

# --- (II) 離散時間・連続状態: ランダムウォーク ---
n_steps = 200
for path in range(5):
    steps = np.random.choice([-1, 1], size=n_steps)
    walk = np.cumsum(steps)
    axes[0, 1].plot(range(n_steps), walk, alpha=0.6, linewidth=1)

axes[0, 1].axhline(y=0, color='k', linewidth=0.5)
axes[0, 1].set_xlabel('Time step $n$')
axes[0, 1].set_ylabel('$X(n)$')
axes[0, 1].set_title('(II) Random Walk (Discrete Time, Continuous State)')
axes[0, 1].grid(True, alpha=0.3)

# --- (III) 連続時間・離散状態: ポアソン過程 ---
lam = 2.0  # 到着率
T_max = 10.0

for path in range(3):
    # 指数分布の到着間隔を生成
    inter_arrivals = np.random.exponential(1/lam, size=100)
    arrival_times = np.cumsum(inter_arrivals)
    arrival_times = arrival_times[arrival_times <= T_max]

    # 階段関数としてプロット
    t_plot = np.sort(np.concatenate([[0], arrival_times, arrival_times, [T_max]]))
    n_plot = np.sort(np.concatenate([[0, 0], np.arange(1, len(arrival_times)+1),
                                      np.arange(1, len(arrival_times)+1),
                                      [len(arrival_times)]]))
    # より簡潔な方法
    t_events = np.concatenate([[0], arrival_times])
    counts = np.arange(len(t_events))

    axes[1, 0].step(t_events, counts, where='post', alpha=0.7, linewidth=1.5,
                    label=f'パス {path+1}')

axes[1, 0].set_xlabel('Time $t$')
axes[1, 0].set_ylabel('$N(t)$')
axes[1, 0].set_title(f'(III) Poisson Process (Continuous Time, Discrete State, $\\lambda={lam}$)')
axes[1, 0].legend(fontsize=8)
axes[1, 0].grid(True, alpha=0.3)

# --- (IV) 連続時間・連続状態: ブラウン運動 ---
dt = 0.01
T_max = 5.0
t = np.arange(0, T_max, dt)

for path in range(5):
    dW = np.sqrt(dt) * np.random.randn(len(t))
    W = np.cumsum(dW)
    W = np.insert(W, 0, 0)[:-1]  # W(0) = 0
    axes[1, 1].plot(t, W, alpha=0.6, linewidth=1)

axes[1, 1].axhline(y=0, color='k', linewidth=0.5)
axes[1, 1].set_xlabel('Time $t$')
axes[1, 1].set_ylabel('$W(t)$')
axes[1, 1].set_title('(IV) Brownian Motion (Continuous Time, Continuous State)')
axes[1, 1].grid(True, alpha=0.3)

plt.tight_layout()
plt.savefig('stochastic_process_intro.png', dpi=150, bbox_inches='tight')
plt.show()

# 各確率過程の性質をまとめて表示
print("=== 代表的な確率過程の性質 ===")
print(f"{'過程':<20} {'マルコフ性':>10} {'独立増分':>10} {'定常':>10}")
print("-" * 55)
properties = [
    ('マルコフ連鎖', 'Yes', 'No', '条件付き'),
    ('ランダムウォーク', 'Yes', 'Yes', 'No'),
    ('ポアソン過程', 'Yes', 'Yes', 'Yes(増分)'),
    ('ブラウン運動', 'Yes', 'Yes', 'Yes(増分)'),
    ('AR(1)過程', 'Yes', 'No', '条件付き'),
]
for name, markov, indep, stat in properties:
    print(f"{name:<20} {markov:>10} {indep:>10} {stat:>10}")