import numpy as np
import matplotlib.pyplot as plt
from matplotlib.patches import Circle
from matplotlib_venn import venn2

# 同時確率表の可視化
joint_prob = np.array([[0.30, 0.10],
                       [0.10, 0.50]])

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

# 左: 同時確率表のヒートマップ
ax1 = axes[0]
im = ax1.imshow(joint_prob, cmap='Blues', vmin=0, vmax=0.5)
ax1.set_xticks([0, 1])
ax1.set_xticklabels(['Rain (B)', 'Sunny (not B)'], fontsize=11)
ax1.set_yticks([0, 1])
ax1.set_yticklabels(['Umbrella (A)', 'No umbrella (not A)'], fontsize=11)
for i in range(2):
    for j in range(2):
        ax1.text(j, i, f'{joint_prob[i, j]:.2f}', ha='center', va='center',
                 fontsize=14, color='black' if joint_prob[i, j] < 0.3 else 'white')
ax1.set_title('Joint Probability P(A, B)', fontsize=13)
plt.colorbar(im, ax=ax1, shrink=0.8)

# 中央: 条件付き確率の比較
ax2 = axes[1]
P_A = joint_prob.sum(axis=1)  # P(A), P(not A)
P_B = joint_prob.sum(axis=0)  # P(B), P(not B)

# P(A|B) vs P(A|not B)
P_A_given_B = joint_prob[0, 0] / P_B[0]
P_A_given_notB = joint_prob[0, 1] / P_B[1]

categories = ['P(A,B)\n(Joint)', 'P(A|B)\n(Conditional)', 'P(A)\n(Marginal)']
values = [joint_prob[0, 0], P_A_given_B, P_A[0]]
colors = ['steelblue', 'coral', 'green']
bars = ax2.bar(categories, values, color=colors, alpha=0.7)
for bar, val in zip(bars, values):
    ax2.text(bar.get_x() + bar.get_width()/2, bar.get_height() + 0.02,
             f'{val:.2f}', ha='center', fontsize=13)
ax2.set_ylabel('Probability', fontsize=12)
ax2.set_title('Comparison: Joint vs Conditional vs Marginal', fontsize=13)
ax2.set_ylim(0, 1)
ax2.grid(True, alpha=0.3, axis='y')

# 右: モンテカルロシミュレーション
ax3 = axes[2]
np.random.seed(42)
n_sim = 50000
samples = np.random.choice(4, size=n_sim, p=joint_prob.flatten())
# 0: A&B, 1: A&notB, 2: notA&B, 3: notA&notB
is_A = (samples == 0) | (samples == 1)
is_B = (samples == 0) | (samples == 2)

# 条件付き確率の推定
n_values = np.arange(100, n_sim, 100)
P_AB_est = []
P_A_given_B_est = []
P_A_est = []

for n in n_values:
    a = is_A[:n]
    b = is_B[:n]
    P_AB_est.append(np.mean(a & b))
    P_A_given_B_est.append(np.mean(a[b]) if b.sum() > 0 else 0)
    P_A_est.append(np.mean(a))

ax3.plot(n_values, P_AB_est, 'b-', alpha=0.6, linewidth=1, label='P(A,B) sim')
ax3.plot(n_values, P_A_given_B_est, 'r-', alpha=0.6, linewidth=1, label='P(A|B) sim')
ax3.plot(n_values, P_A_est, 'g-', alpha=0.6, linewidth=1, label='P(A) sim')
ax3.axhline(y=0.30, color='b', linestyle='--', alpha=0.5)
ax3.axhline(y=0.75, color='r', linestyle='--', alpha=0.5)
ax3.axhline(y=0.40, color='g', linestyle='--', alpha=0.5)
ax3.set_xlabel('Number of samples', fontsize=12)
ax3.set_ylabel('Estimated probability', fontsize=12)
ax3.set_title('Monte Carlo Estimation', fontsize=13)
ax3.legend(fontsize=9)
ax3.grid(True, alpha=0.3)

plt.tight_layout()
plt.show()