#!/usr/bin/env python3
from tkinter import *
from scipy.integrate import odeint
from numpy.linalg import norm
from datetime import timedelta
import numpy as np
import time
# ............................................................. Lists
QtKey=['name','rad','m','col','x','y','vx','vy']
QtLab=['Name','Radius','Mass','Color','x','y','vx','vy']
nQt=len(QtKey)
PrLab=['dt','\u03C4/ms','Trail length']
PrForm=['{:.3e}','{:d}','{:d}']
# .................................................. Global Variables
RunAll=True
GetPr=GetQt=RunMotion=False
# ................................................... Physical values
G=6.67408e-11
tHour=3600
tDay=86400
#tYear=3.1557600e7   # Julian
#tYear=3.1556952e7   # Gregorian
tYear=3.1558149504e7 # Sidereal
dt=2.16e4            # 0.25 day
# .................................. Drawing and Animation Parameters
tau=20                # ms
scale=1.54e-9         # px/m
cw=ch=900             # px
Ox=cw/2.0
Oy=ch/2.0
bcrad=2               # px
TrailLength=400
SelB=SelPr=SelQt=0
# ......................................................... Parameters
param=[dt,tau,TrailLength]
# ...................................................... Class CelBody
class CelBody:
  def __init__(self,name,radius,mass,color,x,y,vx,vy):
    self.name=name
    self.rad=radius
    self.m=mass
    self.col=color
    self.x=x
    self.y=y
    self.vx=vx
    self.vy=vy
    self.x0,self.y0,self.vx0,self.vy0=self.x,self.y,self.vx,self.vy
    self.image=canvas.create_oval(cvx(self.x)-self.rad,\
      cvy(self.y)+self.rad,cvx(self.x)+self.rad,\
        cvy(self.y)-self.rad,fill=self.col,outline=self.col)
    self.trail=[self.x,self.y]*TrailLength
    self.ScaledTrail=[cvx(self.x),cvy(self.y)]*TrailLength
    self.TrailImg=canvas.create_line(self.ScaledTrail,fill=self.col)
  # ....................................................... Move Body
  def redraw(self):
    canvas.coords(self.image,cvx(self.x)-self.rad,\
      cvy(self.y)+self.rad,cvx(self.x)+self.rad,\
        cvy(self.y)-self.rad)
    if norm([self.x-self.trail[-2],self.y-self.trail[-1]])*scale>10:
      del self.trail[:2]
      self.trail.extend([self.x,self.y])
      del self.ScaledTrail[:2]
      self.ScaledTrail.extend([cvx(self.x),cvy(self.y)])
    canvas.coords(self.TrailImg,self.ScaledTrail)
# ............................................................... cvx
def cvx(x):
  global Ox,scale
  return Ox+scale*x
# ............................................................... cvy
def cvy(x):
  global Oy,scale
  return Oy-scale*x
# ............................................... Start/Stop Function
def StartStop():
  global RunMotion
  RunMotion=not RunMotion
  if RunMotion:
    butt[0]['text']='Stop'
    for b in butt[1:]+QtEntry+PrEntry:
      b['state']=DISABLED
  else:
    butt[0]['text']='Restart'
    for b in butt[1:]+QtEntry+PrEntry:
      b['state']=NORMAL
    SelectBody(0)
# ..................................................... Exit Function
def StopAll():
  global RunAll
  RunAll=False
# ........................................... Read Particol Variables
def ReadQt(WhichEntry):
  global GetQt,SelQt
  SelQt=WhichEntry
  GetQt=True  
# ................................................... Read Parameters
def ReadPr(WhichEntry):
  global GetPr,SelPr
  SelPr=WhichEntry
  GetPr=True
# ....................................................... Select Body
def SelectBody(delta):
  global SelB
  SelB=(SelB+delta)%nB
  SelLab.config(text=body[SelB].name,fg=body[SelB].col)
  for i,vv in enumerate(list(body[SelB].__dict__.values())[4:nQt]):
    QtEntry[i].delete(0,'end')
    QtEntry[i].insert(0,'{:.3e}'.format(vv))
