from __future__ import print_function, division
from visual.controls import *
import os
from time import clock

# This program creates orbit files the first time it is run.

print("""
Press to choose initial position for green star.
Then drag to choose its initial velocity.
During an orbit, click to stop, then click
  again to define new initial conditions.
The files x=-2.76.orb and x=-2.765.orb differ
  only very slightly in the initial value of x.
The intent of the program is to demonstrate the
  great complexity of motion with just 3 objects,
  despite simple rules (Newtonian mechanics), and
  the high sensitivity to initial conditions in
  such a system.
If you first rotate the scene before dragging,
  you can make a nonplanar orbit.
Bruce Sherwood, January 2003.
""")

scene.width = sw = 600
scene.height = sh = 600
scene.x = scene.y = 0
xmax = 2.8
G = 1. # gravitational constant in arbitrary units
M_to_R = 0.07 # convert mass to radius by M_to_R*mass**(1./3.)
vstar = 2 # The star whose velocity we'll choose by dragging
save = zeros(3*7, 'float') # to save pos, vel, mass at start of each orbit

def save_file(file_extensions=None, x=100, y=100, title="Save", mode='w', maxfiles=20):
    if maxfiles < 5: maxfiles = 5
    if mode[0] != 'w':
        raise ValueError("To save a file, mode must start with 'w'.")
    if file_extensions is not None:
        if isinstance(file_extensions, (list,tuple)):
            raise ValueError("Only one file extension can be specified.")
        file_extensions = [file_extensions]
    return filedialog(file_extensions=file_extensions,
                   x=x, y=y, title=title, mode=mode, maxfiles=maxfiles)

def get_file(file_extensions=None, x=100, y=100, title="Open", mode='rU', maxfiles=20):
    if maxfiles < 5: maxfiles = 5
    if mode[0] != 'r' and mode != 'U':
        raise ValueError("To read a file, mode must start with 'r'.")
    if file_extensions is not None:
        if not isinstance(file_extensions, (list,tuple)):
            file_extensions = [file_extensions]
    return filedialog(file_extensions=file_extensions,
                   x=x, y=y, title=title, mode=mode, maxfiles=maxfiles)

