#!/usr/bin/env python3
import os
from math import radians
from numpy import *
from tkinter import *
from random import *
from scipy.integrate import odeint
import time

# ............................................. subscripts for labels
sub=['\u2080','\u2081','\u2082','\u2083','\u2084','\u2085',\
  '\u2086','\u2087','\u2088','\u2089','\u2081\u2080','\u2018\u2081 ']
# .................................................. Global variables
RunAll=True
RunIter=False
ButtWidth=20
NewBaryc=False
GetData=False
ReWrite=False
nRows=32
scale=1.54e-9
cw=900
ch=900
bcrad=2
cycle=20             # ms
dt=2.16e4            # 0.25 day
PathLength=200
G=6.67408e-11
tHour=3600
tDay=86400
#tYear=3.1557600e7   # Julian
#tYear=3.1556952e7   # Gregorian
tYear=3.1558149504e7 # Sidereal
# ............................................................ Origin
Ox=cw/2
Oy=ch/2
# ............................................... Start/Stop function
def StartStop():
  global RunIter
  RunIter=not RunIter
  if RunIter:
    StartButton["text"]="Stop"
  else:
    StartButton["text"]="Restart"
# ..................................................... Exit function
def StopAll():
  global RunAll
  RunAll=False
# ................................................ Read Data function
def ReadData(*arg):
  global GetData
  GetData=True 
# .......................................................... Scale Up
def ScaleUp(*arg):
  global scale
  scale*=sqrt(2.0)
  ScaleLab['text']="%.4e"%(scale)
# ........................................................ Scale Down
def ScaleDown(*arg):
  global scale
  scale/=sqrt(2.0)
  ScaleLab['text']="%.4e"%(scale)
# ........................................... evaluate center of mass
def baryc(bd):
  mtot=sum(zz.m for zz in bd)
  cx=sum(zz.x*zz.m for zz in bd)/mtot
  cy=sum(zz.y*zz.m for zz in bd)/mtot
  cvx=sum(zz.vx*zz.m for zz in bd)/mtot
  cvy=sum(zz.vy*zz.m for zz in bd)/mtot
  return [[cx,cy],[cvx,cvy]]
# ..................................... move origin to center of mass
def SetBaryc():
  global NewBaryc
  [xcm,ycm],[cvx,cvy]=baryc(bd)
  for zz in bd:
    zz.x-=xcm
    zz.y-=ycm
    zz.vx-=cvx
    zz.vy-=cvy
  NewBaryc=True
# ..................................................... Class CelBody
class CelBody:
  def __init__(self,mass,radius,r,v,aphangle,color):
    self.m=mass
    self.rad=radius
    self.r=r
    self.phi=radians(aphangle)
    self.x=r*cos(self.phi)
    self.y=r*sin(self.phi)
    self.pathmin=r*0.05
    self.vx=-v*sin(self.phi)
    self.vy=v*cos(self.phi)
    self.col=color
    self.image=canvas.create_oval(Ox+int(scale*self.x-self.rad),\
      int(Oy-scale*self.y+self.rad),int(Ox+scale*self.x+self.rad),\
        int(Oy-scale*self.y-self.rad),fill=self.col,outline=self.col)
    self.path=[self.x,self.y]*PathLength
    self.ScaledPath=[0.0,0.0]*PathLength
    self.PathImg=canvas.create_line(self.ScaledPath,fill=self.col)
   # ....................................................... move body
  def move(self):
    canvas.coords(self.image,Ox+scale*self.x-self.rad,\
      Oy-scale*self.y+self.rad,Ox+scale*self.x+self.rad,\
        Oy-scale*self.y-self.rad)
  def UpdatePath(self):
    if abs(self.x-self.path[-2])+abs(self.y-self.path[-1])>self.pathmin:
      del self.path[:2]
      self.path.append(self.x)
      self.path.append(self.y)
  def DrawPath(self):
    self.ScaledPath[::2]=[Ox+scale*zz for zz in self.path[::2]]
    self.ScaledPath[1::2]=[Oy-scale*zz for zz in self.path[1::2]]
    canvas.coords(self.PathImg,self.ScaledPath)
# ....................................................... root window
root=Tk()
root.title('AphelPlanets')
root.bind('<Return>',ReadData)
root.bind('<Control-plus>',ScaleUp)
root.bind('<Control-minus>',ScaleDown)
# ............................................................ 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)
# ............................................................ buttons
StartButton=Button(toolbar,text="Start",command=StartStop,width=11)
StartButton.grid(row=0,column=0,sticky=W)
AdjustButton=Button(toolbar, text="Set Barycenter",\
  command=SetBaryc,width=11)
