import streamlit as st
import numpy as np
import pandas as pd
import time
import plotly.graph_objects as go

st.set_page_config(page_title="Urn Sampling Demo", layout="wide")

st.title("Urn Sampling and Relative Frequencies")

# --------------------------------------------------
# Controls
# --------------------------------------------------

col1, col2, col3 = st.columns(3)

with col1:
    n0 = st.number_input("Number of 0-balls", 0, 1000, 2)

with col2:
    n1 = st.number_input("Number of 1-balls", 0, 1000, 3)

with col3:
    n2 = st.number_input("Number of 2-balls", 0, 1000, 5)

draws = st.slider("Number of draws", 10, 5000, 500)

replacement = st.checkbox("Draw with replacement", value=True)

speed = st.slider(
    "Animation delay (seconds)",
    0.0,
    0.2,
    0.01,
    0.01
)

# --------------------------------------------------
# Build urn
# --------------------------------------------------

urn = [0] * n0 + [1] * n1 + [2] * n2

if len(urn) == 0:
    st.warning("Add at least one ball.")
    st.stop()

p0 = n0 / len(urn)
p1 = n1 / len(urn)
p2 = n2 / len(urn)

st.markdown(
    f"""
    **True probabilities**

    - P(0) = {p0:.3f}
    - P(1) = {p1:.3f}
    - P(2) = {p2:.3f}
    """
)

start = st.button("Start Simulation")

# --------------------------------------------------
# Simulation
# --------------------------------------------------

if start:

    count0 = 0
    count1 = 0
    count2 = 0

    freq0 = []
    freq1 = []
    freq2 = []

    draw_index = []

    latest_draw_placeholder = st.empty()
    stats_placeholder = st.empty()
    chart_placeholder = st.empty()

    current_urn = urn.copy()

    for i in range(1, draws + 1):

        if replacement:
            draw = np.random.choice(urn)

        else:
            if len(current_urn) == 0:
                break

            idx = np.random.randint(len(current_urn))
            draw = current_urn.pop(idx)

        if draw == 0:
            count0 += 1
        elif draw == 1:
            count1 += 1
        else:
            count2 += 1

        freq0.append(count0 / i)
        freq1.append(count1 / i)
        freq2.append(count2 / i)

        draw_index.append(i)

        latest_draw_placeholder.metric(
            "Latest Draw",
            int(draw)
        )

        stats_placeholder.write(
            {
                "Count(0)": count0,
                "Count(1)": count1,
                "Count(2)": count2,
            }
        )

        fig = go.Figure()

        fig.add_trace(
            go.Scatter(
                x=draw_index,
                y=freq0,
                mode="lines",
                name="Relative freq(0)"
            )
        )

        fig.add_trace(
            go.Scatter(
                x=draw_index,
                y=freq1,
                mode="lines",
                name="Relative freq(1)"
            )
        )

        fig.add_trace(
            go.Scatter(
                x=draw_index,
                y=freq2,
                mode="lines",
                name="Relative freq(2)"
            )
        )

        fig.add_hline(
            y=p0,
            line_dash="dash",
            annotation_text="P(0)"
        )

        fig.add_hline(
            y=p1,
            line_dash="dash",
            annotation_text="P(1)"
        )

        fig.add_hline(
            y=p2,
            line_dash="dash",
            annotation_text="P(2)"
        )

        fig.update_layout(
            title="Convergence of Relative Frequencies",
            xaxis_title="Draw Number",
            yaxis_title="Relative Frequency",
            yaxis=dict(range=[0, 1]),
            height=500
        )

        chart_placeholder.plotly_chart(
            fig,
            use_container_width=True
        )

        if speed > 0:
            time.sleep(speed)