def filedialog(file_extensions=None, x=100, y=100, title="Open", mode='rU', maxfiles=20):
    # file_extensions is a list of types (reading) or a 1-element list (writing)
    writing =  not (mode[0] == 'r' or mode == 'U')
    filecolor = color.black
    dircolor = (0,0.5,0.5)
    selectcolor = (.7,1,1)
    inactivecolor = (0.8,0.8,0.8)
    activecolor = (0.6,0.6,0.6)
    winwidth = 300
    hitem = 20 # pixel height of each menu listing
    ctrls = 20 # button area
    hcanvas = hitem*(maxfiles+3)+ctrls # approx height without title bar
    currentdisplay = display.get_selected()
    menus = display(background=color.white, foreground=color.black, exit=0, range=100,
                       x=x, y=y, title=title, fov=0.001,
                       # allow 30 pixels for title bar
                       width=winwidth, height=30+hcanvas+ctrls)
    if menus.width >= hcanvas+ctrls:
        hmenu = 200.*hitem/menus.width
        ytop = 100.*hcanvas/menus.width-2*hmenu
        ybottom = ytop-maxfiles*hmenu
        ytopscroll = 100.*(hcanvas+ctrls)/menus.width
        ybottomscroll = -ytopscroll
        xmax = 100.
        xmin = -xmax
        scroll_edge = 87.
        okaypos = (20,ybottom-1.2*hmenu,0.1)
        cancelpos = (60,ybottom-1.2*hmenu,0.1)
        okaysize = (35,1.1*hmenu,0.1)
    else:
        hmenu = 200.*hitem/hcanvas
        ytop = 100.-2*hmenu
        ybottom = ytop-maxfiles*hmenu
        ytopscroll = 100.
        ybottomscroll = -ytopscroll
        xmax = 100.*menus.width/(hcanvas+ctrls)
        xmin = -xmax
        scroll_edge = xmax-14*menus.width/hcanvas
        okaypos = (20.*menus.width/hcanvas,ybottom-1.2*hmenu,0.1)
        cancelpos = (60.*menus.width/hcanvas,ybottom-1.2*hmenu,0.1)
        okaysize = (35.*menus.width/hcanvas,1.1*hmenu,0.1)
    okay = box(pos=okaypos, size=okaysize, color=inactivecolor)
    if writing:
        t = 'Save'
    else:
        t = 'Open'
    okaylabel = label(pos=okaypos, text=t, color=color.black, opacity=0, box=0)
    cancel = box(pos=cancelpos, size=okaysize, color=inactivecolor)
    label(pos=cancelpos, text='Cancel', color=color.black, opacity=0, box=0)
    go_up = box(pos=(cancel.pos.x,ytop+hmenu,0), size=okaysize, color=inactivecolor)
    go_up_arr = arrow(pos=(go_up.pos.x,go_up.pos.y-0.4*hmenu,0), axis=(0,0.9*hmenu,0),
          fixedwidth=1, shaftwidth=0.3*hmenu, color=(0.9,0,0))
    showdir = label(pos=(xmin+0.2*(xmax-scroll_edge),ytop+hmenu,0), text='',
                    color=dircolor, opacity=0, xoffset=1, box=0, line=0)
    if writing:
        # A bug in Visual 3 makes it necessary not to allow len(getname.text) < 1
        getname = label(pos=(xmin+0.5*(xmax-scroll_edge),okay.y,0), border = 2,
                        text='|', box=0, xoffset=1, line=0, opacity=0)
        z = 0.1
        curve(pos=[(getname.x,getname.y-0.5*hmenu,z), (getname.x,getname.y+0.5*hmenu,z),
                   (okay.x-1.2*0.5*okay.length,getname.y+0.5*hmenu,z), 
                   (okay.x-1.2*0.5*okay.length,getname.y-0.5*hmenu,z),
                   (getname.x,getname.y-0.5*hmenu,z)])

    labels = []
    for n in range(maxfiles):
        labels.append(label(pos=(xmin+0.5*(xmax-scroll_edge),ytop-n*hmenu),
                            text='', opacity=0, box=0,
                            xoffset=1, line=0))
        
    scrolltrack = box(pos=(0.5*(scroll_edge+xmax),0,0.02),
                      color=color.white, size=(xmax-scroll_edge,200,0.001), visible=0)
    scrollside = curve(pos=[(scroll_edge,100,.03),(scroll_edge,-100,.03)],
                       color=(0.8,0.8,0.8), visible=0)
    scroll = box(pos=(0.5*(scroll_edge+xmax),0,0.04), offset=vector(0,0,0),
                 color=inactivecolor, size=((xmax-scroll_edge+1),0,0.1), visible=0)
    
    shade = box(pos=(0,0,0), color=(.95,.95,.95), size=(200,hmenu,0.01), visible=0)
    select = box(pos=(0,0,0.01), color=selectcolor, size=(200,hmenu,0.01), visible=0)
    clicktime = -1

    while menus.visible:
        shade.visible = 0
        select.visible = 0
        highlighted = None
        selected = None
        changedir = False
        drag = False
        topmenu = 0
        showdir.text = os.path.split(os.getcwd())[-1]
        allfiles = os.listdir(os.curdir)
        files = []
        for f in allfiles:
            is_a_dir = os.path.isdir(f)
            if is_a_dir or (file_extensions is None):
                files.append([f,is_a_dir,False]) # file name, whether a directory, whether selected
            else:
                period = f.rfind('.')
                if period:
                    if f[period:] in file_extensions:
                        files.append([f,is_a_dir,False])
        Nfiles = len(files)
        need_to_scroll = (Nfiles > maxfiles)
        if need_to_scroll:
            Nfiles = maxfiles
            hscroll = (ytopscroll-ybottomscroll)*maxfiles/len(files)
            if hscroll < hmenu:
                hscroll = hmenu
            dy = (ytopscroll-ybottomscroll-hscroll)/(len(files)-maxfiles)
            scrolltrack.visible = 1
            scrollside.visible = 1
            scroll.y = ytopscroll-0.5*hscroll
            scroll.height = hscroll
            scroll.visible = 1
        else:
            scrolltrack.visible = 0
            scrollside.visible = 0
            scroll.visible = 0

        for n in labels:
            n.text = ''
        for n, f in enumerate(files[:Nfiles]):
            lcolor = filecolor
            if f[1]: lcolor = dircolor
            labels[n].text = f[0]
            labels[n].color = lcolor

        if writing:
            getname.text = '|'
            ending = '|'
        blink = clock()
        blinkon = True
        while menus.visible:
            rate(50)
            # A bug in Visual 3 makes it necessary not to allow len(getname.text) < 1
            if writing and clock()-blink > 0.5:
                blink = clock()
                blinkon = not blinkon
                if blinkon:
                    ending = '|'
                else:
                    ending = ' '
                if getname.text == '':
                    getname.text = ending
                elif getname.text == '|' or getname.text == ' ':
                    getname.text = ending
                elif getname.text[-1] == '|' or getname.text[-1] == ' ':
                    getname.text = getname.text[:-1]+ending 
                else:
                    getname.text += ending
            mpos = menus.mouse.pos
            if writing and menus.kb.keys: # event waiting to be processed?
                s = menus.kb.getkey() # get keyboard info; make sure string length never 0
                if s == '\n':
                    shade.visible = 0
                    select.visible = 0
                    ret = finish_save(getname.text, file_extensions, mode, menus, labels, currentdisplay)
                    if ret:
                        return ret
                    if highlighted: shade.visible = 1
                    if selected: select.visible = 1
                elif len(s) == 1 and s != '|':
                    if getname.text == '': # should never happen
                        getname.text = s+ending
                    elif getname.text[-1] == ending:
                        getname.text = getname.text[:-1]+s+ending
                    else:
                        getname.text = getname.text+s+ending # add new character
                elif s == 'backspace' or s == 'delete':
                    if getname.text == '': # should never happen
                        getname.text = ending
                    elif getname.text[-1] == ending:
                        getname.text = getname.text[:-2]+ending # erase character
                    else:
                        if len(getname.text) <= 1: # should never happen
                            getname.text = ending
                        else:
                            getname.text = getname.text[:-1]+ending # erase character
                elif s == 'shift+delete' or s == 'shift+backspace':
                    getname.text = ending # erase all text
            if menus.mouse.events:
                m = menus.mouse.getevent()
                mpos = m.pos
                nmenu = int((ytop+0.5*hmenu-mpos.y)/hmenu)
                if mpos.y > ytop+0.5*hmenu: nmenu = -1
                if drag and m.release == 'left':
                    drag = False
                elif need_to_scroll and m.pick == scroll:
                    scroll.color = dircolor
                    scroll.offset = scroll.y-mpos.y
                    drag = True
                elif m.click == 'left':

                    # Check for clicking go up or open or cancel
                    if m.pick == go_up or m.pick == go_up_arr:
                        os.chdir('../')
                        changedir = True
                        break
                    if m.pick == cancel:
                        return finish_get(None, mode, menus, currentdisplay)
                    elif m.pick == okay:
                        if writing and getname.text != '|':
                            shade.visible = 0
                            select.visible = 0
                            ret = finish_save(getname.text, file_extensions, mode, menus, labels, currentdisplay)
                            if ret:
                                return ret
                            if highlighted: shade.visible = 1
                            if selected: select.visible = 1
                        if selected is not None:
                            filename, is_a_dir, s = files[selected]
                            if is_a_dir:
                                os.chdir(filename)
                                changedir = True
                                break
                            elif not writing:
                                return finish_get(filename, mode, menus, currentdisplay)
                        
                    # Handle doubleclick on a name
                    clicktime = clock()-clicktime
                    if topmenu+nmenu == selected and clicktime < 0.5:
                        filename, is_a_dir, s = files[selected]
                        if is_a_dir:
                            os.chdir(filename)
                            changedir = True
                            break
                        else:
                            if writing:
                                getname.text = filename+'|'
                            else:
                                return finish_get(filename, mode, menus, currentdisplay)
                        
                    # Handle singleclick on a name
                    if (0 <= nmenu <= Nfiles-1):
                        select.y = ytop-nmenu*hmenu
                        select.visible = 1
                        if selected:
                            files[selected][2] = False
                        selected = topmenu+nmenu
                        files[selected][2] = True
                        if writing:
                            if files[selected][1]:
                                okaylabel.text = 'Open'
                            else:
                                okaylabel.text = 'Save'
                        clicktime = clock()
                    elif selected:
                        select.visible = 0
                        files[selected][2] = False
                        selected = 0
                        if writing:
                            okaylabel.text = 'Save'
                        
            if drag:
                newy = mpos.y+scroll.offset
                if newy+0.5*hscroll >= ytopscroll:
                    scroll.y = ytopscroll-0.5*hscroll
                    scroll.offset = scroll.y-mpos.y
                elif newy-0.5*hscroll <= ybottomscroll:
                    scroll.y = ybottomscroll+0.5*hscroll
                    scroll.offset = scroll.y-mpos.y
                else:
                    scroll.y = newy
                newtopmenu = int((ytopscroll-0.5*hscroll-scroll.y)/dy)
                if newtopmenu != topmenu:
                    topmenu = newtopmenu
                    select.visible = 0
                    for n, lab in enumerate(labels):
                        lab.text = files[n+topmenu][0]
                        if files[n+topmenu][2]:
                            select.y = ytop-n*hmenu
                            select.visible = 1
                        elif files[n+topmenu][1]:
                            lab.color = dircolor
                        else:
                            lab.color = filecolor
            else:
                if need_to_scroll:
                    if menus.mouse.pick == scroll:
                        scroll.color = activecolor
                    else:
                        scroll.color = inactivecolor
                        
                if menus.mouse.pick == go_up or menus.mouse.pick == go_up_arr:
                    go_up.color = activecolor
                else:
                    go_up.color = inactivecolor
                    
                if menus.mouse.pick == cancel:
                    cancel.color = activecolor
                else:
                    cancel.color = inactivecolor
                    
                if menus.mouse.pick == okay:
                    okay.color = activecolor
                else:
                    okay.color = inactivecolor
                    
                if need_to_scroll and mpos.x >= scroll_edge:
                    shade.visible = 0
                    highlighted = None
                else:
                    nmenu = int((ytop+0.5*hmenu-mpos.y)/hmenu)
                    if mpos.y > ytop+0.5*hmenu: nmenu = -1
                    if mpos.y < ytop+0.5*hmenu-hmenu*Nfiles: nmenu = -1
                    if (0 <= nmenu <= Nfiles):
                        shade.y = ytop-nmenu*hmenu
                        shade.visible = 1
                        highlighted = nmenu
                    else:
                        shade.visible = 0
                        highlighted = None
        if changedir: continue
        break
    currentdisplay.select()
    return None

