from vpython import *
import numpy as np
import matplotlib.pyplot as plt
from matplotlib.animation import FuncAnimation
from matplotlib.ticker import MultipleLocator, FuncFormatter

# Constants
g = 9.81  # acceleration due to gravity, in m/s^2
L1 = 1.0  # length of first pendulum in meters
L2 = 1.0  # length of second pendulum in meters
m1 = 1.0  # mass of first pendulum in kg
m2 = 1.0  # mass of second pendulum in kg

# Initial conditions
theta1 = np.pi / 2  # initial angle of first pendulum
theta2 = np.pi / 2  # initial angle of second pendulum
omega1 = 0.0  # initial angular velocity of first pendulum
omega2 = 0.0  # initial angular velocity of second pendulum

# Time parameters
dt = 0.01  # time step in seconds
t = 0  # initial time

# Create the pendulum rods and bobs
rod1 = cylinder(pos=vector(0, 0, 0), axis=vector(L1*np.sin(theta1), -L1*np.cos(theta1), 0), radius=0.02)
bob1 = sphere(pos=rod1.axis, radius=0.05, color=color.red)
rod2 = cylinder(pos=bob1.pos, axis=vector(L2*np.sin(theta2), -L2*np.cos(theta2), 0), radius=0.02)
bob2 = sphere(pos=rod2.pos, radius=0.05, color=color.blue)

# Create a fixed reference frame
x_axis = arrow(pos=vector(-2, 0, 0), axis=vector(4, 0, 0), shaftwidth=0.01, color=color.yellow)
y_axis = arrow(pos=vector(0, -2, 0), axis=vector(0, 4, 0), shaftwidth=0.01, color=color.yellow)
z_axis = arrow(pos=vector(0, 0, -2), axis=vector(0, 0, 4), shaftwidth=0.01, color=color.yellow)
grid_lines = []
for i in range(-10, 11):
    grid_lines.append(curve(pos=[vector(i/5, -2, 0), vector(i/5, 2, 0)], radius=0.005, color=color.gray(0.5)))
    grid_lines.append(curve(pos=[vector(-2, i/5, 0), vector(2, i/5, 0)], radius=0.005, color=color.gray(0.5)))

# Set up the camera
scene.camera.pos = vector(0, -1, 2)  # Position of the camera
scene.camera.axis = vector(0, 1, -1)  # Direction the camera is pointing
scene.camera.up = vector(0, 1, 0)  # Up direction of the camera

# Function to update the positions of the pendulum
def update_positions():
    rod1.axis = vector(L1*np.sin(theta1), -L1*np.cos(theta1), 0)
    bob1.pos = rod1.pos + rod1.axis
    rod2.pos = bob1.pos
    rod2.axis = vector(L2*np.sin(theta2), -L2*np.cos(theta2), 0)
    bob2.pos = rod2.pos + rod2.axis

# Function to compute the derivatives
def derivatives(theta1, theta2, omega1, omega2):
    delta = theta2 - theta1
    den1 = (m1 + m2) * L1 - m2 * L1 * np.cos(delta) * np.cos(delta)
    den2 = (L2/L1) * den1
    
    domega1 = (m2 * L1 * omega1 * omega1 * np.sin(delta) * np.cos(delta) + 
               m2 * g * np.sin(theta2) * np.cos(delta) + 
               m2 * L2 * omega2 * omega2 * np.sin(delta) - 
               (m1 + m2) * g * np.sin(theta1)) / den1
    domega2 = (-L2/L1 * omega2 * omega2 * np.sin(delta) * np.cos(delta) + 
               (m1 + m2) * g * np.sin(theta1) * np.cos(delta) - 
               (m1 + m2) * L1 * omega1 * omega1 * np.sin(delta) - 
               (m1 + m2) * g * np.sin(theta2)) / den2
    
    return domega1, domega2

# Set up Matplotlib for real-time plotting
fig, ax = plt.subplots()
ax.set_xlim(0, 10)
ax.set_ylim(-np.pi, np.pi)
line1, = ax.plot([], [], label='theta1')
line2, = ax.plot([], [], label='theta2')
ax.legend()

# Format y-axis to display in multiples of pi
def format_func(value, tick_number):
    N = int(np.round(value / np.pi))
    if N == 0:
        return "0"
    elif N == 1:
        return r"$\pi$"
    elif N == -1:
        return r"-$\pi$"
    else:
        return r"${0}\pi$".format(N)

ax.yaxis.set_major_locator(MultipleLocator(base=np.pi/2))  # Set major ticks in multiples of pi/2
ax.yaxis.set_major_formatter(FuncFormatter(format_func))

time_data = []
theta1_data = []
theta2_data = []

def init():
    line1.set_data([], [])
    line2.set_data([], [])
    return line1, line2

def update(frame):
    global theta1, theta2, omega1, omega2, t
    # Compute derivatives
    domega1, domega2 = derivatives(theta1, theta2, omega1, omega2)
    
    # Update angular velocities and angles
    omega1 += domega1 * dt
    omega2 += domega2 * dt
    theta1 += omega1 * dt
    theta2 += omega2 * dt
    
    # Update positions
    update_positions()
    
    # Increment time
    t += dt

    # Update data for plotting
    time_data.append(t)
    theta1_data.append(theta1)
    theta2_data.append(theta2)
    
    line1.set_data(time_data, theta1_data)
    line2.set_data(time_data, theta2_data)
    
    return line1, line2

# Animation function
ani = FuncAnimation(fig, update, frames=np.arange(0, 1000), init_func=init, blit=True, interval=10)

# Show the plot
plt.show()

# Main loop for VPython
while True:
    rate(100)
    update(0)  # Call update function to synchronize with the VPython loop