# boids_streamlit_advanced.py

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

# Define the Boid class
class Boid:
    def __init__(self, position, velocity):
        self.position = np.array(position, dtype='float64')
        self.velocity = np.array(velocity, dtype='float64')
        self.acceleration = np.zeros(2, dtype='float64')
        self.history = []

    def update(self, boids, width, height, params):
        self.flock(boids, params)
        self.velocity += self.acceleration
        speed = np.linalg.norm(self.velocity)
        if speed > params['max_speed']:
            self.velocity = (self.velocity / speed) * params['max_speed']
        self.position += self.velocity
        self.acceleration = np.zeros(2, dtype='float64')

        # Screen wrapping
        self.position[0] = self.position[0] % width
        self.position[1] = self.position[1] % height

        # Add current position to history
        if params['draw_trail']:
            self.history.append(self.position.copy())
            if len(self.history) > params['max_history']:
                self.history.pop(0)

    def flock(self, boids, params):
        self.acceleration = np.zeros(2, dtype='float64')
        self.fly_towards_center(boids, params)
        self.avoid_others(boids, params)
        self.match_velocity(boids, params)
        self.limit_speed(params)
        self.avoid_boundaries(width=params['width'], height=params['height'], params=params)

    def fly_towards_center(self, boids, params):
        centering_factor = params['centering_factor']
        center_x = 0.0
        center_y = 0.0
        num_neighbors = 0

        for other in boids:
            if other is not self and self.distance(other) < params['visual_range']:
                center_x += other.position[0]
                center_y += other.position[1]
                num_neighbors += 1

        if num_neighbors > 0:
            center_x /= num_neighbors
            center_y /= num_neighbors
            direction = np.array([center_x, center_y]) - self.position
            self.acceleration += centering_factor * direction

    def avoid_others(self, boids, params):
        min_distance = params['min_distance']
        avoid_factor = params['avoid_factor']
        move = np.zeros(2, dtype='float64')

        for other in boids:
            if other is not self and self.distance(other) < min_distance:
                move += self.position - other.position

        if np.linalg.norm(move) > 0:
            move = move / np.linalg.norm(move) * avoid_factor
            self.acceleration += move

    def match_velocity(self, boids, params):
        matching_factor = params['matching_factor']
        avg_velocity = np.zeros(2, dtype='float64')
        num_neighbors = 0

        for other in boids:
            if other is not self and self.distance(other) < params['visual_range']:
                avg_velocity += other.velocity
                num_neighbors += 1

        if num_neighbors > 0:
            avg_velocity /= num_neighbors
            self.acceleration += matching_factor * (avg_velocity - self.velocity)

    def limit_speed(self, params):
        speed = np.linalg.norm(self.velocity)
        if speed > params['max_speed']:
            self.velocity = (self.velocity / speed) * params['max_speed']

    def avoid_boundaries(self, width, height, params):
        margin = params['boundary_margin']
        turn_factor = params['boundary_turn_factor']

        if self.position[0] < margin:
            self.acceleration[0] += turn_factor
        elif self.position[0] > width - margin:
            self.acceleration[0] -= turn_factor

        if self.position[1] < margin:
            self.acceleration[1] += turn_factor
        elif self.position[1] > height - margin:
            self.acceleration[1] -= turn_factor

    def distance(self, other):
        return np.linalg.norm(self.position - other.position)

# Simulation parameters
params = {
    'num_boids': 100,
    'visual_range': 75.0,
    'min_distance': 20.0,
    'centering_factor': 0.005,
    'avoid_factor': 0.05,
    'matching_factor': 0.05,
    'max_speed': 15.0,
    'draw_trail': False,
    'max_history': 50,
    'width': 800,
    'height': 600,
    'boundary_margin': 100.0,
    'boundary_turn_factor': 0.05
}