# ........................................................ Scale Down
def ScaleUpDown(ud):
  global scale
  scale*=np.sqrt(2)**ud
  ScaleLab['text']='{:10.3e}'.format(scale)
  for b in body:
    b.ScaledTrail[::2]=[cvx(x) for x in b.trail[::2]]
    b.ScaledTrail[1::2]=[cvy(y) for y in b.trail[1::2]]
# .......................... Evaluate Center of Mass and its Velocity
def baryc(body):
  mtot=sum(b.m for b in body)
  cx=sum(b.x*b.m for b in body)/mtot
  cy=sum(b.y*b.m for b in body)/mtot
  cvx=sum(b.vx*b.m for b in body)/mtot
  cvy=sum(b.vy*b.m for b in body)/mtot
  return [cx,cy,cvx,cvy]
# ..................................... Move Origin to Center of Mass
def SetBaryc():
  global posvect,velvect
  xcm,ycm,cvx,cvy=baryc(body)
  for i,b in enumerate(body):
    b.x-=xcm
    b.y-=ycm
    b.vx-=cvx
    b.vy-=cvy
    posvect[2*i:2*i+2]=[b.x,b.y]
    velvect[2*i:2*i+2]=[b.vx,b.vy]
  SelectBody(0)
# ....................................................... Acceleration
def accel(posvect):
  global g,masses
  x=posvect[::2]
  y=posvect[1::2]
  distx=np.atleast_2d(x).T-x
  disty=np.atleast_2d(y).T-y
  alpha=np.arctan2(disty,distx)
  r2=np.square(distx)+np.square(disty)
  np.fill_diagonal(r2,1.0)
  ff=np.divide(G,r2)
  np.fill_diagonal(ff,0.0)
  fx=ff*np.cos(alpha)
  fy=ff*np.sin(alpha)
  fx=(fx.T*masses).T
  fy=(fy.T*masses).T
  accvect=[0]*len(posvect)
  accvect[::2]=fx.sum(axis=0)
  accvect[1::2]=fy.sum(axis=0)
  return np.array(accvect)
# .............................................................. dfdt
def dfdt(InpVect,t):
  pos=InpVect[:2*nB]
  vel=InpVect[2*nB:]
  acc=accel(pos)
  return np.append(vel,acc)
# ...................................................... Reinitialize
def reinit():
  global body,CurrTime,nIter,posvect,TrailLength,velvect
  for i,b in enumerate(body):
    b.x=b.x0
    b.y=b.y0
    b.vx=b.vx0
    b.vy=b.vy0
    posvect[2*i:2*i+2]=[b.x,b.y]
    velvect[2*i:2*i+2]=[b.vx,b.vy]
    b.trail=[b.x,b.y]*TrailLength
    b.ScaledTrail=[cvx(b.x),cvy(b.y)]*TrailLength
  nIter=0
  CurrTime=0.0
  TimeLab.config(text=str(CurrTime/tDay)+' d')
  IterLab.config(text=str(nIter))
  SelectBody(0)
# ....................................................... Root Window
root=Tk()
root.title('Planets. Start: 1 January 2019')
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.option_add('*Font','Helvetica 11')
# ........................................................... Buttons
nr=0
butt=[]
ButtLab=['Start','Set Barycenter','Reinitialize','Exit']
ButtComm=[StartStop,SetBaryc,reinit,StopAll]
for i,(ll,cc) in enumerate(zip(ButtLab,ButtComm)):
  butt.append(Button(toolbar,text=ll,command=cc,width=11))
  butt[i].grid(row=nr,column=0,sticky=W)
  nr+=1
# ..................................................... Create Bodies
body=[]
body.append(CelBody('Sun',10.0,1.9891e30,'#ffff00',0.0,0.0,0.0,0.0))
body.append(CelBody('Mercury',4.0,3.30e23,'#888888',-3.64259e10,\
  -3.7153e10,3.06117e4,-4.32992e4,))
body.append(CelBody('Venus',4.0,4.87e24,'#dfdf00',-8.2943e10,\
  6.89927e10,-2.22843e4,-2.71505e4))