def finish_get(filename, mode, menus, currentdisplay):
    menus.visible = 0
    del menus
    currentdisplay.select()
    if filename is None:
        return None
    return open(str(filename), mode)

def finish_save(filename, file_extensions, mode, menus, labels, currentdisplay):
    if filename == '':
        return None
    if filename[-1] == '|' or filename[-1] == ' ':
        filename = filename[:-1]
    if filename == '':
        return None
    t = filename.split(".")
    ext = ''
    if len(t) > 0:
        ext = '.'+t[-1]
    elif t[0] == '':
        return None
    if file_extensions is not None:
        if ext != file_extensions[0]:
            filename += file_extensions[0]
    try:
        fd = open(str(filename), 'r') # see whether file already exists
    except:
        menus.visible = 0
        del menus
        currentdisplay.select()
        return open(str(filename), mode)
    fd.close()
    for a in labels:
        a.visible = 0
    inactivecolor = (0.8,0.8,0.8)
    activecolor = (0.6,0.6,0.6)
    templabel = label(pos=(0,15,0.3), text="%s already exists" % filename, box=0, opacity=0)
    overwrite = box(pos=(-30,-5,0.4), size=(50,15,.1), color=inactivecolor)
    overlabel = label(pos=overwrite.pos, text="Overwrite", box=0, opacity=0)
    cancel = box(pos=(30,-5,0.4), size=(50,15,.1), color=inactivecolor)
    cancellabel = label(pos=cancel.pos, text="Cancel", box=0, opacity=0)
    while menus.visible:
        rate(50)
        if menus.mouse.events:
            m = menus.mouse.getevent()
            if m.click == 'left':
                if m.pick == overwrite:
                    menus.visible = 0
                    del menus
                    currentdisplay.select()
                    return open(str(filename), mode)
                if m.pick == cancel:
                    templabel.visible = 0
                    overwrite.visible = 0
                    overlabel.visible = 0
                    cancel.visible = 0
                    cancellabel.visible = 0
                    del templabel
                    del overwrite
                    del overlabel
                    del cancel
                    del cancellabel
                    for a in labels:
                        a.visible = 1
                    return None
        if menus.mouse.pick == overwrite:
            overwrite.color = activecolor
        else:
            overwrite.color = inactivecolor
        if menus.mouse.pick == cancel:
            cancel.color = activecolor
        else:
            cancel.color = inactivecolor
    return None

