#!/usr/bin/env python3
from tkinter import *
from datetime import timedelta
import numpy as np
import time
# .................................................. Global variables
RunAll=True
GetStep=GetTau=RunMotion=False
# ................................................... Physical values
ME=5.9722e24          # Earth mass/kg
G=6.674e-11           # Gravitational constant [m^3/(kg s^2)]
GM=ME*G
dt=5.0                # s
# .................................. Drawing and Animation Parameters
tau=5                 # ms
scale=5.0e-6          # px/m
cw=600                # px
ch=500                # px
Ox=150                # px
Oy=ch/2.0             # px
rad=4                 # px
ms=1.0e3              # satellite mass / kg
TrailLength=400
x0=4.0e7              # m
y0=0.0
vx0=0.0
vy0=0.5*np.sqrt(GM/x0) # m/s
col='red'
# ...................................................................
INITENER,EULERENER,ITER,ELAPSED,ORBITS,PERIOD,CYCLE,SCALE=range(8)
quant=['Initial Energy','EulerEnergy','Iterations',\
  'Elapsed Time','Orbits','Orbital Period','Cycle','Scale']
# ............................................... Start/Stop function
def StartStop():
  global RunMotion
  RunMotion=not RunMotion
  if RunMotion:
    StartButton['text']='Stop'
    CloseButton['state']=DISABLED
  else:
    StartButton['text']='Restart'
    CloseButton['state']=NORMAL
# ..................................................... Exit function
def StopAll():
  global RunAll
  RunAll=False
# ........................................................ Scale Down
def ScaleUpDown(ud):
  global scale
  scale*=np.sqrt(2)**ud
  QtLab[SCALE]['text']='{:10.3e}'.format(scale)
  ScaledTrail[::2]=[Ox+scale*x for x in Trail[::2]]
  ScaledTrail[1::2]=[Oy-scale*y for y in Trail[1::2]]
# ........................................................ Read Entry
def ReadData(tx):
  global GetStep,GetTau
  if tx==0:
    GetStep=True
  elif tx==1:
    GetTau=True
# ...................................................... Acceleration  
def accel(x,y):
  r2=x**2+y**2
  aa=-GM/r2
  alpha=np.arctan2(y,x)
  ax=aa*np.cos(alpha)
  ay=aa*np.sin(alpha)
  return [ax,ay]
# ............................................................ Energy
def ener(x,y,vx,vy):
  r=np.sqrt(x**2+y**2)
  pot=-GM*ms/r
  kin=0.5*ms*(vx**2+vy**2)
  return pot+kin
# ..................................................... Time Reversal
def TimeReversal():
  global dt
  dt=-dt
# ....................................................... Root Window
root=Tk()
root.title('Euler Satellite')
root.bind('<Control-plus>',lambda event,num=1:ScaleUpDown(num))
root.bind('<Control-minus>',lambda event,num=-1:ScaleUpDown(num))
# ............................................................ Canvas
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)
toolbar.columnconfigure(1,minsize=120)
# ............................................................ Buttons
nr=0
StartButton=Button(toolbar,text='Start',command=StartStop,width=11)
StartButton.grid(row=nr,column=0,sticky=W)
nr+=1
TimeRevButton=Button(toolbar,text='Time Reversal',\
  command=TimeReversal,width=11)
TimeRevButton.grid(row=nr,column=0,sticky=W)
nr+=1
CloseButton=Button(toolbar, text='Exit', command=StopAll,width=11)
CloseButton.grid(row=nr,column=0,columnspan=2,sticky=W)
nr+=1
# ...................................................... Time Entries
TimEntry=[]
TimeTxt=['Time Step/s','\u03C4/ms']
tVal=[dt,tau]
tfor=['{:.2f}','{:d}']
for i,tt in enumerate(TimeTxt):
  lb=Label(toolbar,text=tt,font=('Helvetica',11))
  lb.grid(row=nr,column=0)
  TimEntry.append(Entry(toolbar,bd=5,width=11))
  TimEntry[i].grid(row=nr,column=1)
  TimEntry[i].insert(0,tfor[i].format(tVal[i]))
  TimEntry[i].bind('<Return>',lambda event,num=i:ReadData(num))
  nr+=1
# ............................................................ Labels
QtLab=[]
for i,qt in enumerate(quant):
  lab=Label(toolbar,text=qt,font=('Helvetica',11))
  lab.grid(row=nr,column=0)
  QtLab.append(Label(toolbar,text='0',font=('Helvetica',11)))
  QtLab[i].grid(row=nr,column=1)
  nr+=1