body.append(CelBody('Earth',4.0,5.97e24,'#4444ff',-2.64516e10,\
  1.46764e11,-2.93198e4,-5.77993e3,))
body.append(CelBody('Mars',4.0,6.42e23,'#ff0000',1.60991e11,\
  1.44349e11,-1.83819e4,1.75399e4,))
body.append(CelBody('Ceres',2.0,9.39e20,'#000000',4.4541e11,\
  0.0,0.0,1.65967e4,))
body.append(CelBody('Jupiter',6.0,1.90e27,'#ffaa50',-3.19386e11,\
  -7.45836e11,1.13936e4,-5.20503e3,))
body.append(CelBody('Saturn',6.0,5.68e26,'#dfdf00',3.09446e11,\
  -1.42781e12,9.32053e3,1.51707e3,))
body.append(CelBody('Uranus',6.0,8.68e25,'#8888ff',2.5355e12,\
  1.57957e12,-3.59205e3,5.46011e3,))#  Uranus
body.append(CelBody('Neptune',6.0,1.02e26,'#8888ff',4.39613e12,\
  -1.09375e12,1.27283e3,5.24317e3))
body.append(CelBody('Pluto',2.0,1.46e22,'#000000',7.37593e12,\
  0.0,0.0,3.67702e3,))
nB=len(body)
# ....................................... Position and Velocity Lists
posvect=np.array([0.0,0.0]*nB)
velvect=np.array([0.0,0.0]*nB)
masses=np.array([0.0]*nB)
for i,b in enumerate(body):
  posvect[2*i:2*i+2]=[b.x,b.y]
  velvect[2*i:2*i+2]=[b.vx,b.vy]
  masses[i]=b.m
# ........................................... Selected Particle Label
Label(toolbar,text='Selected Body:',pady=20).grid(row=nr,column=0)
SelLab=Label(toolbar,text=body[0].name,width=15,bg='#ffffff')
SelLab.grid(row=nr,column=1)
SelLab.bind('<Button-5>',lambda event,num=-1:SelectBody(num))
SelLab.bind('<Button-1>',lambda event,num=-1:SelectBody(num))
SelLab.bind('<Button-4>',lambda event,num=1:SelectBody(num))
SelLab.bind('<Button-3>',lambda event,num=1:SelectBody(num))
nr+=1
# .................. Entries for Physical Quantities of Selected Body
QtEntry=[]
for i,kk in enumerate(QtKey[4:]):
  Label(toolbar,text=QtLab[i+4]).grid(row=nr,column=0)
  QtEntry.append(Entry(toolbar,bd=3,width=16))
  QtEntry[i].grid(row=nr,column=1)
  QtEntry[i].insert(0,'{:.3e}'.format(body[0].__dict__[str(kk)]))
  QtEntry[i].bind('<Return>',lambda event,num=i:ReadQt(num))
  nr+=1
# ......................................................... Separator
Label(toolbar,text='  ').grid(row=nr,column=0)
nr+=1
# ............................................ Entries for Parameters
PrEntry=[]
for i,pl in enumerate(PrLab):
  Label(toolbar,text=pl).grid(row=nr,column=0)
  PrEntry.append(Entry(toolbar,bd=3,width=16))
  PrEntry[i].grid(row=nr,column=1)
  PrEntry[i].insert(0,PrForm[i].format(param[i]))
  PrEntry[i].bind('<Return>',lambda event,num=i:ReadPr(num))
  nr+=1
