#!/usr/bin/env python3
import numpy as np
from tkinter import *
#from math import degrees
from scipy.integrate import odeint
import scipy.optimize as optimize
import sys
import time

# ............................................. subscripts for labels
sub=['\u2080','\u2081','\u2082','\u2083','\u2084','\u2085',\
  '\u2086','\u2087','\u2088','\u2089','\u2081\u2080','\u2018\u2081 ']
# .................................................. Global variables
RunAll=True
ButtWidth=20
nRows=30
scale=1.3e-9
cw=700
ch=700
bcrad=2
cycle=5             # ms
#dt=2.16e4            # 0.25 day
dt=1.08e4            # 0.125 day
PathLength=200
G=6.67408e-11
tHour=3600
tDay=86400
#tYear=3.1557600e7   # Julian
#tYear=3.1556952e7   # Gregorian
tYear=3.1558149504e7 # Sidereal
# ..................................................................
StartYear=2020
StartTuple=(2020,1,1,0,0,0,0,0,0)
StartSec=time.mktime(StartTuple)
# ...................................................................
BestStart=[]
BestArrival=[]
BestDist=1.0e30
BestVel=0.0
# ............................................................ 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*=np.sqrt(2.0)
  ScaleLab['text']="%.4e"%(scale)
# ........................................................ Scale Down
def ScaleDown(*arg):
  global scale
  scale/=np.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]]
# ..................................................... Class CelBody
class CelBody:
  def __init__(self,mass,radius,x,y,vx,vy,color,name):
    self.m=mass
    self.rad=radius
    self.x=self.xx=x
    self.y=self.yy=y
    self.r=np.sqrt(x*x+y*y)
    angle=np.arctan2(self.y,self.x)
    self.phi=(angle+2.0*np.pi)%(2*np.pi)
    self.pathmin=np.sqrt(self.x**2+self.y**2)*0.05
    self.vx=vx
    self.vy=vy
    self.col=color
    self.name=name
    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):
    angle=np.arctan2(self.y,self.x)
    self.phi=(angle+2.0*np.pi)%(2*np.pi)
    self.r=np.sqrt(self.x*self.x+self.y*self.y)
    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('Mars Project 2020')
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)
# ........................................................... canvas2
canvas2=Canvas(root,width=cw,height=100,background="#ffffff")
canvas2.grid(row=1,column=0)
# ........................................................... toolbar
toolbar=Frame(root)
toolbar.grid(row=0,column=1, sticky=N)
# ............................................................ buttons
CloseButton=Button(toolbar, text="Exit", command=StopAll,width=11)
CloseButton.grid(row=0,column=0,columnspan=2,sticky=W)
nr=1
# ........................................................ 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]
# ....................................................... bodies2vect
def bodies2vect(bodies):
  vect=[]
  for zz in bodies:
    vect.append(zz.x)
    vect.append(zz.y)
    vect.append(zz.vx)
    vect.append(zz.vy)
  return vect