QtLab[SCALE]['text']='{:10.3e}'.format(scale)
nr+=1
# .............................................. Draw Coordinate Axes
canvas.create_line(0,Oy,cw,Oy,fill='black')
canvas.create_line(Ox,0,Ox,ch,fill='black')
canvas.create_oval(Ox-6,Oy-8,Ox+8,Oy+8,outline='#50a0ff',fill='#50a0ff')
# .................................................... Initial Values
x,y,vx,vy=x0,y0,vx0,vy0
# ................................................... Euler Satellite
SatImg=canvas.create_oval(Ox+scale*x-rad,Oy-scale*y+rad,\
  Ox+scale*x+rad,Oy-scale*y-rad,fill=col,outline=col)
Trail=[x,y]*TrailLength
ScaledTrail=[Ox+scale*x,Oy-scale*y]*TrailLength
ImTrail=canvas.create_line(ScaledTrail,fill=col)
# .................................................... Initial Energy
en=ener(x0,y0,vx0,vy0)
QtLab[INITENER].config(text='{:.6e}'.format(en))
QtLab[EULERENER].config(text='{:.6e}'.format(en))
# ...................................................................
nIter=0
tcount=0
nOrbits=0
Telaps=0.0
ypair=[y0,y0]
tt0=time.time()
# ......................................................... Main Loop
while RunAll:
  StartIter=time.time()
  # ............................................ Draw Euler Satellite
  canvas.coords(SatImg,Ox+scale*x-rad,Oy-scale*y+rad,\
    Ox+scale*x+rad,Oy-scale*y-rad)
  canvas.coords(ImTrail,ScaledTrail)
  canvas.update()
  # .......................................................... Motion
  if RunMotion:
    Telaps+=dt
    # ............................................... Euler Algorithm
    ax,ay=accel(x,y)
    x+=vx*dt
    y+=vy*dt
    vx+=ax*dt
    vy+=ay*dt
     # ................................................. Update Trail
    if scale*np.linalg.norm([x-Trail[-2],y-Trail[-1]])>10:
      del Trail[:2]
      Trail.extend([x,y])
      del ScaledTrail[:2]
      ScaledTrail.extend([Ox+scale*x,Oy-scale*y])
    # ................................................. Update ypair
    ypair[0]=ypair[1]
    ypair[1]=y
    if ypair[0]<0 and ypair[1]>0:
      nOrbits+=1
      QtLab[ORBITS].config(text=str(nOrbits))
      OrbPeriod=Telaps/nOrbits
      QtLab[PERIOD].config(text=str(timedelta(seconds=int(OrbPeriod))))
    # ........................................ Show Iteration Counter
    nIter+=1
    if nIter%20==0:
      QtLab[ELAPSED].config(text=str(timedelta(seconds=int(Telaps))))
      QtLab[ITER].config(text=str(nIter))
      en=ener(x,y,vx,vy)
      QtLab[EULERENER].config(text='{:.6e}'.format(en))
  # ................................................... New Time Step
  elif GetStep:
    try:
      dt=float(TimEntry[0].get())
    except ValueError:
      pass
    TimEntry[0].delete(0,END)
    TimEntry[0].insert(0,'{:.2f}'.format(dt))
    # ...................................... Reset Position and Trail
    x,vx,y,vy=x0,vx0,y0,vy0
    Trail=[x0,y0]*TrailLength
    ScaledTrail=[Ox+scale*x0,Oy-scale*y0]*TrailLength
    ypair=[y0,y0]
    nOrbits=0
    Telaps=0.0
    GetStep=FALSE
  elif GetTau:
    try:
      tau=int(TimEntry[1].get())
    except ValueError:
      pass
    TimEntry[1].delete(0,END)
    TimEntry[1].insert(0,'{:d}'.format(tau))
    GetTau=FALSE
  # ................................................ Cycle Duration
  tcount+=1
  if tcount>=10:
    tcount=0
    ttt=time.time()
    elapsed=ttt-tt0
    QtLab[CYCLE]['text']='%8.3f'%(elapsed*100.0)+' ms'
    tt0=ttt
  ElapsIter=int((time.time()-StartIter)*1000.0)
  canvas.after(tau-ElapsIter)
#----------------------------------------------------------------------
root.destroy()