#!/usr/bin/env python3
#coding: utf8
import os
from numpy import *
from tkinter import *
from random import *
from scipy.integrate import odeint

global RunAll
global RunIter
RunAll=True
RunIter=False

deg=180.0/pi

def StartStop():
  global RunIter
  RunIter=not RunIter
  if RunIter:
    StartButton["text"]="Stop"
  else:
    StartButton["text"]="Restart"
    
def StopAll():
  global RunAll
  RunAll=False

root=Tk()
root.title("Double Pendulum")
cw=800
ch=800
canvas=Canvas(root,width=cw,height=ch,background="#ffffff")
canvas.grid(row=0,column=0)

toolbar=Frame(root)
toolbar.grid(row=0,column=1, sticky=N)
# ................................................................... buttons
StartButton=Button(toolbar,text="Start",command=StartStop,width=5)
StartButton.grid(row=0,column=0,sticky=W)

CloseButton=Button(toolbar, text="Close", command=StopAll)
CloseButton.grid(row=0,column=1,sticky=E)
# .......................................................... Entries for mass 1
LabMass1=Label(toolbar,text= "Mass 1")
LabMass1.grid(row=1,column=0)
EntryMass1 = Entry(toolbar,bd =5,width=8)
EntryMass1.grid(row=1,column=1)

Lab1=Label(toolbar,text="Length 1")
Lab1.grid(row=2,column=0)
Entry1 = Entry(toolbar,bd =5,width=8)
Entry1.grid(row=2,column=1)

LabTheta1=Label(toolbar,text= "\u03D1 1",font=("Helvetica",12))
LabTheta1.grid(row=3,column=0)
EntryTheta1 = Entry(toolbar,bd =5,width=8)
EntryTheta1.grid(row=3,column=1)

LabOmega1=Label(toolbar,text= "\u03C9 1",font=("Helvetica",12))
LabOmega1.grid(row=4,column=0)
EntryOmega1 = Entry(toolbar,bd =5,width=8)
EntryOmega1.grid(row=4,column=1)

LabVisco1=Label(toolbar,text= "Friction 1")
LabVisco1.grid(row=5,column=0)
EntryVisco1 = Entry(toolbar,bd =5,width=8)
EntryVisco1.grid(row=5,column=1)
# .......................................................... Entries for mass 2
LabMass2=Label(toolbar,text= "Mass 2")
LabMass2.grid(row=6,column=0)
EntryMass2 = Entry(toolbar,bd =5,width=8)
EntryMass2.grid(row=6,column=1)

Lab2=Label(toolbar,text="Length 2")
Lab2.grid(row=7,column=0)
Entry2 = Entry(toolbar,bd =5,width=8)
Entry2.grid(row=7,column=1)

LabTheta2=Label(toolbar,text= "\u03D1 2",font=("Helvetica",12))
LabTheta2.grid(row=8,column=0)
EntryTheta2 = Entry(toolbar,bd =5,width=8)
EntryTheta2.grid(row=8,column=1)

LabOmega2=Label(toolbar,text= "\u03C9 2",font=("Helvetica",12))
LabOmega2.grid(row=9,column=0)
EntryOmega2 = Entry(toolbar,bd =5,width=8)
EntryOmega2.grid(row=9,column=1)

