#!/usr/bin/env python3
from tkinter import *
import numpy as np
from scipy.integrate import odeint
import time
# .................................................. Global variables
RunAll=True
RunIter=GetData=False
ButtWidth=11
rad=10
x=300.0
y=0.0
vx=0.0
vy=2.83    # circular
dt=2.0     # s
tau=20     # ms
nTrail=800
params=[rad,x,y,vx,vy,dt,tau,nTrail]
nPar=len(params)
K=8.0      # force constant
# .................................................. Button functions
def StartStop():
  global RunIter
  RunIter=not RunIter
  if RunIter:
    StartButton['text']='Stop'
  else:
    StartButton['text']='Restart'
# .............................................
def StopAll():
  global RunAll
  RunAll=False
# ..............................................
def ReadData(*arg):
  global GetData
  GetData=True 
#................................................. special velocities
def SpecialVel(vvv):
  global x,y,vx,vy,yy,params,EntryPar,TrailCoords,nPar,nTrail
  x=params[1]
  y=params[2]
  vx=params[3]=0
  vy=params[4]=vvv
  yy=[x,vx,y,vy]
  for i in range(nPar):
    EntryPar[i].delete(0,'end')
    EntryPar[i].insert(0,'{:.4f}'.format(params[i]))
  TrailCoords=[Ox+x,Oy-y]*nTrail
# ...................................................................
root=Tk()
root.title('Logarithmic Potential')
root.bind('<Return>',ReadData)
# ....................................................... open canvas
cw=ch=800
Ox=Oy=cw/2
canvas=Canvas(root,width=cw,height=ch,background='white')
canvas.grid(row=0,column=0)
# ........................................................... toolbar
toolbar=Frame(root)
toolbar.grid(row=0,column=1,sticky=N)
# ........................................................... buttons
nr=0
StartButton=Button(toolbar,text='Start',command=StartStop,\
  width=ButtWidth)
StartButton.grid(row=nr,column=0,sticky=W)
nr+=1
# ................................................ special velocities
ButtonCirc=Button(toolbar,text='Circular',\
  command=lambda:SpecialVel(2.83),width=ButtWidth)
ButtonCirc.grid(row=nr,column=0,sticky=W)
nr+=1
ButtonLobe3=Button(toolbar,text='3 lobes',\
  command=lambda:SpecialVel(0.893),width=ButtWidth)
ButtonLobe3.grid(row=nr,column=0,sticky=W)
nr+=1
ButtonLobe5=Button(toolbar,text='5 lobes',\
  command=lambda:SpecialVel(0.18416),width=ButtWidth)
ButtonLobe5.grid(row=nr,column=0,sticky=W)
nr+=1
ButtonLobe7=Button(toolbar,text='7 lobes',\
  command=lambda:SpecialVel(6.149e-2),width=ButtWidth)
ButtonLobe7.grid(row=nr,column=0,sticky=W)
nr+=1
ButtonLobe9=Button(toolbar,text='9 lobes',\
  command=lambda:SpecialVel(2.358e-2),width=ButtWidth)
ButtonLobe9.grid(row=nr,column=0,sticky=W)
nr+=1
# ....................................................... Exit button
CloseButton=Button(toolbar,text='Exit',command=StopAll,\
  width=ButtWidth)
CloseButton.grid(row=nr,column=0,sticky=W)
nr+=1
# ........................................................... Entries
EntryPar=[]
LabPar=[]
ParList=['rad','x','y','vx','vy','dt','Cycle/ms','Trail length']
nPar=len(ParList)
# ...................................................................
for i in range (nPar):
  LabPar.append(Label(toolbar,text=str(ParList[i]),
                      font=('Times',11)))
  LabPar[i].grid(row=nr,column=0)
  EntryPar.append(Entry(toolbar,bd=5,width=10))
  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
# ........................................................ Parameters
for i in range(nPar):
  EntryPar[i].delete(0,'end')
  EntryPar[i].insert(0,'{:.4f}'.format(params[i]))
# .............................................. create axes and ball
ball=canvas.create_oval(Ox+x-rad,Oy-y-rad,Ox+x+rad,Oy+y+rad,\
  fill='red',outline='red')
baryc=canvas.create_oval(Ox-2,Oy-2,Ox+2,Oy+2,fill='black',\
  outline='black')
canvas.create_line(0,Oy,cw,Oy,fill='black')
canvas.create_line(Ox,0,Ox,ch,fill='black')
TrailCoords=[Ox+x,Oy-y]*nTrail
trail=canvas.create_line(TrailCoords,fill='red')
# ...........................................................
yy=[x,vx,y,vy]
t=[0,dt]
# ........................................... function for odeint()
def dfdt(yInp,t):
  x,vx,y,vy=yInp
  r=np.sqrt(x**2+y**2)
  alpha=np.arctan2(y,x)
  f=-K/r
  ax=f*np.cos(alpha)
  ay=f*np.sin(alpha)
  return [vx,ax,vy,ay]
# ................................................... initialize time
tt0=time.time()
tcount=0
# ......................................................... main loop
while RunAll:
  StartIter=time.time()
  canvas.coords(ball,Ox+x-rad,Oy-y-rad,Ox+x+rad,Oy-y+rad)
  canvas.coords(trail,TrailCoords)
  canvas.update()
  if RunIter:
    # ..................................................... move ball
    psoln=odeint(dfdt,yy,t)
    yy=psoln[1,:]
    x,vx,y,vy=yy
    # .................................................. update trail
    xtr=Ox+x
    ytr=Oy-y
    if np.fabs(xtr-TrailCoords[-2])>3 or\
      np.fabs(ytr-TrailCoords[-1])>3:
        del TrailCoords[:2]
        TrailCoords.append(xtr)
        TrailCoords.append(ytr)
  elif GetData:
    i=0
    while i<nPar:
      try:
        params[i]=float(EntryPar[i].get())
      except ValueError:
        pass
      i+=1
    rad,x,y,vx,vy,dt,tau,nTrail=params
    tau=int(tau)
    nTrail=int(nTrail)
    TrailCoords=[Ox+x,Oy-y]*nTrail
    yy=[x,vx,y,vy]
    t=[0,dt]
    for i in range(nPar):
      EntryPar[i].delete(0,'end')
      EntryPar[i].insert(0,'{:.4f}'.format(params[i]))
    GetData=False
  # .................................................. cycle duration
  tcount+=1
  if tcount==10:
    tcount=0
    ttt=time.time()
    elapsed=ttt-tt0
    CycleLab['text']='{:8.1f}'.format(elapsed*100.0)+' ms'
    tt0=ttt
  # ........................................... wait for cycle duration
  ElapsIter=int((time.time()-StartIter)*1000.0)
  canvas.after(tau-ElapsIter)