#!/usr/bin/env python3
#coding: utf8
import os
from numpy import *
from tkinter import *
from random import *
from scipy.integrate import odeint
import time
# .................................................. Global variables
RunAll=True
RunIter=False
NewBaryc=False
GetData=False
ReWrite=False
# ...................................................................
wake1x=[]
wake1y=[]
wake2x=[]
wake2y=[]
nWake=200
deg=180.0/pi
# ...................................................................
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 

def SetBaryc():                     # evaluate barycenter
  global NewBaryc
  global y
  global params
  x1,vx1,y1,vy1,x2,vx2,y2,vy2=y     # unpack current values of y
  m1,m2,KG,cycle,nWake=params
  mtot=m1+m2
  xc=(m1*x1+m2*x2)/mtot
  yc=(m1*y1+m2*y2)/mtot
  vcx=(m1*vx1+m2*vx2)/mtot
  vcy=(m1*vy1+m2*vy2)/mtot
  x1-=xc
  y1-=yc
  x2-=xc
  y2-=yc
  vx1-=vcx
  vy1-=vcy
  vx2-=vcx
  vy2-=vcy
  y=[x1,vx1,y1,vy1,x2,vx2,y2,vy2]
  NewBaryc=True
  
root=Tk()
root.title("Gravity 2 bodies")
root.bind('<Return>',ReadData)
# ...................................................................... canvas
cw=800
ch=800
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=1,column=0,columnspan=2,sticky=W)

ReadButton=Button(toolbar, text="Read Data", command=ReadData,width=11)
ReadButton.grid(row=2,column=0,columnspan=2,sticky=W)

CloseButton=Button(toolbar, text="Exit", command=StopAll,width=11)
CloseButton.grid(row=3,column=0,sticky=W)

# ...............................................................................
LabVar=[]
EntryVar=[]
LabPar=[]
EntryPar=[]
VarList=["x\u2081","vx\u2081","y\u2081","vy\u2081","x\u2082","vx\u2082",\
  "y\u2082","vy\u2082"]
nVar=len(VarList)
ParList=["m\u2081","m\u2082","G","Cycle/ms","Wake Length"]
nPar=len(ParList)

nr=4
for i in range(len(VarList)):
  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=8))
  EntryVar[i].grid(row=nr,column=1)
  nr+=1
for i in range(len(ParList)):
  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=8))
  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
# ........................................................... frames per second
fps=25.0
cycle=20   #milliseconds
# ...................................................................... Origin
Ox=cw/2
Oy=ch/2
# ..................................................................... gravity
KG=5000.0
# ...................................................................... mass 1
m1=10.0
x1=300.0
y1=0.0;
vx1=0.0
vy1=20.0
col1="red"
rad1=10
# ...................................................................... mass 2
m2=200.0
x2=0
y2=0.0;
vx2=0.0
vy2=0.0
col2="blue"
rad2=10
# ........................................................... barycenter radius
bcrad=2
#............................................................... initial values
y=[x1,vx1,y1,vy1,x2,vx2,y2,vy2]
# .................................................................. parameters
params=[m1,m2,KG,cycle,nWake]
# .......................................................... initialize entries
for i in range(len(VarList)):
  EntryVar[i].insert(0,str(y[i]))

for i in range(len(ParList)):
  EntryPar[i].insert(0,str(params[i]))
  
# .................................................................... function
def f(yDer, t, params):
  x1,vx1,y1,vy1,x2,vx2,y2,vy2=yDer     # unpack current values of y
  m1,m2,KG,cycle,nWake=params                # unpack parameters
  deltax=x2-x1
  deltay=y2-y1
  r2=deltax**2+deltay**2
  alpha=arctan2(deltay,deltax)
  force=KG*m1*m2/r2
  fx=force*cos(alpha)
  fy=force*sin(alpha)
  ax1=fx/m1
  ay1=fy/m1
  ax2=-fx/m2
  ay2=-fy/m2
  derivs = [vx1,ax1,vy1,ay1,vx2,ax2,vy2,ay2] # list of dy/dt=f functions
  return derivs
# ........................................... numerical time interval
t=[0.0,0.1]
tt0=time.time()
tcount=0
while RunAll:
  StartIter=time.time()
  # ..................................................... draw bodies
  canvas.delete(ALL)
  canvas.create_line(0,Oy,cw,Oy,fill="green")
  canvas.create_oval(Ox-bcrad,Oy-bcrad,Ox+bcrad,Oy+bcrad,fill="black")
  xx1=Ox+x1
  yy1=Oy-y1
  xx2=Ox+x2
  yy2=Oy-y2
  cx=(m1*xx1+m2*xx2)/(m1+m2)
  cy=(m1*yy1+m2*yy2)/(m1+m2)
  canvas.create_line(xx1,yy1,xx2,yy2,fill="yellow")
  canvas.create_oval(cx-bcrad,cy-bcrad,cx+bcrad,cy+bcrad,fill="black")
  canvas.create_oval(xx1-rad1,yy1-rad1,xx1+rad1,yy1+rad1,fill=col1)
  canvas.create_oval(xx2-rad2,yy2-rad2,xx2+rad2,yy2+rad2,fill=col2)
  # ....................................................... draw wake
  k=1
  while k<len(wake1x):
    canvas.create_line(Ox+wake1x[k-1],Oy-wake1y[k-1],Ox+wake1x[k],\
      Oy-wake1y[k],fill=col1)
    k+=1
  k=1
  while k<len(wake2x):
    canvas.create_line(Ox+wake2x[k-1],Oy-wake2y[k-1],Ox+wake2x[k],\
      Oy-wake2y[k],fill=col2)
    k+=1
  if RunIter:
    # ..................................................... next step
    psoln = odeint(f,y,t,args=(params,))
    x1=psoln[1,0]
    vx1=psoln[1,1]
    y1=psoln[1,2]
    vy1=psoln[1,3]
    x2=psoln[1,4]
    vx2=psoln[1,5]
    y2=psoln[1,6]
    vy2=psoln[1,7]
    # ...............................................................
    y=psoln[1,:]
    # .................................................. update wakes
    wake1x.append(int(x1))
    wake1y.append(int(y1))
    wake2x.append(int(x2))
    wake2y.append(int(y2))
    while len(wake1x)>nWake:
      pip=wake1x.pop(0)
      pip=wake1y.pop(0)
      pip=wake2x.pop(0)
      pip=wake2y.pop(0)
  else:
    if NewBaryc:
      ReWrite=True
      NewBaryc=False
    elif GetData:
      for i in range(nVar):
        try:
          y[i]=float(EntryVar[i].get())
        except ValueError:
          pass
      for i in range(nPar):
        try:
          params[i]=float(EntryPar[i].get())
        except ValueError:
          pass
      ReWrite=True
      GetData=False
    if ReWrite:
      x1,vx1,y1,vy1,x2,vx2,y2,vy2=y
      m1,m2,KG,cycle,nWake=params
      for i in range(nVar):
        EntryVar[i].delete(0,"end")
        EntryVar[i].insert(0,"{:.2f}".format(y[i]))
      for i in range(nPar):
        EntryPar[i].delete(0,"end")
        EntryPar[i].insert(0,"{:.2f}".format(params[i]))
      cycle=int(cycle)
      nWake=int(nWake)
      params=[m1,m2,KG,cycle,nWake]
      ReWrite=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()
canvas.mainloop()