LabVisco2=Label(toolbar,text= "Friction 2")
LabVisco2.grid(row=10,column=0)
EntryVisco2 = Entry(toolbar,bd =5,width=8)
EntryVisco2.grid(row=10,column=1)
# ........................................................... frame persistence
cycle_period=20 #milliseconds
# ...................................................................... Origin
Ox=cw/2
Oy=ch/2
prad=3
# ..................................................................... gravity
G=9.8
# .................................................................. pendulum 1
L1=150.0
rad1=12
m1=10.0
visc1=0.0
pColor1="red"
# .................................................................. pendulum 2
L2=170.0
rad2=12
m2=10.0
visc2=0.0
pColor2="blue"
#.............................................................. initial values
theta1=pi/2.0
omega1=0.0
theta2=0.75*pi
omega2=0.0
y=[theta1,omega1,theta2,omega2]
# ..................................................................... degrees
# .................................................................. parameters
params=[L1,m1,visc1,L2,m2,visc2,G]
# .......................................................... initialize entries
EntryMass1.insert(0,str(m1))
Entry1.insert(0,str(L1))
EntryTheta1.insert(0,str(round(theta1*deg,1)))
EntryOmega1.insert(0,str(omega1))
EntryVisco1.insert(0,str(visc1))
EntryMass2.insert(0,str(m2))
Entry2.insert(0,str(round(L2,1)))
EntryTheta2.insert(0,str(round(theta2*deg,1)))
EntryOmega2.insert(0,str(omega2))
EntryVisco2.insert(0,str(visc2))
# .................................................................... function
def f(y, t, params):
  theta1,omega1,theta2,omega2=y        # unpack current values of y
  L1,m1,visc1,L2,m2,visc2,G=params     # unpack parameters
  mu=1.0+(m1/m2)
  dtheta=theta1-theta2
  csdtheta=cos(dtheta)
  omega1dot=(G*(sin(theta2)*csdtheta-mu*sin(theta1))\
    -(L2*(omega2**2)+L1*(omega1**2)*csdtheta)*sin(dtheta))\
      /(L1*(mu-(csdtheta**2)))-omega1*visc1
  omega2dot=(G*mu*(sin(theta1)*csdtheta-sin(theta2))\
    +(mu*L1*(omega1**2)+(L2*omega2**2)*csdtheta)*sin(dtheta))\
      /(L2*(mu-(csdtheta)**2))-omega2*visc2
  derivs = [omega1,omega1dot,omega2,omega2dot] # list of dy/dt=f functions
  return derivs

# ..................................................... numerical time interval
t=[0.0,0.1]
while RunAll:
  # ...................................................... draw double pendulum
  canvas.delete(ALL)
  canvas.create_line(0,Oy,cw,Oy,fill="green")
  canvas.create_oval(Ox-prad,Oy-prad,Ox+prad,Oy+prad,fill="black")
  x1=Ox+L1*sin(theta1)
  y1=Oy+L1*cos(theta1)
  x2=x1+L2*sin(theta2)
  y2=y1+L2*cos(theta2)
  canvas.create_line(Ox,Oy,x1,y1,fill="black")
  canvas.create_line(x1,y1,x2,y2,fill="black")
  canvas.create_oval(x1-rad1,y1-rad1,x1+rad1,y1+rad1,fill=pColor1)
  canvas.create_oval(x2-rad2,y2-rad2,x2+rad2,y2+rad2,fill=pColor2)
  canvas.update()
  canvas.after(cycle_period)
  if RunIter:
    # ............................................................... next step
    psoln = odeint(f,y,t,args=(params,))
    theta1=psoln[1,0]
    omega1=psoln[1,1]
    theta2=psoln[1,2]
    omega2=psoln[1,3]
    # ............................................................................
    y=[theta1,omega1,theta2,omega2]
  else:
    try:
      L1=float(Entry1.get())
    except ValueError:
      pass
    try:
      m1=float(EntryMass1.get())
    except ValueError:
      pass
    try:
      L2=float(Entry2.get())
    except ValueError:
      pass
    try:
      m2=float(EntryMass2.get())
    except ValueError:
      pass
    try:
      theta1=float(EntryTheta1.get())/deg
    except ValueError:
      pass
    try:
      omega1=float(EntryOmega1.get())
    except ValueError:
      pass
    try:
      visc1=float(EntryVisco1.get())
    except ValueError:
      pass
    try:
      theta2=float(EntryTheta2.get())/deg
    except ValueError:
      pass
    try:
      omega2=float(EntryOmega2.get())
    except ValueError:
      pass
    try:
      visc2=float(EntryVisco2.get())
    except ValueError:
      pass
    
    y=[theta1,omega1,theta2,omega2]
    params=[L1,m1,visc1,L2,m2,visc2,G]
  #----------------------------------------------------------------------------
root.destroy()
canvas.mainloop()