# Streamlit sidebar for parameter adjustments
st.sidebar.title("Boids Simulation Parameters")
params['num_boids'] = st.sidebar.slider("Number of Boids", 10, 300, 100)
params['visual_range'] = st.sidebar.slider("Visual Range", 10.0, 200.0, 75.0)
params['min_distance'] = st.sidebar.slider("Minimum Separation Distance", 5.0, 100.0, 20.0)
params['centering_factor'] = st.sidebar.slider("Centering Factor", 0.001, 0.02, 0.005)
params['avoid_factor'] = st.sidebar.slider("Avoidance Factor", 0.01, 0.1, 0.05)
params['matching_factor'] = st.sidebar.slider("Matching Factor", 0.01, 0.1, 0.05)
params['max_speed'] = st.sidebar.slider("Maximum Speed", 5.0, 30.0, 15.0)
params['draw_trail'] = st.sidebar.checkbox("Draw Trails")
if params['draw_trail']:
    params['max_history'] = st.sidebar.slider("Trail Length", 10, 100, 50)
params['boundary_margin'] = st.sidebar.slider("Boundary Margin", 50.0, 300.0, 100.0)
params['boundary_turn_factor'] = st.sidebar.slider("Boundary Turn Factor", 0.01, 0.2, 0.05)

# Simulation screen size
width, height = 800, 600
params['width'] = width
params['height'] = height

# Initialize Boids
boids = []
for _ in range(params['num_boids']):
    position = [np.random.uniform(0, width), np.random.uniform(0, height)]
    angle = np.random.uniform(0, 2 * np.pi)
    velocity = [np.cos(angle), np.sin(angle)]
    boids.append(Boid(position, velocity))

# Plotly graph setup
fig = go.Figure(
    layout=go.Layout(
        xaxis=dict(range=[0, width], autorange=False, showgrid=False, zeroline=False),
        yaxis=dict(range=[0, height], autorange=False, showgrid=False, zeroline=False),
        width=width,
        height=height,
        margin=dict(l=0, r=0, t=0, b=0)
    )
)

# Plot initial positions of Boids
scatter = go.Scatter(
    x=[boid.position[0] for boid in boids],
    y=[boid.position[1] for boid in boids],
    mode='markers',
    marker=dict(size=8, color='blue')
)

fig.add_trace(scatter)

# Trail trace
if params['draw_trail']:
    trail_scatter = go.Scatter(
        x=[],
        y=[],
        mode='lines',
        line=dict(color='rgba(0,0,255,0.2)', width=1),
        showlegend=False
    )
    fig.add_trace(trail_scatter)

# Simplified Title
st.title("Boids Simulation")

# Animation display area
animation_placeholder = st.empty()

# Explanation Section Title
st.header("Mathematical Background of the Boids Algorithm")

# Explanation Section
st.markdown("### **Overview of the Boids Algorithm**")
st.markdown("""
The Boids algorithm, proposed by Craig Reynolds in 1986, is a method for simulating flocking behavior in groups of agents called Boids. Each agent follows simple rules to recreate complex group dynamics. The three fundamental rules are:

1. **Separation**: Maintain a suitable distance from nearby Boids to avoid collisions.
2. **Alignment**: Align velocity with the average velocity of neighboring Boids.
3. **Cohesion**: Move towards the average position of neighboring Boids.
""")

st.markdown("### **Mathematical Model**")
st.markdown("""
The movement of each Boid is represented by its position vector \(\mathbf{p}_i(t)\) and velocity vector \(\mathbf{v}_i(t)\). The position and velocity of Boid \(i\) at time \(t\) are described by the following differential equations:
""")

st.latex(r"""
\frac{d\mathbf{p}_i(t)}{dt} = \mathbf{v}_i(t)
""")
st.latex(r"""
\frac{d\mathbf{v}_i(t)}{dt} = \mathbf{a}_i(t)
""")

st.markdown("""
Here, the acceleration \(\mathbf{a}_i(t)\) is the sum of three forces:
""")
st.latex(r"""
\mathbf{a}_i(t) = \mathbf{a}_{\text{separation}} + \mathbf{a}_{\text{alignment}} + \mathbf{a}_{\text{cohesion}}
""")

