from bokeh.plotting import figure, output_file, show
from bokeh.io import output_notebook
from bokeh.models import ColumnDataSource

# Set Bokeh output to notebook
output_notebook()

import numpy as np
from scipy.integrate import solve_ivp
from bokeh.plotting import figure, show

# Double pendulum parameters
g = 9.8  # gravity (m/s^2)
L1 = 1.0  # length of the first pendulum (m)
L2 = 1.0  # length of the second pendulum (m)
m1 = 1.0  # mass of the first pendulum (kg)
m2 = 1.0  # mass of the second pendulum (kg)

# Double pendulum equations
def double_pendulum(t, y):
    theta1, z1, theta2, z2 = y
    delta = theta2 - theta1

    denominator1 = (m1 + m2) * L1 - m2 * L1 * np.cos(delta)**2
    denominator2 = (L2 / L1) * denominator1

    theta1_dot = z1
    theta2_dot = z2

    z1_dot = (
        m2 * L1 * z1**2 * np.sin(delta) * np.cos(delta)
        + m2 * g * np.sin(theta2) * np.cos(delta)
        + m2 * L2 * z2**2 * np.sin(delta)
        - (m1 + m2) * g * np.sin(theta1)
    ) / denominator1

    z2_dot = (
        -m2 * L2 * z2**2 * np.sin(delta) * np.cos(delta)
        + (m1 + m2) * g * np.sin(theta1) * np.cos(delta)
        - (m1 + m2) * L1 * z1**2 * np.sin(delta)
        - (m1 + m2) * g * np.sin(theta2)
    ) / denominator2

    return [theta1_dot, z1_dot, theta2_dot, z2_dot]

# Initial conditions and time steps
theta1_0 = np.pi / 2  # initial angle of pendulum 1
theta2_0 = np.pi / 4  # initial angle of pendulum 2
z1_0 = 0.0  # initial angular velocity of pendulum 1
z2_0 = 0.0  # initial angular velocity of pendulum 2
y0 = [theta1_0, z1_0, theta2_0, z2_0]

t = np.linspace(0, 20, 10000)  # time array

# Solve the system
solution = solve_ivp(double_pendulum, [t[0], t[-1]], y0, t_eval=t, method="RK45")
theta1 = solution.y[0]
theta2 = solution.y[2]

# Convert to Cartesian coordinates
x1 = L1 * np.sin(theta1)
y1 = -L1 * np.cos(theta1)
x2 = x1 + L2 * np.sin(theta2)
y2 = y1 - L2 * np.cos(theta2)

# Use the trajectories of the second pendulum for plotting
num_segments = 7
xs = np.array_split(x2, num_segments)
ys = np.array_split(y2, num_segments)

# Define a color palette
colors = ["#F7FCF0", "#E0F3DB", "#CCEBC5", "#A8DDB5", "#7BCCC4", "#43A2CA", "#0868AC"]

# Create the Bokeh figure
p = figure(title="Double Pendulum Visualization",
           background_fill_color="#f9f9f9",
           x_axis_label="X",
           y_axis_label="Y",
           match_aspect=True)

# Add the multi_line glyph
p.multi_line(xs, ys, line_color=colors, line_alpha=0.8, line_width=1.5)

# Show the plot
show(p)