def buildfiles(): # If a standard *.orb file doesn't exist, create it
    orbitdata = [
            ('commute.orb', [-0.5, 0, 0, -0.0465, 0.663, 0, 2, 0.5, 0, 0, 0, -1.2, 0, 1, -0.65, 1.82, 0, 0.93, -1.26, 0, 0.1]),
            ('moon1.orb', [-0.5, 0, 0, 0, 0.68, 0, 2, 0.5, 0, 0, 0, -1.2, 0, 1, 0.12, 0.86, 0, 0, -1.6, 0, 0.1]),
            ('moon2.orb', [-0.5, 0, 0, 0.0525, 0.643, 0, 2, 0.5, 0, 0, 0, -1.2, 0, 1, 0.54, 0.49, 0, -1.05, -0.86, 0, 0.1]),
            ('moon3.orb', [-0.5, 0, 0, -0.0475, 0.6755, 0, 2, 0.5, 0, 0, 0, -1.2, 0, 1, -0.32, 1.01, 0, 0.95, -1.51, 0, 0.1]),
            ('moon4.orb', [-0.5, 0, 0, 0.0195, 0.7135, 0, 2, 0.5, 0, 0, 0, -1.2, 0, 1, 1.22, 0.05, 0, -0.39, -2.27, 0, 0.1]),
            ('rosette.orb', [-0.5, 0, 0, 0.0325, 0.63, 0, 2, 0.5, 0, 0, 0, -1.2, 0, 1, 0.91, 0.77, 0, -0.65, -0.6, 0, 0.1]),
            ('sevenloops.orb', [-0.5, 0, 0, -0.000450452, 0.516665, -2.96059e-017, 2, 0.5, 0, 0, -0.000450452, -1.28333, -2.96059e-017, 1, -0.108108, 0.0168919, -4.44089e-016, 0.0135136, 2.5, 8.88178e-016, 0.1]),
            ('shuttle1.orb', [-0.5, 0, 0, 0.0395, 0.6585, 0, 2, 0.5, 0, 0, 0, -1.2, 0, 1, 0.71, 0.8, 0, -0.79, -1.17, 0, 0.1]),
            ('shuttle2.orb', [-0.5, 0, 0, 0.0375, 0.654, 0, 2, 0.5, 0, 0, 0, -1.2, 0, 1, 0.7, 0.83, 0, -0.75, -1.08, 0, 0.1]),
            ('trefoil.orb', [-0.214352881110285, -0.368303127453297, 0, -0.477702690351465, -0.0690629362757092, 0, 2, -0.195743825051338, 0.735857211504369, 0, 1.13140538070293, 0.0241258725514185, 0, 1, 0.22, 0.04, 0, -1.76, 1.14, 0, 0.1]),
            ('weave.orb', [-0.5, 0, 0, 0, 0.5305, 0, 2, 0.5, 0, 0, 0, -1.2, 0, 1, 0.13, -0.81, 0, 0, 1.39, 0, 0.1]),
            ('weave2.orb', [-0.5, 0, 0, -0.038, 0.5775, 0, 2, 0.5, 0, 0, 0, -1.2, 0, 1, -1.92, -1.4, 0, 0.76, 0.45, 0, 0.1]),
            ('wild1.orb', [-0.5, 0, 0, -0.0275, 0.6005, 0, 2, 0.5, 0, 0, 0, -1.2, 0, 1, -2.79, -0.22, 0, 0.55, -0.01, 0, 0.1]),
            ('wildnonplanar.orb', [-0.5, 0, 0, -0.0324781, 0.577146, -0.0510757, 2, 0.5, 0, 0, -0.0324781, -1.22285, -0.0510757, 1, -0.490672, 0.438966, 0.647793, 0.974342, 0.685622, 1.53227, 0.1]),
            ('x=-2.76.orb', [-0.5, 0, 0, -0.023, 0.5995, 0, 2, 0.5, 0, 0, 0, -1.2, 0, 1, -2.76, -0.14, 0, 0.46, 0.01, 0, 0.1]),
            ('x=-2.765.orb', [-0.5, 0, 0, -0.023, 0.5995, 0, 2, 0.5, 0, 0, 0, -1.2, 0, 1, -2.765, -0.14, 0, 0.46, 0.01, 0, 0.1]),
        ]
    error = 0
    for item in orbitdata:
        file, data = item
        
        try:
            fd = open(file, 'r')
        except:
            pass # file does not exist
        else:
            continue # file does exist
        
        try:
            fd = open(file, 'w')
        except:
            error = 1
            continue
        
        for nn in range(3*7):
            try:
                fd.write('%g' % data[nn])
            except:
                error = 1
                print('Cannot write to file '+file)
            if ((nn+1) % 7) == 0:
                c = '\n'
            else:
                c = '\t'
            try:
                fd.write('\n')
            except:
                error = 1
                print('Cannot write to file '+file)
    if error:
        print("""
If you want to create data files for some interesting orbits,
first save the program in your own folder, then run it again.
""")

