Py.Cafe

import dash
import dash_bootstrap_components as dbc
import dash_core_components as dcc
import dash_html_components as html
from dash.dependencies import Input, Output, State
from PIL import Image, ImageDraw
import base64
from io import BytesIO
import numpy as np
import random
import plotly.graph_objs as go

# Initialize the app
app = dash.Dash(__name__, external_stylesheets=[dbc.themes.BOOTSTRAP])

# Canvas dimensions
canvas_width = 600
canvas_height = 400
ball_radius = 20

# Initial ball state
num_balls = 5
initial_ball_state = [
    {
        'x': random.randint(ball_radius, canvas_width - ball_radius),
        'y': random.randint(ball_radius, canvas_height - ball_radius),
        'vx': random.choice([-3, 3]),
        'vy': random.choice([-3, 3])
    } for _ in range(num_balls)
]

# Define the layout
app.layout = dbc.Container(
    [
        dbc.Row(
            dbc.Col(html.H1("Bouncing Ball Simulation"), className="text-center")
        ),
        dbc.Row(
            dbc.Col(
                html.Div([
                    html.Img(id='ball-canvas', width=canvas_width, height=canvas_height, style={"border": "1px solid #000"}),
                ]),
                width=12,
            )
        ),
        dbc.Row(
            dbc.Col(
                dcc.Slider(
                    id='num-balls-slider',
                    min=1,
                    max=10,
                    step=1,
                    value=num_balls,
                    marks={i: str(i) for i in range(1, 11)},
                    className="mb-4"
                ),
                width=12,
            )
        ),
        dbc.Row(
            dbc.Col(
                dcc.Slider(
                    id='velocity-slider',
                    min=1,
                    max=10,
                    step=1,
                    value=3,
                    marks={i: str(i) for i in range(1, 11)},
                    className="mb-4"
                ),
                width=12,
            )
        ),
        dbc.Row(
            dbc.Col(html.Div(id='collision-counter', className="text-center"), width=12)
        ),
        dbc.Row(
            dbc.Col(dcc.Graph(id='collision-histogram'), width=12)
        ),
        dcc.Interval(id='interval-component', interval=30, n_intervals=0),
        dcc.Store(id='ball-state', data=initial_ball_state),
        dcc.Store(id='collision-count', data=0),
        dcc.Store(id='collision-locations', data=[])
    ],
    fluid=True,
)

def detect_collision(ball1, ball2):
    dx = ball1['x'] - ball2['x']
    dy = ball1['y'] - ball2['y']
    distance = np.sqrt(dx**2 + dy**2)
    return distance < 2 * ball_radius

def handle_collision(ball1, ball2):
    dx = ball1['x'] - ball2['x']
    dy = ball1['y'] - ball2['y']
    distance = np.sqrt(dx**2 + dy**2)

    if distance == 0:
        return ball1, ball2  # Avoid division by zero

    nx = dx / distance
    ny = dy / distance

    kx = (ball1['vx'] - ball2['vx'])
    ky = (ball1['vy'] - ball2['vy'])
    p = 2.0 * (nx * kx + ny * ky) / 2

    ball1['vx'] = ball1['vx'] - p * nx
    ball1['vy'] = ball1['vy'] - p * ny
    ball2['vx'] = ball2['vx'] + p * nx
    ball2['vy'] = ball2['vy'] + p * ny

    return ball1, ball2

def update_ball_state(ball_state, velocity_factor, collision_locations):
    collision_count = 0
    for ball in ball_state:
        x, y = ball['x'], ball['y']
        vx, vy = ball['vx'], ball['vy']

        # Update ball position
        x += vx * velocity_factor
        y += vy * velocity_factor

        # Check for collisions with the walls
        if x - ball_radius < 0 or x + ball_radius > canvas_width:
            vx = -vx
        if y - ball_radius < 0 or y + ball_radius > canvas_height:
            vy = -vy

        ball.update({'x': x, 'y': y, 'vx': vx, 'vy': vy})

    # Check for collisions between balls
    for i in range(len(ball_state)):
        for j in range(i + 1, len(ball_state)):
            if detect_collision(ball_state[i], ball_state[j]):
                ball_state[i], ball_state[j] = handle_collision(ball_state[i], ball_state[j])
                collision_count += 1
                collision_locations.append((ball_state[i]['x'], ball_state[i]['y']))

    return ball_state, collision_count, collision_locations

@app.callback(
    [Output('ball-canvas', 'src'),
     Output('ball-state', 'data'),
     Output('collision-counter', 'children'),
     Output('collision-count', 'data'),
     Output('collision-locations', 'data'),
     Output('collision-histogram', 'figure')],
    [Input('interval-component', 'n_intervals'),
     Input('num-balls-slider', 'value'),
     Input('velocity-slider', 'value')],
    [State('ball-state', 'data'),
     State('collision-count', 'data'),
     State('collision-locations', 'data')]
)
def update_canvas(n_intervals, num_balls, velocity, ball_state, total_collisions, collision_locations):
    velocity_factor = velocity / 3  # Adjust velocity based on slider value

    if len(ball_state) != num_balls:
        ball_state = [
            {
                'x': random.randint(ball_radius, canvas_width - ball_radius),
                'y': random.randint(ball_radius, canvas_height - ball_radius),
                'vx': random.choice([-3, 3]),
                'vy': random.choice([-3, 3])
            } for _ in range(num_balls)
        ]

    ball_state, collision_count, collision_locations = update_ball_state(ball_state, velocity_factor, collision_locations)
    total_collisions += collision_count

    # Create an empty image with white background
    img = Image.new('RGB', (canvas_width, canvas_height), color=(255, 255, 255))
    draw = ImageDraw.Draw(img)

    # Draw the balls
    for ball in ball_state:
        x, y = ball['x'], ball['y']
        draw.ellipse((x - ball_radius, y - ball_radius, x + ball_radius, y + ball_radius), fill=(0, 0, 0))

    # Convert the image to a base64 string
    buffered = BytesIO()
    img.save(buffered, format="PNG")
    img_str = base64.b64encode(buffered.getvalue()).decode("utf-8")

    collision_text = f"Total Collisions: {total_collisions}"

    if collision_locations:
        collision_x, collision_y = zip(*collision_locations)
    else:
        collision_x, collision_y = [], []

    # Create the 3D histogram
    histogram_fig = go.Figure(data=[go.Histogram2d(
        x=collision_x,
        y=collision_y,
        colorscale='Viridis',
        autobinx=False,
        autobiny=False,
        nbinsx=20,
        nbinsy=20
    )])
    histogram_fig.update_layout(
        scene=dict(
            xaxis=dict(title='X'),
            yaxis=dict(title='Y')
        ),
        margin=dict(l=0, r=0, b=0, t=0)
    )

    return f"data:image/png;base64,{img_str}", ball_state, collision_text, total_collisions, collision_locations, histogram_fig

# Run the app
if __name__ == "__main__":
    app.run_server(debug=True)