# ....................................................... Cycle Label
Label(toolbar,text='Period:',).grid(row=nr,column=0)
CycleLab=Label(toolbar,text='     ')
CycleLab.grid(row=nr,column=1,sticky=W)
nr+=1
# ............................................. Elapsed Physical Time
Label(toolbar,text='Elapsed time:').grid(row=nr,column=0)
TimeLab=Label(toolbar,text='0     ')
TimeLab.grid(row=nr,column=1,sticky=W)
nr+=1
# ........................................................ Iterations
Label(toolbar,text='Iterations:').grid(row=nr,column=0)
IterLab=Label(toolbar,text='0     ')
IterLab.grid(row=nr,column=1,sticky=W)
nr+=1
# ....................................................... Scale Label
Label(toolbar,text='Scale:').grid(row=nr,column=0)
ScaleLab=Label(toolbar,text='{:10.3e}'.format(scale))
ScaleLab.grid(row=nr,column=1,sticky=W)
nr+=1
# .............................................. Draw Coordinate Axes
canvas.create_line(0,Oy,cw,Oy,fill='black')
canvas.create_line(Ox,0,Ox,ch,fill='black')
# ................................. Create Barycenter Image on Canvas
bc=canvas.create_oval(Ox-bcrad,Oy-bcrad,Ox+bcrad,Oy+bcrad,fill='black')
# ................................................... Initialize Time
tt0=time.time()
tcount=0
t=[0,dt]
nIter=0
CurrTime=0.0
# .................................................... Animation Loop
while RunAll:
  StartIter=time.time()
  # .................................................. Draw Particles
  for b in body:
    b.redraw()
  # ................................................. Draw Barycenter
  cx,cy=baryc(body)[:2]
  canvas.coords(bc,cvx(cx)-bcrad,cvy(cy)-bcrad,cvx(cx)+bcrad,\
    cvy(cy)+bcrad)
  canvas.update()
  # .......................................................... motion
  if RunMotion:
    nIter+=1
    CurrTime+=dt
    # ................................................... Call odeint
    sol=odeint(dfdt,np.append(posvect,velvect),t)
    posvect=sol[1,:][:2*nB]
    velvect=sol[1,:][2*nB:]
    for i,b in enumerate(body):
      b.x,b.y=posvect[2*i:2*i+2]
      b.vx,b.vy=velvect[2*i:2*i+2]
    if nIter%20==0:
      TimeLab.config(text=str(CurrTime/tDay)+' d')
      IterLab.config(text=str(nIter))
  elif GetPr: # ................................... Read Parameters
    try:
      vv=float(PrEntry[SelPr].get())
    except ValueError:
      pass
    else:
      if SelPr==1 or SelPr==2:
        vv=int(vv)
      param[SelPr]=vv
      PrEntry[SelPr].delete(0,'end')
      PrEntry[SelPr].insert(0,PrForm[SelPr].format(vv))
      # ...................................... If TrailLength Changed
      dTrail=param[2]-TrailLength
      if dTrail<0:
        for b in body:
          del b.trail[:2*abs(dTrail)]
          del b.ScaledTrail[:2*abs(dTrail)]
      elif dTrail>0:
        for b in body:
          NewPoints=[b.trail[0],b.trail[1]]*dTrail
          b.trail=NewPoints+b.trail
          NewPoints=[b.ScaledTrail[0],b.ScaledTrail[1]]*dTrail
          b.ScaledTrail=NewPoints+b.ScaledTrail
      # .............................................................
      dt,tau,TrailLength=param
      t=[0,dt]
    GetPr=False
  elif GetQt:  # ............................ Read Particle Variables
    try:
      body[SelB].__dict__[str(QtKey[SelQt+4])]=vv=\
        float(QtEntry[SelQt].get())
    except ValueError:
      pass
    else:
      QtEntry[SelQt].delete(0,'end')
      QtEntry[SelQt].insert(0,'{:.3e}'.format(vv))
      posvect[2*SelB:2*SelB+2]=[body[SelB].x,body[SelB].y]
      velvect[2*SelB:2*SelB+2]=[body[SelB].vx,body[SelB].vy]
      masses[SelB]=body[SelB].m
    GetQt=False
  # ................................................ Cycle Duration
  tcount+=1
  if tcount>=10:
    tcount=0
    ttt=time.time()
    elapsed=ttt-tt0
    CycleLab['text']='{:8.3f}'.format(elapsed*100.0)+' ms'
    tt0=ttt
  ElapsIter=int((time.time()-StartIter)*1000.0)
  canvas.after(tau-ElapsIter)
#--------------------------------------------------------------- Exit
root.destroy()