import streamlit as st
import matplotlib.pyplot as plt
import numpy as np
import matplotlib.animation as animation
import tempfile
import base64

# Set up page configuration
st.set_page_config(page_title="Hithesh Rai - Interactive Bohr Model of the Atom", page_icon="🔬")

# Header
st.title("Interactive Bohr Model of the Atom")
st.subheader("Visualize electrons orbiting around a nucleus")

# About Section
st.write("""
Explore a simple interactive model of an atom based on the Bohr model, where electrons orbit a nucleus in defined circular paths. Use the controls to set the number of electrons and watch them orbit in real time!
""")

# Input for number of electrons
num_electrons = st.slider("Select the number of electrons", 1, 5, 3)

# Create figure and axis
fig, ax = plt.subplots(figsize=(6, 6))
ax.set_xlim(-10, 10)
ax.set_ylim(-10, 10)
ax.set_aspect('equal')
ax.axis('off')

# Set up nucleus
nucleus, = ax.plot(0, 0, 'o', color='orange', markersize=10, label="Nucleus")

# Colors for electron orbits
orbit_colors = ['blue', 'green', 'red', 'purple', 'cyan']

# Set up electron orbits based on selected number
electrons = []
orbits = []
for i in range(num_electrons):
    radius = 2 + 2 * i
    theta = np.linspace(0, 2 * np.pi, 100)
    x = radius * np.cos(theta)
    y = radius * np.sin(theta)
    orbit, = ax.plot(x, y, linestyle="--", color=orbit_colors[i % len(orbit_colors)], alpha=0.5)
    electron, = ax.plot([], [], 'o', color=orbit_colors[i % len(orbit_colors)], markersize=6)
    electrons.append(electron)
    orbits.append(orbit)

# Animation update function
def update(frame):
    for i, electron in enumerate(electrons):
        radius = 2 + 2 * i
        x = radius * np.cos(frame * 0.1 + i * np.pi / num_electrons)
        y = radius * np.sin(frame * 0.1 + i * np.pi / num_electrons)
        electron.set_data([x], [y])  # Wrap x and y in lists to avoid RuntimeError
    return electrons

# Generate animation and save as temporary HTML file
with tempfile.NamedTemporaryFile(delete=False, suffix=".html") as tmpfile:
    ani = animation.FuncAnimation(fig, update, frames=range(200), interval=50, blit=True)
    ani.save(tmpfile.name, writer="html", fps=20)
    tmpfile.seek(0)
    html_data = tmpfile.read().decode("utf-8")

# Display animation in Streamlit
st.components.v1.html(html_data, height=500)