import numpy as np
import matplotlib.pyplot as plt

# 1. Simulation of the sequential arrival of data
data = [1, 1, 0, 1, 1, 1, 0, 1, 1, 1]  # 1 - increase, 0 - decrease

# 2. Initial prior parameters (Uniform Prior)
alpha_current = 1
beta_current = 1

history = []

print(
    f"{'Step':<5} | {'Data':<7} | {'Alpha':<6} | {'Beta':<5} | {'MAP':<6} | {'Mean (Forecast)':<10}"
)
print("-" * 60)

# 3. Sequential updating loop (Bayesian Updating)
for i, y in enumerate(data):

    # UPDATE: Adding a success or failure
    if y == 1:
        alpha_current += 1
    else:
        beta_current += 1

    # MAP (mode)
    if alpha_current + beta_current > 2:
        map_val = (alpha_current - 1) / (alpha_current + beta_current - 2)
    else:
        map_val = 0.5  # For the uniform distribution, the mode is not uniquely defined

    mean_val = alpha_current / (alpha_current + beta_current)

    history.append((alpha_current, beta_current, mean_val))

    print(
        f"{i+1:<5} | {y:<7} | {alpha_current:<6} | {beta_current:<5} | {map_val:<6.3f} | {mean_val:<10.3f}"
    )

# 4. Visualizing the Learning Process
steps = np.arange(1, len(data) + 1)
means = [h[2] for h in history]

plt.figure(figsize=(8, 4), dpi=100)
plt.step(
    steps, means, where="post", color="#3498db", lw=2, label="Posterior Mean (Forecast)"
)
plt.axhline(0.5, color="gray", linestyle="--", alpha=0.5, label="Initial Prior (0.5)")

plt.title("Step-by-step Bayesian Updating", fontsize=12)
plt.xlabel("Observation Number (Time)", fontsize=10)
plt.ylabel("Probability of Success", fontsize=10)
plt.ylim(0, 1)
plt.grid(True, linestyle=":", alpha=0.6)
plt.legend()
plt.show()