import sys
import matplotlib.pyplot as plt
import matplotlib.animation as animation
import numpy as np
import random


# パラメータ設定：段数 n=10, ボール数 100, 速度 50ms
n = 20
num_balls = 1000
speed_ms = 10 # ms

argv = sys.argv
nargs = len(argv)
if nargs > 1: n = int(argv[1])
if nargs > 2: num_balls = int(argv[2])
if nargs > 3: speed_ms = int(argv[3])


def simulate_galton_board(n=10, num_balls=100, speed_ms=50):
    """
    n: 釘の段数
    num_balls: 落下させる鉄球の総数
    speed_ms: アニメーションの間隔（ミリ秒）。小さくするほど高速。
    """
    # グラフの設定
    fig, (ax1, ax2) = plt.subplots(2, 1, figsize=(6, 8), gridspec_kw={'height_ratios': [2, 1]})
    
    # 1. 釘の座標を設定（n段の三角形）
    pin_x = []
    pin_y = []
    for row in range(n):
        y = n - row
        for i in range(row + 1):
            pin_x.append(i - row * 0.5)
            pin_y.append(y)
            
    # 2. ヒストグラム用の設定
    # n段の場合、底には n+1 個のポケットができる
    bins = np.arange(-(n/2 + 0.5), (n/2 + 1.5), 1)
    bin_counts = np.zeros(n + 1)
    bin_centers = (bins[:-1] + bins[1:]) / 2

    # 全ての軌跡を事前に生成（高速化のため）
    all_paths = []
    final_positions = []
    for _ in range(num_balls):
        path_x = [0]
        path_y = [n + 0.5]
        curr_x = 0
        for row in range(n):
            step = 0.5 if random.random() > 0.5 else -0.5
            curr_x += step
            path_x.append(curr_x)
            path_y.append(n - row)
        
        path_x.append(curr_x)
        path_y.append(0)
        all_paths.append((path_x, path_y))
        final_positions.append(curr_x)

    def update(frame):
        ax1.clear()
        ax2.clear()
        
        # 上段：釘と現在の球の軌跡
        ax1.scatter(pin_x, pin_y, color='#555555', s=30, zorder=3) # 釘（少し小さめに）
        px, py = all_paths[frame]
        ax1.plot(px, py, color='#FF4B4B', linewidth=2, marker='o', markersize=5, zorder=4) # 軌跡
        
        ax1.set_xlim(-(n/2 + 1), (n/2 + 1))
        ax1.set_ylim(-0.5, n + 1)
        ax1.set_title(f"Galton Board (n={n}, Ball {frame+1}/{num_balls})", fontsize=14)
        ax1.axis('off')

        # 下段：累積分布
        current_final = final_positions[frame]
        idx = np.digitize(current_final, bins) - 1
        bin_counts[idx] += 1
        
        colors = plt.cm.viridis(np.linspace(0, 1, n + 1))
        ax2.bar(bin_centers, bin_counts, width=0.8, color=colors, edgecolor='white', alpha=0.8)
        ax2.set_xlim(-(n/2 + 1), (n/2 + 1))
        # y軸のスケールを自動調整
        ax2.set_ylim(0, max(bin_counts) + 2)
        ax2.set_title("Cumulative Binomial Distribution", fontsize=12)
        ax2.set_ylabel("Frequency")

    # intervalで速度を調整（50ms = 秒間20フレーム）
    ani = animation.FuncAnimation(fig, update, frames=num_balls, repeat=False, interval=speed_ms)
    plt.tight_layout()
    plt.show()

simulate_galton_board(n=n, num_balls=num_balls, speed_ms=speed_ms)