AdjustButton.grid(row=0,column=1,sticky=W)
CloseButton=Button(toolbar, text="Exit", command=StopAll,width=11)
CloseButton.grid(row=0,column=2,columnspan=2,sticky=W)
nr=0
# .................................................... create bodies
bd=[]
bd.append(CelBody(1.9891e30,10.0,0.0,0.0,0.0,'#ffff00'))#   Sun
bd.append(CelBody(3.30e23,4.0,6.9816e10,3.8875e4,-25.49,'#888888'))#  Mercury
bd.append(CelBody(4.87e24,4.0,1.08939e11,3.4785e4,28.61681,'#dfdf00'))#  Venus
bd.append(CelBody(5.97e24,4.0,1.52097e11,2.9297e4,0.0,'#4444ff'))# Earth
bd.append(CelBody(6.42e23,4.0,2.492e11,2.1979e4,282.94719,'#ff0000'))# Mars
bd.append(CelBody(9.39e20,2.0,4.4541e11,1.65967e4,49.90411,'#000000'))# Ceres
bd.append(CelBody(1.90e27,6.0,8.1662e11,1.24719e4,271.3838,'#ffaa50'))# Jupiter
bd.append(CelBody(5.68e26,6.0,1.5145e12,9.10217e3,350.10981,'#dfdf00'))#  Saturn
bd.append(CelBody(8.68e25,6.0,3.008e12,6.48871e3,0.809504,'#8888ff'))#  Uranus
bd.append(CelBody(1.02e26,6.0,4.54e12,5.3842e3,305.17281,'#8888ff'))#  Neptune
bd.append(CelBody(1.46e22,2.0,7.37593e12,3.67702e3,121.1858,'#000000'))#  Pluto
nB=len(bd)
# ........................................ coordinates and barycenter
canvas.create_line(0,Oy,cw,Oy,fill="black")
canvas.create_line(Ox,0,Ox,ch,fill="black")
bc=canvas.create_oval(Ox-bcrad,Oy-bcrad,Ox+bcrad,Oy+bcrad,fill="black")
# ............................................................ masses
masses=[]
for zz in bd:
  masses.append(zz.m)
# ........................................................ value list
values=[dt,cycle,PathLength]
t=[0.0,dt]
# ....................................................... vect2bodies
def vect2bodies(vect,bodies):
  for i,zz in enumerate(bodies):
    zz.x=vect[4*i]
    zz.y=vect[4*i+1]
    zz.vx=vect[4*i+2]
    zz.vy=vect[4*i+3]
# ............................................... create input vector
def WriteInput(bodies,val,InpV,vect):
  nn=5*len(bodies)
  InpV[:nn:5]=[zz.m for zz in bodies]
  InpV[1:nn:5]=vect[::4]=[zz.x for zz in bodies]
  InpV[2:nn:5]=vect[1::4]=[zz.y for zz in bodies]
  InpV[3:nn:5]=vect[2::4]=[zz.vx for zz in bodies]
  InpV[4:nn:5]=vect[3::4]=[zz.vy for zz in bodies]
  InpV[nn::]=val
# ....................................................................
def ReadInput(InpV,bodies,masses,val,vect):
  nn=5*len(bodies)
  for i,zz in enumerate(bodies):
    zz.m=masses[i]=InpV[i*5]
    zz.x=vect[i*4]=InpV[i*5+1]
    zz.y=vect[i*4+1]=InpV[i*5+2]
    zz.vx=vect[i*4+2]=InpV[i*5+3]
    zz.vy=vect[i*4+3]=InpV[i*5+4]
  val[::]=InpV[nn::]
# ...................................................... input vector
InpV=[0]*(5*nB+len(values))
y=[0]*4*nB
WriteInput(bd,values,InpV,y)
nInp=len(InpV)
# ........................................................ Input list
InputList=[]
i=0
while i<nB:
  InputList.append('m'+sub[i])
  InputList.append('x'+sub[i])
  InputList.append('y'+sub[i])
  InputList.append('vx'+sub[i])
  InputList.append('vy'+sub[i])
  i+=1
