#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
Created on Sun Apr 18 20:18:45 2021

@author: olivier
"""

#Author: Chris Greening
#Date: 7/15/19
#Purpose: Another crack at the double pendulum to convert it to OOP
#to support multiple pendula

import numpy as np
import pandas as pd
from pandas import Series, DataFrame
from numpy import sin, cos
from scipy.integrate import odeint
import matplotlib.pyplot as plt
import matplotlib.animation as animation

from IPython.display import HTML # Added by OCornes for Jupyter  compatibility

fig = plt.figure()

class System:
    def __init__(self):
        self.double_pendulums = []

    def add_double_pendulum(self, double_pendulum):
        pass

class DoublePendulum:
    def __init__(
            self, 
            L1=1, 
            L2=1, 
            m1=1, 
            m2=1, 
            g=-9.81,
            y0=[90, 0, -10, 0],
            tmax = 180,
            dt = .05,
            color = "g"
        ):

        self.tmax = tmax
        self.dt = dt
        self.t = np.arange(0, self.tmax+self.dt, self.dt)
        self.color = color
        self.g = g
        self.pendulum1 = Pendulum(L1, m1)
        self.pendulum2 = Pendulum(L2, m2)
        
        # Initial conditions: theta1, dtheta1/dt, theta2, dtheta2/dt.
        self.y0 = np.array(np.radians(y0))

        # Do the numerical integration of the equations of motion
        self.y = odeint(self.derivative, self.y0, self.t, args=(self.pendulum1.L, self.pendulum2.L, self.pendulum1.m, self.pendulum2.m))

        self.pendulum1.calculate_path(self.y[:, 0])
        self.pendulum2.calculate_path(self.y[:, 2], self.pendulum1.x, self.pendulum1.y)

        self.w = self.y[:, 1]

        self.fig = fig
        self.ax_range = self.pendulum1.L + self.pendulum2.L
        self.ax = self.fig.add_subplot(111, autoscale_on=False, xlim=(-self.ax_range, self.ax_range), ylim=(-self.ax_range, self.ax_range))
        self.ax.set_aspect('equal')
        self.ax.grid()

        self.pendulum1.set_axes(self.ax)
        self.pendulum2.set_axes(self.ax)

        self.line, = self.ax.plot([], [], 'o-', lw=2,color=self.color)
        self.time_template = 'time = %.1fs'
        self.time_text = self.ax.text(0.05, 0.9, '', transform=self.ax.transAxes)

    def derivative(self, y, t, L1, L2, m1, m2):
        """Return the first derivatives of y = theta1, z1, theta2, z2."""
        theta1, z1, theta2, z2 = y

        c, s = np.cos(theta1-theta2), np.sin(theta1-theta2)

        theta1dot = z1
        z1dot = (m2*self.g*np.sin(theta2)*c - m2*s*(L1*z1**2*c + L2*z2**2) -
                (m1+m2)*self.g*np.sin(theta1)) / L1 / (m1 + m2*s**2)
        theta2dot = z2
        z2dot = ((m1+m2)*(L1*z1**2*s - self.g*np.sin(theta2) + self.g*np.sin(theta1)*c) + 
                m2*L2*z2**2*s*c) / L2 / (m1 + m2*s**2)
        return theta1dot, z1dot, theta2dot, z2dot

    def init(self):
        self.line.set_data([], [])
        self.time_text.set_text('')
        return self.line, self.time_text

class Pendulum:
    def __init__(self, L, m):
        self.L = L
        self.m = m

    def set_axes(self, ax):
        self.ax = ax

        # defines line that tracks where pendulum's have gone
        self.p, = self.ax.plot([], [], color='r-')
        self.w = self.ax.plot([], [])

    def calculate_path(self, y, x0=0, y0=0):
        self.x = self.L*np.sin(y) + x0
        self.y = self.L*np.cos(y) + y0

def animate(i):
    arr = pendula #pendulum2, pendulum3, pendulum4, pendulum5, pendulum6]
    return_arr = []
    for double_pendulum in arr:
        thisx = [0, double_pendulum.pendulum1.x[i], double_pendulum.pendulum2.x[i]]
        thisy = [0, double_pendulum.pendulum1.y[i], double_pendulum.pendulum2.y[i]]
        
        double_pendulum.line.set_data(thisx, thisy)
        double_pendulum.time_text.set_text(double_pendulum.time_template % (i*double_pendulum.dt))

        return_arr.append(double_pendulum.line)
        return_arr.append(double_pendulum.time_text)
        return_arr.append(double_pendulum.pendulum1.p)
        return_arr.append(double_pendulum.pendulum2.p)

    return return_arr

def random_hex():
    
    hex_chars = "0123456789ABCDEF"
    
    hex_string = "#"
    for i in range(6):
        index = np.random.randint(0, len(hex_chars))
        hex_string += hex_chars[index]

    return hex_string
        
L1 = 5
L2 = 5

pendula = []
initial_dtheta = 0
initial_theta = 90
dtheta = .5

#creates pendula 
for _ in range(10):
    pendula.append(DoublePendulum(L1=L1,L2=L2,y0=[initial_theta-initial_dtheta, 0,-10,0], tmax=15, color=random_hex()))
    initial_dtheta += dtheta

# plt.plot(pendula[0].x2, pendula[0].y2, color=pendula[0].color)

ani = animation.FuncAnimation(fig, animate, np.arange(1, len(pendula[0].y)),
                            interval=25, blit=True, init_func=pendula[0].init)

# ani.save('line.gif', dpi=80, writer='imagemagick')

HTML(ani.to_html5_video())

# plt.show() # removed by OCornes for Jupyter adaptation