#!/usr/bin/env python3
import numpy as np
from tkinter import *
from tkinter.messagebox import *
from random import *
from scipy.integrate import odeint
import time

RunAll=True
GetData=RunIter=ReWrite=False

ButtWidth=9

def StartStop():
  global RunIter
  RunIter=not RunIter
  if RunIter:
    StartButton["text"]="Stop"
  else:
    StartButton["text"]="Restart"
    
def ReadData(*arg):
  global GetData
  GetData=True
    
def StopAll():
  global RunAll
  RunAll=False
# .............................................................................
root=Tk()
root.title("Elastic Oscillator")
root.bind('<Return>',ReadData)
# ...................................................................... canvas
cw=800
ch=640
canvas=Canvas(root,width=cw,height=ch,background="#ffffff")
canvas.grid(row=0,column=0)
# ..................................................................... toolbar
toolbar=Frame(root)
toolbar.grid(row=0,column=1,sticky=N)
# ..................................................................... buttons
nr=1
StartButton=Button(toolbar,text="Start",command=StartStop,width=ButtWidth)
StartButton.grid(row=nr,column=0,sticky=W)
nr+=1
ReadButton=Button(toolbar, text="Read Data", command=ReadData,width=ButtWidth)
ReadButton.grid(row=nr,column=0,columnspan=2,sticky=W)
nr+=1
CloseButton=Button(toolbar, text="Exit", command=StopAll,width=ButtWidth)
CloseButton.grid(row=nr,column=0,sticky=W)
nr+=1
# ...................................................................... arrays
LabVar=[]
EntryVar=[]
VarList=["x\u2080","vx\u2080","y\u2080","vy\u2080"]
nVar=len(VarList)
LabPar=[]
EntryPar=[]
ParList=["Hooke","Mass","cycle/ms"]
nPar=len(ParList)
# ..................................................................... Entries
for i in range(nVar):
  LabVar.append(Label(toolbar,text=str(VarList[i]),font=("Helvetica",12)))
  LabVar[i].grid(row=nr,column=0)
  EntryVar.append(Entry(toolbar,bd =5,width=ButtWidth))
  EntryVar[i].grid(row=nr,column=1)
  nr+=1
for i in range(nPar):
  LabPar.append(Label(toolbar,text=str(ParList[i]),font=("Helvetica",12)))
  LabPar[i].grid(row=nr,column=0)
  EntryPar.append(Entry(toolbar,bd =5,width=ButtWidth))
  EntryPar[i].grid(row=nr,column=1)
  nr+=1
# ........................................................ time label
CycleLab0=Label(toolbar,text="Period:",font=("Helvetica",11))
CycleLab0.grid(row=nr,column=0)
CycleLab=Label(toolbar,text="     ",font=("Helvetica",11))
CycleLab.grid(row=nr,column=1,sticky=W)
nr+=1
# .............................................................................
cycle=20 #milliseconds
# .................................................................. Parameters
Ox=cw/2
Oy=ch/2
prad=3
rad=12
m=10.0
Hooke=10.0
pColor="red"
params=[Hooke,m,cycle]
# ................................................ initial values for variables
xx=300.0
vx=0.0
yy=0
vy=200.0
y=[xx,vx,yy,vy]
# .......................................................... initialize entries
for i in range(len(VarList)):
  EntryVar[i].insert(0,"{:.2f}".format(y[i]))
for i in range(len(ParList)):
  EntryPar[i].insert(0,"{:.2f}".format(params[i]))
# .................................................................... function
def f(y, t, params):
  xx,vx,yy,vy = y                  # unpack current values of y
  Hooke,m,cycle = params           # unpack parameters
  stretch=np.sqrt(xx**2+yy**2)
  theta=np.arctan2(yy,xx)
  force=-Hooke*stretch
  fx=force*np.cos(theta)
  fy=force*np.sin(theta)
  ax=(fx/m)
  ay=(fy/m)
  derivs=[vx,ax,vy,ay]
  return derivs
# .............................................................................
t=[0.0,0.05]
tt0=time.time()
tcount=0
# ........................................................................ loop
while RunAll:
  StartIter=time.time()
  # ............................................................. draw pendulum
  canvas.delete(ALL)
  canvas.create_line(0,ch-Oy,cw,ch-Oy,fill="green")
  canvas.create_oval(Ox-prad,ch-(Oy+prad),Ox+prad,ch-(Oy-prad),fill="black")
  canvas.create_line(Ox,ch-Oy,Ox+xx,Oy-yy,fill="black")
  canvas.create_oval(Ox+xx-rad,Oy-(yy+rad),Ox+xx+rad,Oy-(yy-rad),fill=pColor)
  if RunIter:
    # ............................................................... next step
    psoln = odeint(f,y,t,args=(params,))
    xx=psoln[1,0]
    vx=psoln[1,1]
    yy=psoln[1,2]
    vy=psoln[1,3]
    # .........................................................................
    y=[xx,vx,yy,vy]
  elif GetData:
    i=0
    while i<nVar:
      try:
        y[i]=float(EntryVar[i].get())
      except ValueError:
        pass
      i+=1
    i=0
    while i<nPar:
      try:
        params[i]=float(EntryPar[i].get())
      except ValueError:
        pass
      i+=1
      ReWrite=True
      GetData=False
  if ReWrite:
    i=0
    while i<nVar:
      EntryVar[i].delete(0,'end')
      EntryVar[i].insert(0,"{:.2f}".format(y[i]))
      i+=1
    i=0
    while i<nPar:
      EntryPar[i].delete(0,'end')
      EntryPar[i].insert(0,"{:.2f}".format(params[i]))
      i+=1
    ReWrite=False
    xx,vx,yy,vy=y
    Hooke,m,cycle=params
    cycle=int(cycle)
    GetData=False
  # .................................................. cycle duration
  tcount+=1
  if tcount==10:
    tcount=0
    ttt=time.time()
    elapsed=ttt-tt0
    CycleLab['text']="%8.2f"%(elapsed*100.0)+" ms"
    tt0=ttt
  ElapsIter=int((time.time()-StartIter)*1000.0)
  canvas.update()
  canvas.after(cycle-ElapsIter)
  #----------------------------------------------------------------------------
root.destroy()
root.mainloop()