st.markdown("#### **1. Separation**")
st.markdown("""
To prevent collisions, the separation force is calculated based on the distance \(d_{ij}\) between Boid \(i\) and its neighboring Boids \(j\):
""")
st.latex(r"""
\mathbf{a}_{\text{separation}} = \sum_{j \in N(i)} \frac{\mathbf{p}_i - \mathbf{p}_j}{d_{ij}^2}
""")
st.markdown("where \(N(i)\) is the set of neighboring Boids around Boid \(i\).")

st.markdown("#### **2. Alignment**")
st.markdown("""
The alignment force encourages Boid \(i\) to match the average velocity \(\mathbf{v}_{\text{avg}}\) of its neighbors:
""")
st.latex(r"""
\mathbf{a}_{\text{alignment}} = \frac{\mathbf{v}_{\text{avg}} - \mathbf{v}_i}{\tau}
""")
st.markdown("where \(\tau\) is a scaling parameter.")

st.markdown("#### **3. Cohesion**")
st.markdown("""
The cohesion force steers Boid \(i\) towards the average position \(\mathbf{C}_{\text{avg}}\) of its neighbors:
""")
st.latex(r"""
\mathbf{a}_{\text{cohesion}} = \frac{\mathbf{C}_{\text{avg}} - \mathbf{p}_i}{\sigma}
""")
st.markdown("where \(\sigma\) is a scaling parameter.")

st.markdown("### **Update Rules**")
st.markdown("""
Each Boid's position and velocity are updated based on discrete time steps \(\Delta t\) as follows:
""")
st.latex(r"""
\mathbf{v}_i(t + \Delta t) = \mathbf{v}_i(t) + \mathbf{a}_i(t) \Delta t
""")
st.latex(r"""
\mathbf{p}_i(t + \Delta t) = \mathbf{p}_i(t) + \mathbf{v}_i(t + \Delta t) \Delta t
""")
st.markdown("""
These equations ensure that each Boid updates its velocity and position based on the combined separation, alignment, and cohesion forces.
""")

st.markdown("### **Additional Features**")
st.markdown("""
This simulation includes the following additional features:

1. **Boundary Avoidance**: When a Boid approaches the edge of the simulation area, it receives a steering force to remain within bounds, preventing it from moving off-screen.
2. **Trail Drawing**: The past positions of each Boid are displayed as trails, allowing visualization of their movement patterns.
""")

st.markdown("### **Parameters and Their Roles**")
st.markdown("""
- **Boundary Margin (\(M\))**: The distance from the edge of the simulation area at which Boids begin to steer away.
- **Boundary Turn Factor (\(\gamma\))**: The strength of the steering force applied when avoiding boundaries.

These parameters allow fine-tuning of Boid behavior near the edges of the simulation area.
""")


# Animation settings
frame_rate = 30  # Frames per second
sleep_time = 1.0 / frame_rate

# Reset button
if st.sidebar.button("Reset Simulation"):
    boids = []
    for _ in range(params['num_boids']):
        position = [np.random.uniform(0, width), np.random.uniform(0, height)]
        angle = np.random.uniform(0, 2 * np.pi)
        velocity = [np.cos(angle), np.sin(angle)]
        boids.append(Boid(position, velocity))

# Animation loop
while True:
    # Update Boids
    for boid in boids:
        boid.update(boids, width, height, params)

    # Update Boids' positions
    scatter.x = [boid.position[0] for boid in boids]
    scatter.y = [boid.position[1] for boid in boids]

    # Update trails
    if params['draw_trail']:
        trail_x = []
        trail_y = []
        for boid in boids:
            trail_x.extend([pos[0] for pos in boid.history])
            trail_y.extend([pos[1] for pos in boid.history])
        trail_scatter.x = trail_x
        trail_scatter.y = trail_y

    # Update the Plotly figure
    fig.data[0].x = scatter.x
    fig.data[0].y = scatter.y
    if params['draw_trail']:
        fig.data[1].x = trail_scatter.x
        fig.data[1].y = trail_scatter.y

    # Display the animation
    animation_placeholder.plotly_chart(fig, use_container_width=True)

    # Control the frame rate
    time.sleep(sleep_time)