def mass_to_radius(mass):
    return M_to_R*mass**(1./3.)

def restoreview():
    scene.range = (sw/sh)*xmax
    scene.up = (0,1,0)
    scene.forward = (0,0,-1)

def restore():
    global save
    for nn in range(3):
        stars[nn].pos = vector(save[7*nn+0],save[7*nn+1],save[7*nn+2])
        stars[nn].vel = vector(save[7*nn+3],save[7*nn+4],save[7*nn+5])
        stars[nn].mass = save[7*nn+6]
    stars[0].color = color.red
    stars[1].color = color.blue
    stars[2].color = color.green
    for s in stars:
        s.radius = mass_to_radius(s.mass)
        s.p = s.mass*s.vel
        s.trail.pos = []
        s.trail.color = s.color
        if s.mass > 0:
            s.visible = 1
    pause(state=1)

def showvarrs():
    for nn in range(3):
        varr[nn].pos = stars[nn].pos
        varr[nn].axis = vscale*stars[nn].vel
        varr[nn].visible = 1

def setdefaultsituation():
    global save
    save = [-0.5,0,0,0,0.6,0,2, # star 0: x,y,z,px,py,pz,m
            0.5,0,0,0,-1.2,0,1, # star 1: x,y,z,px,py,pz,m
            -2,0,0,0,1,0,0.1]   # star 2: x,y,z,px,py,pz,m
    restore()

