## solves extended gravitational problem
## mass moving in central potential r^n

IN_IDLE = 0

STEPSIZE = 0.002    # genauer und langsamer
#STEPSIZE = 0.005

from Tkinter import *

# rk5(6) by Butcher  
def rk5(t,x,f,h):
    n = len(x)
    xnew = [0]*n

    k1 = f(t,x)
    for i in range(n): k1[i] = k1[i]*h

    for i in range(n): xnew[i] = x[i]+k1[i]/8.0
    k2 = f(t+h/8.0,xnew)
    for i in range(n): k2[i] = k2[i]*h

    for i in range(n): xnew[i] = x[i]+k2[i]/4.0
    k3 = f(t+h/4.0,xnew)
    for i in range(n): k3[i] = k3[i]*h

    for i in range(n): xnew[i] = x[i]+k1[i]/2.0-k2[i]+k3[i]
    k4 = f(t+h/2.0,xnew)
    for i in range(n): k4[i] = k4[i]*h

    for i in range(n):
        xnew[i] = x[i]+(k1[i]*3.0+k4[i]*9.0)/16.0
    k5 = f(t+h*3.0/4.0,xnew)
    for i in range(n): k5[i] = k5[i]*h

    for i in range(n):
        xnew[i] = x[i]+(-5.0*k1[i]+4.0*k2[i]+12.0*k3[i]-12.0*k4[i]+8.0*k5[i])/7.0
    k6 = f(t+h,xnew)
    for i in range(n): k6[i] = k6[i]*h

    for i in range(n): 
        xnew[i] = x[i]+(7.0*k1[i]+32.0*k3[i]+12.0*k4[i]+32.0*k5[i]+7.0*k6[i])/90.0

    tnew = t+h
    return tnew,xnew


class ForceWindow:
    def __init__(self,size=500):
        self.root = Tk()
        self.root.title("HIB - Mass Moving In A Central Force Field")
        font1 = ("Arial",10,"bold")
        font2 = ("Arial",10)

        commands = Frame()
        b1 = Button(commands,text="Start",font=font1,command=self.start).pack(side=LEFT,padx=8)
        b2 = Button(commands,text="Stop",font=font1,command=self.stop).pack(side=LEFT,padx=8)
        b3 = Button(commands,text="Quit",font=font1,command=self.quit).pack(side=LEFT,padx=8)
         
        self.c = Canvas(self.root,width=size+6,height=size+6,bg="#f0f0f0")
        self.size = size
        
        params = Frame()

        self.v_exponent = StringVar()
        self.v_vx = StringVar()
        self.v_x0 = StringVar()
        self.v_vy = StringVar()
        self.v_zoom = StringVar()
        Label(params,text="\nF = -r**n\n\n",font=font1).grid(column=0,row=0,columnspan=2)
        Label(params,text="n = ",font=font2).grid(column=0,row=1)
        Entry(params,width=4,textvariable=self.v_exponent,font=font2).grid(column=1,row=1)
        Label(params,text="x0 = ",font=font2).grid(column=0,row=2)
        Entry(params,width=4,textvariable=self.v_x0,font=font2).grid(column=1,row=2)
        Label(params,text="vx = ",font=font2).grid(column=0,row=3)
        Entry(params,width=4,textvariable=self.v_vx,font=font2).grid(column=1,row=3)
        Label(params,text="vy = ",font=font2).grid(column=0,row=4)
        Entry(params,width=4,textvariable=self.v_vy,font=font2).grid(column=1,row=4)
        Label(params,text="\nzoom out",font=font2).grid(column=0,row=5,columnspan=2)
        Entry(params,width=4,textvariable=self.v_zoom,font=font2).grid(column=0,row=6,columnspan=2)
        self.v_exponent.set("-2")
        self.v_x0.set("1.0")
        self.v_vx.set("0.0")
        self.v_vy.set("1.0")
        self.v_zoom.set("1.0")
        
        commands.pack(side=BOTTOM,padx=4,pady=4)
        params.pack(side=LEFT,padx=4,pady=4)
        self.c.pack(side=RIGHT)

        self.rsun = 0.075
        self.rplanet = 0.02
        self.rmax = 2
        self.startx,self.starty = float(self.v_x0.get()),0.0

        self.WORKING = 0
        self.RUNNING = False 
        self.setup()
        if not IN_IDLE: self.root.mainloop()

    def setup(self):
        self.c.delete("all")
        self.MID = self.size/2+3
        self.scale = self.size/2.0/self.rmax/float(self.v_zoom.get())
        # show sun
        r = self.scale*self.rsun
        self.c.create_oval(self.MID-r,self.MID-r,self.MID+r,self.MID+r,
                           outline="#c0c000",fill="#c0c000")
        # show planet
        r = self.scale*self.rplanet
        self.startx,self.starty = float(self.v_x0.get()),0.0
        xx,yy = self.scale*self.startx,self.scale*self.starty
        self.c.create_oval(self.MID+xx-r,self.MID-yy-r,self.MID+xx+r,self.MID-yy+r,
                           outline="#0000f0",fill="#0000f0")
        self.oldxx,self.oldyy = xx,yy

    def force(self,t,values):
        vx,vy,x,y = values
        r = (x*x+y*y)**0.5
        f = -r**self.exponent
        fx = f*x/r
        fy = f*y/r
        return [fx,fy,vx,vy]

    def start(self):
        self.setup()
        self.RUNNING = True
        vx,vy,x,y = float(self.v_vx.get()),float(self.v_vy.get()),float(self.v_x0.get()),0.0
        self.exponent = float(self.v_exponent.get())
        t,initial = 0.0,[vx,vy,x,y]
        h = STEPSIZE
        min_r2 = self.rsun**2
        while self.RUNNING:
            x,y = initial[2],initial[3]
            self.plot(x,y)
	    
	    ## make two steps 
            if (x*x+y*y)<min_r2: break 
            t,initial = rk5(t,initial,self.force,h)
            if (x*x+y*y)<min_r2: break 
            t,initial = rk5(t,initial,self.force,h)
            
        self.RUNNING = False

    def plot(self,x,y):
        xx,yy = self.scale*x,self.scale*y
        self.c.create_line(self.MID+xx,self.MID-yy,self.MID+xx+1,self.MID-yy)
        if (self.oldxx-xx)**2+(self.oldyy-yy)**2>1000:
            ## force too big - nonsense integration result, mass shoots away
            self.RUNNING=False
            return
        self.oldxx,self.oldyy = xx,yy
        self.c.update()
    
             
    def stop(self):
        self.RUNNING = False

    def quit(self):
        self.RUNNING = False
        self.root.after(50)
        self.root.destroy()

w = ForceWindow(500)