# ........................................................ Input list
InputList=[]
InputList.append('dt')
InputList.append('Cycle/ms')
InputList.append('Tail')
# ................................................ Labels and Entries
LabInput=[]
EntryInput=[]
nr=0
for i,zz in enumerate(InputList):
  LabInput.append(Label(toolbar,text=zz,font=('Helvetica',11)))
  LabInput[i].grid(row=2+nr%nRows,column=2*(nr//nRows))
  EntryInput.append(Entry(toolbar,bd=3,width=13,font='Helvetica 9'))
  EntryInput[i].grid(row=2+nr%nRows,column=1+2*(nr//nRows))
  EntryInput[i].insert(0,"{:.5e}".format(values[i]))
  nr+=1
# ........................................................ time label
CycleLab0=Label(toolbar,text="Period:",font=("Helvetica",11))
CycleLab0.grid(row=2+nr%nRows,column=2*(nr//nRows))
CycleLab=Label(toolbar,text="     ",font=("Helvetica",11))
CycleLab.grid(row=2+nr%nRows,column=1+2*(nr//nRows),sticky=W)
nr+=1
# ...................................................... elapsed time
YearLab0=Label(toolbar,text="Years:",font=("Helvetica",11))
YearLab0.grid(row=2+nr%nRows,column=2*(nr//nRows))
YearLab=Label(toolbar,text="     ",font=("Helvetica",11))
YearLab.grid(row=2+nr%nRows,column=1+2*(nr//nRows),sticky=W)
nr+=1
DaysLab0=Label(toolbar,text="Days:",font=("Helvetica",11))
DaysLab0.grid(row=2+nr%nRows,column=2*(nr//nRows))
DaysLab=Label(toolbar,text="     ",font=("Helvetica",11))
DaysLab.grid(row=2+nr%nRows,column=1+2*(nr//nRows),sticky=W)
nr+=1
# ....................................................... scale label
ScaleLab0=Label(toolbar,text="Scale:",font=("Helvetica",11))
ScaleLab0.grid(row=2+nr%nRows,column=2*(nr//nRows))
ScaleLab=Label(toolbar,text="%.4e"%(scale),font=("Helvetica",11))
ScaleLab.grid(row=2+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=(np.tile(x,(len(x),1))).T-x
  disty=(np.tile(y,(len(y),1))).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*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
# .....................................................................
def MarsFunc(var):
  global BestDist,method,t0_opt
  IterLaunch,vlaunch=var
  bd=[]
  bd.append(CelBody(1.9891e30,10.0,0.0,0.0,0.0,0.0,'#ffff00','Sun'))
  # ...................................................... 2020.01.01
  bd.append(CelBody(5.97e24,4.0,-2.651727e10,1.4665045e11,-2.931961e4,\
    -5.7949783e3,'#4444ff','Earth'))
  bd.append(CelBody(6.42e23,4.0,-8.8964577e10,-2.2819209e11,2.0381955e4,\
    -9.3177689e3,'#ff0000','Mars'))
  nB=len(bd)
  masses=[]
  for zz in bd:
    masses.append(zz.m)
  y=bodies2vect(bd)
  IsBack=launched=False
  tt0=time.time()
  iter=0
  tcount=0
  # ........................................ 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")
  canvas.create_text(10,10,anchor=NW,text='01.JAN.2020',font=('Helvetica','10'))
  canvas.create_text(10,30,anchor=NW,text='Method: '+method,\
    font=('Helvetica','12'))
  while RunAll:
    StartIter=time.time()
    # ............................................. Launch spacecraft
    if tcount>=IterLaunch and (not launched):
      dist0=2.0e8
      if tcount>IterLaunch:
        deltat=tcount-IterLaunch
        vEarth=np.sqrt(bd[1].vx**2+bd[1].vy**2)
        dist0+=vEarth*deltat
      xsc=bd[1].x-dist0*np.sin(bd[1].phi)
      ysc=bd[1].y+dist0*np.cos(bd[1].phi)
      vxsc=bd[1].vx-vlaunch*np.sin(bd[1].phi)
      vysc=bd[1].vy+vlaunch*np.cos(bd[1].phi)
      bd.append(CelBody(1.0e5,1.0,xsc,ysc,vxsc,vysc,'#000000','Probe'))
      nB=len(bd)
      y=bodies2vect(bd)
      masses.append(1.0e5)
      scphi0=bd[3].phi
      rmax=bd[3].r
      distSM=np.sqrt((bd[2].x-bd[3].x)**2+(bd[2].y-bd[3].y)**2)
      distSMmin=distSM
      LaunchSec=StartSec+tcount*dt
      LaunchDate=time.ctime(LaunchSec)
      launched=True
    # ................................................... draw bodies
    i=0
    while i<nB:
      bd[i].move()
      bd[i].UpdatePath()
      bd[i].DrawPath()
      i+=1
    # ...................................... distance spacecraft-Mars
    if launched:
      distSM=np.sqrt((bd[2].x-bd[3].x)**2+(bd[2].y-bd[3].y)**2)
      if distSM<distSMmin:
        distSMmin=distSM
        arrival=tcount
    # ........................................... 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
    psoln = odeint(dfdt,y,t,args=(masses,))
    y=psoln[1,:]
    vect2bodies(y,bd)
    # ...................................................
    if launched:
      scphi=bd[3].phi
      DeltaPhi=(scphi+2*np.pi-scphi0)%(2*np.pi)
      if DeltaPhi>3.5 and DeltaPhi<5.0:    #6.1:
        canvas.delete(ALL)
        if distSMmin<BestDist:
          BestDist=distSMmin
          BestStart=IterLaunch
          BestArrival=arrival
          BestVel=vlaunch
          TravelTime=(BestArrival-BestStart)*dt/tDay
          StartDate=time.ctime(StartSec+BestStart*dt)
          ArrDate=time.ctime(StartSec+BestArrival*dt)
          BestEvalTime=time.time()-t0_opt
          canvas2.delete(ALL)
          canvas2.create_text(10,10,anchor=NW,text='Current best launch',\
            font=('Helvetica','10'))
          canvas2.create_text(200,10,anchor=NW,text='Evaluation time: '+\
            '{:.3f} s'.format(BestEvalTime),font=('Helvetica','10'))
          canvas2.create_text(10,30,anchor=NW,text='Departure: '+\
            StartDate+'          Arrival: '+ArrDate,font=('Helvetica','10'))
          canvas2.create_text(10,50,anchor=NW,text='Minimum distance:  '+\
            '{:.3e} m'.format(BestDist)+'     Initial velocity :  '+\
              '{:.2f} m/s'.format(BestVel),font=('Helvetica','10'))
          canvas2.create_text(10,70,anchor=NW,text='Travel time:  '+\
            '{:.3f} days'.format(TravelTime),font=('Helvetica','10'))
          canvas2.update()
          # .........................................................
          print('Evaluation time: {:.3f} s'.format(BestEvalTime))
          print('Departure: '+StartDate)
          print('Arrival: ' +ArrDate)
          print('Initial velocity:  {:.2f} m/s'.format(BestVel))
          print('Travel time: {:.3f} days'.format(TravelTime))
          print('Minimum distance:  {:.3e} m'.format(BestDist))
          print('----------------------------')
          # .........................................................
        return distSMmin
    # ...................................................
    tcount+=1
    # ................................................ 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)
  if not RunAll:
    root.destroy()
    sys.exit()
# ......................................................................
IterLaunch=1520
vlaunch=2.8e3
guess0=[IterLaunch,vlaunch]
'''
#hnd=open('mars.txt','w')
while IterLaunch<=2500:
  x=MarsFunc(guess0)
  hnd.write('{:15.8e}{:5d}\n'.format(x,IterLaunch))
  IterLaunch+=20
hnd.close()
'''
'''
x=optimize.minimize(MarsFunc,guess0,method='nelder-mead',\
  options={'xtol': 1e-5, 'disp': True})
'''
#method='COBYLA'
#method='Nelder-Mead'
method='Powell'
t0_opt=time.time()
x=optimize.minimize(MarsFunc,guess0,method=method,\
  options={'xtol': 1e-5, 'disp': True})
while True:
  if not RunAll:
    canvas.after(300)
    root.destroy()