def getsituation():
    fd = get_file('.orb')
    if not fd:
        return
    data = fd.read()
    words = data.split()
    # File format: three sets of 7 numbers;
    # x, y, z, vx, vy, vz, mass
    for nn in range(3):
        base = 7*nn
        stars[nn].pos = vector(float(words[base+0]),float(words[base+1]),float(words[base+2]))
        stars[nn].vel = vector(float(words[base+3]),float(words[base+4]),float(words[base+5]))
        stars[nn].mass = float(words[base+6])
    for s in stars:
        s.trail.pos = []
        s.p = s.mass*s.vel
    showvarrs()
    bget.state = 1
    restoreview()

def savesituation():
    fd = save_file('.orb')
    # File format: three sets of 7 numbers;
    # x, y, z, vx, vy, vz, mass
    for nn in range(3*7):
        fd.write('%g' % save[nn])
        if ((nn+1) % 7) == 0:
            fd.write('\n')
        else:
            fd.write('\t')

def pause(state=None):
    if state != None:
        bpause.state = state
    else:
        bpause.state = not bpause.state
    if bpause.state:
        bpause.text = 'Run'
    else:
        bpause.text = 'Pause'
        
def energy(state=None):
    if state != None:
        benergy.state = state
    else:
        benergy.state = not benergy.state
    if benergy.state:
        benergy.text = 'Hide Energy'
        Kbar.height = 0
        Ubar.height = 0
        Ebar.height = 0
        Kbar.visible = 1
        Ubar.visible = 1
        Ebar.visible = 1
        Klabel.visible = 1
        Ulabel.visible = 1
        Elabel.visible = 1
    else:
        benergy.text = 'Show Energy'
        Kbar.visible = 0
        Ubar.visible = 0
        Ebar.visible = 0
        Klabel.visible = 0
        Ulabel.visible = 0
        Elabel.visible = 0

