app.py

import streamlit as st
from copy import deepcopy
from random import choice

# Define ball colors and corresponding emojis
ball_colors = {
    "red": "🔴",
    "blue": "🔵",
    "green": "🟢",
    "yellow": "🟡",
    "black": "⚫",
    "white": "⚪",
    "orange": "🟠",
    "pink": "🟣",
    "brown": "🟤",
}

# Initialize session state for the hat
if "hat_balls" not in st.session_state:
    st.session_state.hat_balls = {color: 0 for color in ball_colors.keys()}

# Initialize session state for the expected balls
if "expected_balls" not in st.session_state:
    st.session_state.expected_balls = {color: 0 for color in ball_colors.keys()}


def add_ball(state_key, color):
    """Add a ball to the selected state key (hat_balls or expected_balls)."""
    total_balls = sum(st.session_state[state_key].values())
    if total_balls < 20:  # Limit to 20 balls
        st.session_state[state_key][color] += 1


# Layout for defining the hat and expected balls
st.title("🎩 Probability Calculator")
st.write("This is a probability calculator for picking balls from a hat without replacement.")

hat_col, experiment_col = st.columns(2)

# Hat Configuration
with hat_col:
    st.subheader("Define the Hat")
    st.write("Click a ball to add it to the hat.")
    
    # Arrange buttons in a grid (3 per row)
    cols = st.columns(3)
    for i, (color, emoji) in enumerate(ball_colors.items()):
        with cols[i % 3]:
            if st.button(emoji, key=f"hat_{color}"):
                add_ball("hat_balls", color)

    # Display the hat's contents with emojis
    hat_display = "".join(
        emoji * count for color, (emoji, count) in zip(ball_colors.keys(), zip(ball_colors.values(), st.session_state.hat_balls.values()))
    )
    st.text_area("🎩 Current Hat", value=f"🎩 = {hat_display}", height=100)

# Experiment Parameters
with experiment_col:
    st.subheader("Experiment Parameters")
    st.write("Click a ball to add it to the expected balls.")
    
    # Arrange buttons in a grid (3 per row)
    cols = st.columns(3)
    for i, (color, emoji) in enumerate(ball_colors.items()):
        with cols[i % 3]:
            if st.button(emoji, key=f"expected_{color}"):
                add_ball("expected_balls", color)

    # Display the expected balls with emojis
    expected_display = "".join(
        emoji * count for color, (emoji, count) in zip(ball_colors.keys(), zip(ball_colors.values(), st.session_state.expected_balls.values()))
    )
    st.text_area("🎯 Expected Balls", value=f"🎯 = {expected_display}", height=100)

# Number of balls drawn and experiments
st.write("### Experiment Setup")
draw_col, experiment_count_col = st.columns(2)

with draw_col:
    num_balls_drawn = st.number_input("Number of balls draw from the hat:", min_value=1, value=5)

with experiment_count_col:
    num_experiments = st.number_input("Number of experiments:", min_value=1, value=2000)

# Run Experiment Button
if st.button("Run Experiment"):
    try:
        # Define the Hat class
        class Hat:
            def __init__(self, **kwargs):
                self.contents = [color for color, count in kwargs.items() for _ in range(count)]

            def draw(self, n):
                if n >= len(self.contents):
                    drawn = self.contents[:]
                    self.contents = []
                    return drawn
                drawn = []
                for _ in range(n):
                    ball = choice(self.contents)
                    self.contents.remove(ball)
                    drawn.append(ball)
                return drawn

        def experiment(hat, expected_balls, num_balls_drawn, num_experiments):
            match = 0
            for _ in range(num_experiments):
                # Create a fresh copy of the hat
                hat_copy = deepcopy(hat)

                # Draw balls
                balls_drawn = hat_copy.draw(num_balls_drawn)

                # Check if at least the required count of each color is present
                if all(balls_drawn.count(color) >= count for color, count in expected_balls.items()):
                    match += 1

            # Calculate the probability
            return match / num_experiments

        # Run the experiment
        hat = Hat(**st.session_state.hat_balls)
        expected = {color: count for color, count in st.session_state.expected_balls.items() if count > 0}
        probability = experiment(hat, expected, num_balls_drawn, num_experiments)

        st.subheader(f"The probability is {probability:.4f}%")

    except Exception as e:
        st.error(f"An error occurred: {e}")