InputList.append('dt')
InputList.append('Cycle/ms')
InputList.append('Tail')
# ................................................ Labels and Entries
LabInput=[]
EntryInput=[]
for i,zz in enumerate(InputList):
  LabInput.append(Label(toolbar,text=zz,font=('Helvetica',11)))
  LabInput[i].grid(row=1+nr%nRows,column=2*(nr//nRows))
  EntryInput.append(Entry(toolbar,bd=3,width=13,font='Helvetica 9'))
  EntryInput[i].grid(row=1+nr%nRows,column=1+2*(nr//nRows))
  EntryInput[i].insert(0,"{:.5e}".format(InpV[i]))
  nr+=1
# ........................................................ time label
CycleLab0=Label(toolbar,text="Period:",font=("Helvetica",11))
CycleLab0.grid(row=1+nr%nRows,column=2*(nr//nRows))
CycleLab=Label(toolbar,text="     ",font=("Helvetica",11))
CycleLab.grid(row=1+nr%nRows,column=1+2*(nr//nRows),sticky=W)
nr+=1
# ...................................................... elapsed time
YearLab0=Label(toolbar,text="Years:",font=("Helvetica",11))
YearLab0.grid(row=1+nr%nRows,column=2*(nr//nRows))
YearLab=Label(toolbar,text="     ",font=("Helvetica",11))
YearLab.grid(row=1+nr%nRows,column=1+2*(nr//nRows),sticky=W)
nr+=1
DaysLab0=Label(toolbar,text="Days:",font=("Helvetica",11))
DaysLab0.grid(row=1+nr%nRows,column=2*(nr//nRows))
DaysLab=Label(toolbar,text="     ",font=("Helvetica",11))
DaysLab.grid(row=1+nr%nRows,column=1+2*(nr//nRows),sticky=W)
nr+=1
# ....................................................... scale label
ScaleLab0=Label(toolbar,text="Scale:",font=("Helvetica",11))
ScaleLab0.grid(row=1+nr%nRows,column=2*(nr//nRows))
ScaleLab=Label(toolbar,text="%.4e"%(scale),font=("Helvetica",11))
ScaleLab.grid(row=1+nr%nRows,column=1+2*(nr//nRows),sticky=W)
nr+=1
# .......................................................... function
def dfdt(yInp,t,mm):
  nB=len(mm)
  x=yInp[0::4]
  y=yInp[1::4]
  distx=(tile(x,(len(x),1))).T-x
  disty=(tile(y,(len(y),1))).T-y
  alpha=arctan2(disty,distx)
  r2=square(distx)+square(disty)
  fill_diagonal(r2,1.0)
  ff=divide(G,r2)
  fill_diagonal(ff,0.0)
  fx=ff*cos(alpha)
  fy=ff*sin(alpha)
  fx=(fx.T*mm).T
  fy=(fy.T*mm).T
  ax=fx.sum(axis=0)
  ay=fy.sum(axis=0)
  derivs=[0]*len(yInp)
  derivs[::4]=yInp[2::4]     # vx
  derivs[1::4]=yInp[3::4]    # vy
  derivs[2::4]=ax 
  derivs[3::4]=ay
  return derivs
# ........................................... numerical time interval
tt0=time.time()
iter=0
tcount=0
while RunAll:
  StartIter=time.time()
  # ..................................................... draw bodies
  i=0
  while i<nB:
    bd[i].move()
    bd[i].UpdatePath()
    bd[i].DrawPath()
    i+=1
  # ............................................. draw center of mass
  cx,cy=baryc(bd)[0]
  cx*=scale
  cy*=scale
  canvas.coords(bc,Ox+cx-bcrad,Oy-cy-bcrad,Ox+cx+bcrad,Oy-cy+bcrad)
  canvas.update()
  # .......................................................... motion
  if RunIter:
    # ..................................................... next step
    psoln = odeint(dfdt,y,t,args=(masses,))
    y=psoln[1,:]
    vect2bodies(y,bd)
    tcount+=1
  # .................................................................
  else:
    if NewBaryc:
      ReWrite=True
      WriteInput(bd,values,InpV,y)
      NewBaryc=False
    elif GetData:
      for i,zz in enumerate(EntryInput):
        try:
          InpV[i]=float(zz.get())
        except ValueError:
          pass
      tcount=0
      ReWrite=True
      GetData=False
    if ReWrite:
      ReadInput(InpV,bd,masses,values,y)
      for yy,zz in zip(InpV,EntryInput):
        zz.delete(0,'end')
        zz.insert(0,"{:.5e}".format(yy))
      dt=values[0]
      t=[0.0,dt]
      cycle=int(values[1])
      PathLength=int(values[2])
      for zz in bd:
        zz.path=[zz.x,zz.y]*PathLength
        zz.ScaledPath=[0.0,0.0]*PathLength
      ReWrite=False
  # ................................................ cycle duration
  iter+=1
  if iter%10==0:
    iter=0
    ttt=time.time()
    elapsed=ttt-tt0
    CycleLab['text']="%8.3f"%(elapsed*100.0)+" ms"
    tt0=ttt
  if tcount%40==0:
    totsec=dt*tcount
    ElapsYear=totsec/tYear
    YearLab['text']="%8.3f"%(ElapsYear)
    ElapsDays=totsec/tDay
    DaysLab['text']="%8.3f"%(ElapsDays)
  ElapsIter=int((time.time()-StartIter)*1000.0)
  canvas.after(cycle-ElapsIter)
root.destroy()
root.mainloop()