def reset():
    setdefaultsituation()
    pause(state=1)
    breset.state = 1
    restoreview()

def repeat():
    restore()
    brepeat.state = 1
        
def click():
    # return 1 if click in orbit window
    if scene.mouse.events:
        m = scene.mouse.getevent()
        return m.click
    else:
        return 0

def saveinitialconditions():
    global save
    for nn in range(3):
        save[7*nn] = stars[nn].pos.x
        save[7*nn+1] = stars[nn].pos.y
        save[7*nn+2] = stars[nn].pos.z
        save[7*nn+3] = stars[nn].vel.x
        save[7*nn+4] = stars[nn].vel.y
        save[7*nn+5] = stars[nn].vel.z
        save[7*nn+6] = stars[nn].mass
    
def showenergy():
    K = 0
    U = 0
    for s in stars:
        if s.mass > 0:
            K = K+0.5*(mag(s.p)**2)/s.mass
    for pair in [(0,1), (0,2), (1,2)]:
        i = pair[0]
        j = pair[1]
        if stars[i].mass > 0 and stars[j].mass > 0:
            U = U-G*stars[i].mass*stars[j].mass/mag(stars[i].pos-stars[j].pos)
    Kbar.height = Escale*K
    Kbar.y = Kbar.height/2
    Ubar.height = abs(Escale*U)
    Ubar.y = Escale*U/2
    Ebar.height = abs(Escale*(K+U))
    Ebar.y = Escale*(K+U)/2
                
def orbit():
    dt = 0.01
    for s in stars:
        s.p = s.mass*s.vel
        if s.mass > 0:
            s.visible = 1
    pause(state=0)
    brepeat.state = 0
    saveinitialconditions()
    tclock = clock()
    while 1:
        rate(300)
        ctrl.interact()
        if click() or breset.state or brepeat.state or bget.state:
            pause(state=1)
            return
        if not bpause.state:
            for s1 in stars:  # Find force acting on star s1
                if s1.visible == 0: continue # s1 merged with another star
                F = vector(0,0,0) # We will add up all the forces on s1 
                for s2 in stars:
                    if s2 == s1: continue # All stars but s1 itself
                    if s2.visible == 0: continue # s2 merged with another star
                    r12 = s2.pos-s1.pos
                    if r12.mag <= s1.radius+s2.radius:
                        if s2.mass > s1.mass:
                            s2.color = s1.color
                        s1.mass = s1.mass+s2.mass
                        s1.radius = mass_to_radius(s1.mass)
                        s1.p = s1.p+s2.p
                        s2.visible = 0
                        s2.mass = 0
                        continue
                    else:
                        # Add up all vector forces acting on star s1
                        F = F + (G*s1.mass*s2.mass/mag(r12)**2)*norm(r12)
                s1.p = s1.p+F*dt  # Apply net force to star s1
                
            for s in stars: # After updating momenta, update positions
                if s.visible == 0: continue
                s.pos = s.pos+(s.p/s.mass)*dt
                s.trail.append(pos=s.pos)

            if benergy.state:
                showenergy()

            if clock() > tclock+3:
                instruct.text = ''

#########################################################################

restoreview()
buildfiles()

ctrl = controls(x=scene.x+sw, y=scene.y, width=400, height=200)
bpause = button(pos=(-65,20), width=60, height=30,
             action=lambda: pause())
brepeat = button(pos=(0,20), width=60, height=30, text='Repeat',
             action=lambda: repeat())
benergy = button(pos=(65,20), width=60, height=30,
             action=lambda: energy())
bget = button(pos=(-65,-20), width=60, height=30, text='Get File',
              action=lambda: getsituation())
bsave = button(pos=(0,-20), width=60, height=30, text='Save File',
              action=lambda: savesituation())
breset = button(pos=(65,-20), width=60, height=30, text='Reset',
             action=lambda: reset())
