# Suppress warnings for a cleaner notebook
import warnings
import numpy as np
np.random.seed(42)
warnings.filterwarnings("ignore")
 
 

# Let us try a coin flip example!
import numpy as np
np.random.seed(42)

# We will flip a coin 20 times.
n_flips = 20
flips = np.random.choice([0,1], size=n_flips)  # 0=heads, 1=tails

print("Coin flip results:", flips)
print("Number of tails:", flips.sum())
 
 

# A Beta prior models our belief in the coin's fairness before seeing the data.
from scipy.stats import beta

prior_a, prior_b = 1, 1  # Uniform prior: every probability is equally likely

observed_tails = flips.sum()
observed_heads = n_flips - observed_tails

# Update: Beta posterior = Beta(prior_a + tails, prior_b + heads)
post_a = prior_a + observed_tails
post_b = prior_b + observed_heads

x = np.linspace(0, 1, 100)
prior_pdf = beta.pdf(x, prior_a, prior_b)
posterior_pdf = beta.pdf(x, post_a, post_b)

import matplotlib.pyplot as plt
plt.figure(figsize=(6, 4))
plt.plot(x, prior_pdf, label="Prior", linestyle="dashed")
plt.plot(x, posterior_pdf, label="Posterior")
plt.xlabel("Probability of Tails")
plt.ylabel("Density")
plt.legend()
plt.title("Prior vs. Posterior (Coin Flips)")
plt.show()
 
 

# Let us pretend we have a test for a rare disease.

# Prior: About 1 in 1000 people have the disease.
prior_disease = 1/1000

# The test gives positive for 99% of sick people,
# and 5% false positives for healthy people.
sensitivity = 0.99
false_pos = 0.05

# If you test positive, what is the chance you are truly sick?
num = sensitivity * prior_disease
den = num + false_pos * (1 - prior_disease)
posterior = num / den
print("Chance of disease if you test positive:", round(posterior, 4))
 
 

# Data setup
import seaborn as sns
titanic = sns.load_dataset("titanic")
print("Shape:", titanic.shape)
titanic.head()

# See how many survived and not survived.
titanic["survived"].value_counts()

# Compute the prior: overall chance of survival.
prior_survival = titanic["survived"].mean()
print("Estimated chance of surviving (prior):", round(prior_survival, 2))

# Now let us focus on young passengers under age 10.
young = titanic[titanic["age"] < 10]
n_young = young.shape[0]
n_young_survived = young["survived"].sum()
print("Young passengers:", n_young, "- Survived:", n_young_survived)

# Bayesian update for survival among young passengers.
prior_a, prior_b = 1, 1
post_a = prior_a + n_young_survived
post_b = prior_b + n_young - n_young_survived

posterior_pdf = beta.pdf(x, post_a, post_b)

plt.figure(figsize=(6, 4))
plt.plot(x, beta.pdf(x, prior_a, prior_b), label="Prior", linestyle="dashed")
plt.plot(x, posterior_pdf, label="Posterior (Young)")
plt.xlabel("Survival Probability")
plt.ylabel("Density")
plt.title("Bayesian Inference: Survival of Titanic Kids")
plt.legend()
plt.show()

# Let us do Bayesian inference with multiple groups: male vs female.
for sex in ['male', 'female']:
    group = titanic[titanic['sex'] == sex]
    survived = group['survived'].sum()
    total = group.shape[0]
    group_post_a = 1 + survived
    group_post_b = 1 + total - survived
    mean = group_post_a / (group_post_a + group_post_b)
    print(f"{sex.capitalize()}s posterior mean survival: {mean:.2f}")
    

# Interactive: Set your own prior!
guess = float(input("What is your prior guess of Titanic survival? (a value from 0 to 1): "))
prior_a = guess * 10 + 1
prior_b = (1 - guess) * 10 + 1

post_a = prior_a + titanic['survived'].sum()
post_b = prior_b + titanic.shape[0] - titanic['survived'].sum()

updated_mean = post_a / (post_a + post_b)
print("After seeing the data, updated survival estimate:", round(updated_mean, 2))

# Bayes' rule for the email problem.
prior_spam = 0.2
prob_free_given_spam = 0.7
prob_free_given_not_spam = 0.1

num = prob_free_given_spam * prior_spam
den = num + prob_free_given_not_spam * (1 - prior_spam)
posterior_spam = num / den
print("Chance of spam if 'FREE' appears:", round(posterior_spam, 2))

# Bayesian inference in practice: Avoid overconfidence.
print("If your data is limited, keep your prior broad. Too confident priors can hide the real answer!")

# Extra tip: Try Bayesian updates with more than two outcomes.
# Use the Dirichlet distribution for three or more groups.
print("Ready for a challenge? Explore Dirichlet for more categories, like favorite colors!")

# Your turn! Fill in blanks below to try it.
n_rolls = 12
n_six = 2
prior_a, prior_b = 1, 1
post_a = prior_a + n_six
post_b = prior_b + n_rolls - n_six
mean_prob = post_a / (post_a + post_b)
print("Posterior mean for rolling a six:", round(mean_prob, 3))