# Pygame emulation of the double pendulum
# Author: Brandon Berisford
# Date: 5/12/2021


import pygame
import numpy as np
from numpy import pi
from numpy.linalg import inv
from math import sin, cos
import sys
from pygame.locals import *


# Defining Variables and Constants

m1, m2 = 3.0, 5.0  # massses in kg
l1, l2 = 1.5, 2.0  # rod lengths in m
a1, a2 = 0.0, 0.0  # angles in radians
t = 0.0  # current time
dt = 0.02  # time step
w, h = 800, 480  # height and width of screen
g = 9.81  # gravitational acceleration
WHITE = (255, 255, 255)  # white in RGB
BLACK = (0, 0, 0)  # black in RBG
RED = (255, 0, 0)  # red in RGB
BLUE = (0, 0, 255)  # blue in RBG
LT_BLUE = (230, 230, 255)
offset = (400, 50)  # offset point to place the pendulum on the screen
prev_point = None

# Initial State [a1velocity, a2velocity, angle1, angle2]
y = np.array([0.0, 0.0, pi/2, pi/2])

screen = pygame.display.set_mode((w, h))  # sets the screen width and height
screen.fill(WHITE)  # fills the screen with the color white
trace = screen.copy()  # trace surface for drawing path trace of masses
pygame.display.update()  # updates the display
clock = pygame.time.Clock()  # starts the game clock
pygame.font.init()
myfont = pygame.font.SysFont('Comic Sans MS', 38)


# Calculates the slope of y at time t
def G(y, t):
    a1d, a2d = y[0], y[1]
    a1, a2 = y[2], y[3]

    m11, m12 = (m1+m2)*l1, m2*l2*cos(a1-a2)
    m21, m22 = l1*cos(a1-a2), l2
    m = np.array([[m11, m12], [m21, m22]])

    f1 = -m2*l2*a2d*a2d*sin(a1-a2) - (m1+m2)*g*sin(a1)
    f2 = l1*a1d*a1d*sin(a1-a2) - g*sin(a2)
    f = np.array([f1, f2])

    accel = inv(m).dot(f)

    return np.array([accel[0], accel[1], a1d, a2d])


# implements the runga kutta method to calculate RK4 part of the updated position
def RK4_step(y, t, dt):
    k1 = G(y, t)
    k2 = G(y+0.5*k1*dt, t+0.5*dt)
    k3 = G(y+0.5*k2*dt, t+0.5*dt)
    k4 = G(y+k3*dt, t+dt)

    return dt * (k1 + 2*k2 + 2*k3 + k4) / 6


# update method that runs every frame
def update(angle1, angle2):
    scale = 100  # Scale the points by 100
    # x position of first mass with offset
    x1 = l1*scale*np.sin(angle1) + offset[0]
    # y position of fist amss with offset
    y1 = l1*scale*np.cos(angle1) + offset[1]

    x2 = x1 + l2*scale*np.sin(angle2)  # x position of second mass
    y2 = y1 + l2*scale*np.cos(angle2)  # y position of second mass

    return (x1, y1), (x2, y2)  # return the mass coordinates as tuples


# render method that runs every frame
def render(point1, point2):
    # define x1, y1, x2, y2 from points
    x1, y1 = int(point1[0]), int(point1[1])
    x2, y2 = int(point2[0]), int(point2[1])

    if prev_point:
        xp, yp = prev_point[0], prev_point[1]  # stores the previous point
        # draws trace line from previous point to new point
        pygame.draw.line(trace, LT_BLUE, (xp, yp), (x2, y2), 3)
    # ensures the background stays white
    screen.fill(WHITE)
    screen.blit(trace, (0, 0))  # merge two screens
    # scales the size of the graphics
    scale = 10
    pygame.draw.line(screen, BLACK, offset, (x1, y1), 5)  # first rod
    pygame.draw.line(screen, BLACK, (x1, y1), (x2, y2), 5)  # second rod
    pygame.draw.circle(screen, BLACK, offset, 8)  # pivot point
    pygame.draw.circle(screen, RED, (x1, y1), int(m1*scale))  # first mass
    pygame.draw.circle(screen, BLUE, (x2, y2), int(m1*scale))  # second mass
    return (x2, y2)


while True:
    for event in pygame.event.get():
        if event.type == pygame.QUIT:
            sys.exit()

    point1, point2 = update(y[2], y[3])
    prev_point = render(point1, point2)
    time_string = f'Time: {round(t, 1)} seconds'
    text = myfont.render(time_string, False, (0, 0, 0))
    screen.blit(text, (10, 10))

    t += dt
    y = y + RK4_step(y, t, dt)
    clock.tick(60)
    pygame.display.update()