pause(state=1)
bget.state = 0
breset.state = 0
brepeat.state = 0

xi = 0.7*xmax
spacing = 0.1*xmax
w = 0.8*spacing
offset = w/2
Kbar = box(pos=(xi,0,0), size=(w,0,0.01), color=color.magenta, visible=0)
Ubar = box(pos=(xi+spacing,0,0), size=(w,0,0.01), color=color.cyan, visible=0)
Ebar = box(pos=(xi+2*spacing,0,0), size=(w,0,0.01), color=color.yellow, visible=0)
Klabel = label(pos=Kbar.pos+vector(0,-offset,0), text='K', opacity=0, box=0, line=0, visible=0)
Ulabel = label(pos=Ubar.pos+vector(0,offset,0), text='U', opacity=0, box=0, line=0, visible=0)
Elabel = label(pos=Ebar.pos+vector(0,offset,0), text='K+U', opacity=0, box=0, line=0, visible=0)
Escale = 0.7
energy(state=0)

stars = []
stars.append(sphere(visible=0))
stars.append(sphere(visible=0))
stars.append(sphere(visible=0))
for s in stars:
    s.trail = curve()
setdefaultsituation()
ipress = "Press to position green star, then drag to choose initial velocity"
idrag = "Drag to choose initial velocity"
iorbiting = "Click to stop"
inext = "Click to choose new initial conditions"
irun = "Click to run"
instruct = label(pos=(0,0.94*xmax,0), text=ipress, opacity=0, box=0, line=0)

vscale = 0.5 # Scale velocity vectors to graphics window
varr = []
for s in stars:
    varr.append(arrow(pos=s.pos, axis=vscale*s.vel,
                      shaftwidth=0.03, color=s.color, visible=(s.mass > 0)))
vchoose = arrow(pos=(0,0,0), axis=(0,0,0), shaftwidth=0.05, color=color.yellow)
startorbit = 0
pos = None

while 1: # drag to define initial velocity of vstar
    ctrl.interact()
    if not bpause.state:
        pause(state=0)
        startorbit = 1
    if bget.state:
        bget.state = 0
        startorbit = 1
        instruct.text = irun
        while 1:
            if click():
                break
    if brepeat.state:
        brepeat.state = 0
        startorbit = 1
    if startorbit:
        startorbit = 0
        pos = None
        for a in varr:
            a.visible = 0
        vchoose.visible = 0
        instruct.text = iorbiting
        for s in stars:
            s.trail.pos = []
        orbit()
        instruct.text = inext
        while 1:
            ctrl.interact()
            if brepeat.state:
                restore()
                brepeat.state = 0
                startorbit = 1
                break
            if bget.state:
                break
            if click() or breset.state:
                instruct.text = ipress
                if not breset.state:
                    restore()
                showvarrs()
                breset.state = 0
                break
    if pos:
        vchoose.axis = scene.mouse.pos-pos
        # Adjust other stars' momenta so that total p = 0
        pother = vector(0,0,0)
        mother = 0
        for n in range(3):
            if n == vstar:
                p = stars[n].mass*vchoose.axis/vscale
                m = stars[n].mass
            else:
                pother = pother+stars[n].mass*stars[n].vel
                mother = mother+stars[n].mass
        if p.mag > 0:
            votherold = pother/mother
            vothernew = -p/mother # make total momentum be 0
            for n in range(3):
                if n == vstar: continue
                stars[n].vel = (stars[n].vel-votherold)+vothernew
                varr[n].axis = vscale*stars[n].vel
    if scene.mouse.events:
        ev = scene.mouse.getevent()
        if ev.drag:
            if not pos:
                instruct.text = idrag
            pos = ev.pos
            stars[vstar].pos = pos
        elif ev.press == 'left':
            instruct.text = idrag
            vchoose.pos = ev.pos
            vchoose.axis = (0,0,0)
            vchoose.visible = 1
        elif ev.drop or ev.click:
            if ev.drop:
                stars[vstar].vel = vchoose.axis/vscale
            else:
                stars[vstar].pos = ev.pos
                stars[vstar].vel = vector(0,0,0)
                stars[1].vel = -stars[0].mass*stars[0].vel/stars[1].mass
            startorbit = 1