# -*- coding: utf-8 -*-

from Tkinter import *
import tkSimpleDialog
import tkFileDialog
import math
import string
import random

from math import floor

class vec(object):
    def __init__(self,tvec=[0.,0.]):
        self.__vec=tvec[:]

    def get(self):
        return self.__vec

    def set(self,tvec):
        self.__vec=tvec[:]

    vec=property(get,set)

    def getx(self):
        return self.__vec[0]

    def gety(self):
        return self.__vec[1]

    def setx(self,x):
        self.__vec[0]=x

    def sety(self,y):
        self.__vec[1]=y

    x=property(getx,setx)
    y=property(gety,sety)

    def zero(self):
        self.__vec[0]=0.
        self.__vec[1]=0.

    def __iadd__(self,tvec):
        self.__vec[0]+=tvec.__vec[0]
        self.__vec[1]+=tvec.__vec[1]
        return self

    def __isub__(self,tvec):
        self.__vec[0]-=tvec.__vec[0]
        self.__vec[1]-=tvec.__vec[1]
        return self

    def __neg__(self):
        return vec([-self.__vec[0],-self.__vec[1]])
        
    def __add__(self,tvec):
        return vec([self.__vec[0]+tvec.__vec[0],self.__vec[1]+tvec.__vec[1]])

    def __sub__(self,tvec):
        return vec([self.__vec[0]-tvec.__vec[0],self.__vec[1]-tvec.__vec[1]])

    def __div__(self,s):
        return vec([self.__vec[0]/s,self.__vec[1]/s])

    def __mul__(self,s):
        return vec([self.__vec[0]*s,self.__vec[1]*s])

    def __imul__(self,s):
        self.__vec[0]*=s
        self.__vec[1]*=s
        return self

    def len2(self):
        return self.__vec[0]*self.__vec[0]+self.__vec[1]*self.__vec[1]

    def len(self):
        return self.len2()**0.5

class atom(object):
    def __init__(self,x=vec(),v=vec()):
        self.clear()
        self.__x=x # Position
        self.__v=v # Velocity

    def clear(self):
        self.__x=vec() # Position
        self.__v=vec() # Velocity
        self.__f=vec() # Force
        self.__a=[vec(),vec(),vec()] # Higher derivatives for gear

    def get_pos(self):
        return self.__x

    def set_pos(self,x):
        self.__x=x

    pos=property(get_pos,set_pos)

    def get_vel(self):
        return self.__v

    def set_vel(self,v):
        self.__v=v

    vel=property(get_vel,set_vel)

    def get_force(self):
        return self.__f
    
    def set_force(self,f):
        self.__f=f

    force=property(get_force,set_force)

    def get_higher_deriv(self):
        return self.__a

    higher_deriv=property(get_higher_deriv)

class chemsystem(object):
    def __init__(self,boxlen,atoms,timestep):
        self.atoms=atoms[:]
        self.potential_energy=0.
        self.kinetic_energy=0.
        self.timestep=timestep
        self.__boxlen=boxlen
        self.__cutoff2=(2.**(1./6))**2 # FOR WCA

    # Compute energy and force using the WCA potential
    def compute_force(self):
        natoms=len(self.atoms) # Number of atoms
        self.potential_energy=0. # Initial potential energy
        for atom in self.atoms: # Clear all forces on all atoms
            atom.force.zero()
        # Use a simple pair loop
        for i in xrange(natoms):
            for j in xrange(i+1,natoms):
                dv=self.atoms[i].pos-self.atoms[j].pos # Distance vector
                dv-=vec([self.__boxlen*floor(0.5+dv.x/self.__boxlen), # Min. im.
                        self.__boxlen*floor(0.5+dv.y/self.__boxlen)])
                d2=dv.len2() # Distance squared
                if d2<self.__cutoff2: # Spherical cutoff
                    self.potential_energy+=4*(1./d2**6-1./d2**3)+1
                    force_divided_by_r=4*(12./d2**7-6./d2**4)
                    force=dv*force_divided_by_r
                    self.atoms[i].force+=force # Accumulate forces on atoms
                    self.atoms[j].force-=force

    def compute_kinetic_energy(self):
        self.kinetic_energy=0.
        for atom in self.atoms:
            self.kinetic_energy+=atom.vel.len2()
        self.kinetic_energy*=0.5

    def remove_linear_momentum(self):
        m=vec([0,0])
        for atom in self.atoms:
            m+=atom.vel
        m/=len(self.atoms)
        for atom in self.atoms:
            atom.vel-=m

    def scale_ke(self,target):
        self.compute_kinetic_energy()
        g=(target/self.kinetic_energy)**0.5
        for atom in self.atoms:
            atom.vel=atom.vel*g

    def velocity_verlet(self): # Integrate equations of motion
        ts=self.timestep
        # First find the new positions and first half of velocities
        for atom in self.atoms:
            atom.pos+=atom.vel*ts+atom.force*(ts**2*0.5)
            atom.vel+=atom.force*ts*0.5
            # Make sure the position of the atom is within the central box
            atom.pos-=vec([self.__boxlen*floor(atom.pos.x/self.__boxlen),
                                 self.__boxlen*floor(atom.pos.y/self.__boxlen)])
        # Evaluate the force
        self.compute_force()
        # Complete the velocity
        for atom in self.atoms:
            atom.vel+=atom.force*ts*0.5
        self.compute_kinetic_energy()

    def sprk4(self): # Integrate equations of motion
        ts=self.timestep
        B=[]
        B.append( 0.0792036964311957)
        B.append( 0.353172906049774 )
        B.append( -0.0420650803577195 )
        B.append( 1.-2*(B[0]+B[1]+B[2]))
        B.append( B[2])
        B.append( B[1])
        B.append( B[0])
        b=[]
        b.append( 0.209515106613362)
        b.append( -0.143851773179818)
        b.append( 1./2-(b[0]+b[1]))
        b.append( b[2])
        b.append( b[1])
        b.append( b[0])
        b.append(0.)

        for stage in range(7):            
            for atom in self.atoms:
                atom.pos+=atom.vel*ts*B[stage]
                atom.pos-=vec([self.__boxlen*floor(atom.pos.x/self.__boxlen),
                               self.__boxlen*floor(atom.pos.y/self.__boxlen)])
            # Evaluate the force
            self.compute_force()
            for atom in self.atoms:
                atom.vel+=atom.force*ts*b[stage]
        self.compute_kinetic_energy()

    def sprk6(self): # Integrate equations of motion
        ts=self.timestep
        B=[]
        B.append(0.0502627644003922 )
        B.append(0.413514300428344 )
        B.append(0.0450798897943977 )
        B.append(-0.188054853819569 )
        B.append(0.541960678450780 )
        B.append( 1.-2*(B[0]+B[1]+B[2]+B[3]+B[4]))
        B.append( B[4])
        B.append( B[3])
        B.append( B[2])
        B.append( B[1])
        B.append( B[0])
        b=[]
        b.append(0.148816447901042 )
        b.append(-0.132385865767784 )
        b.append(0.067307604692185 )
        b.append(0.432666402578175 )
        b.append( 1./2-(b[0]+b[1]+b[2]+b[3]))
        b.append( b[4])
        b.append( b[3])
        b.append( b[2])
        b.append( b[1])
        b.append( b[0])
        b.append(0.)

        for stage in range(11):
            for atom in self.atoms:
                atom.pos+=atom.vel*ts*B[stage]
                atom.pos-=vec([self.__boxlen*floor(atom.pos.x/self.__boxlen),
                               self.__boxlen*floor(atom.pos.y/self.__boxlen)])
            # Evaluate the force
            self.compute_force()
            for atom in self.atoms:
                atom.vel+=atom.force*ts*b[stage]
        self.compute_kinetic_energy()

    def gear(self): # Integrate equations of motion
        ts=self.timestep
        gear_constants=(19./120.,3./4., 1., 1./2.,1./12.);
        # First predict
        for atom in self.atoms:
            atom.pos+=(atom.vel*ts+
                       atom.higher_deriv[0]*ts**2/2.+
                       atom.higher_deriv[1]*ts**3/6.+
                       atom.higher_deriv[2]*ts**4/24.)
            atom.vel+=(atom.higher_deriv[0]*ts+
                       atom.higher_deriv[1]*ts**2/2.+
                       atom.higher_deriv[2]*ts**3/6.)
            atom.higher_deriv[0]+=(atom.higher_deriv[1]*ts+
                                   atom.higher_deriv[2]*ts**2/2.)
            atom.higher_deriv[1]+=atom.higher_deriv[2]*ts
            # Make sure the position of the atom is within the central box
            atom.pos-=vec([self.__boxlen*floor(atom.pos.x/self.__boxlen),
                           self.__boxlen*floor(atom.pos.y/self.__boxlen)])
        # Evaluate the force
        self.compute_force()
        # Then Correct
        for atom in self.atoms:
            corr=atom.force-atom.higher_deriv[0]
            atom.pos+=corr*gear_constants[0]*ts**2/2.
            atom.vel+=corr*gear_constants[1]*ts/2.
            atom.higher_deriv[0].vec=atom.force.vec
            atom.higher_deriv[1]+=corr*gear_constants[3]*3./ts
            atom.higher_deriv[2]+=corr*gear_constants[4]*12./ts**2
        self.compute_kinetic_energy()

class AtomDraw(Frame):
    def __init__(self,master=None,size=300):
        Frame.__init__(self,master)
        self.pack(expand=1,fill=BOTH)

        cntr=Frame(self,borderwidth=2,relief=RAISED)
        cntr.pack()

        self.size=size

        self.canvas = Canvas(cntr,
            bd=0, highlightthickness=0,
            width=size,
            height=size)

        self.canvas.pack(expand=1, fill=BOTH, side=TOP)

        # Reset the atomdraw canvas
        self.canvas.xview("moveto", 0)
        self.canvas.yview("moveto", 0)

        self.__themark=None

        self.__markparticle=0
        
        self.current_ovals=None

        self.trajcol=["red","green","yellow"]

        self.trajectory=[[],[],[]]
        self.trajobj=[[],[],[]]

    def set_mark(self,x):
        self.__markparticle=x
        if self.__themark:
            if self.__markparticle:
                self.canvas.itemconfigure(self.__themark,fill="red")
            else:
                self.canvas.itemconfigure(self.__themark,fill="blue")

    def draw_atoms(self,atoms,boxlen,amark):
        cf=self.size/boxlen
        rad=cf*0.5
        if amark:
            acheck=atoms[amark]
        else:
            acheck=None
        if not self.current_ovals:
            self.current_ovals=[]
            for atom in atoms:
                xy=atom.pos.x*cf-rad,atom.pos.y*cf-rad,atom.pos.x*cf+rad,atom.pos.y*cf+rad
                oval=self.canvas.create_oval(xy,fill="blue")
                if atom is acheck:
                    self.__themark=oval
                self.current_ovals.append(oval)
        else:
            for i in xrange(len(self.current_ovals)):
                x=self.current_ovals[i]
                atom=atoms[i]
                xy=atom.pos.x*cf-rad,atom.pos.y*cf-rad,atom.pos.x*cf+rad,atom.pos.y*cf+rad
                self.canvas.coords(x,xy)
                if x is self.__themark:
                    if self.__markparticle:
                        self.canvas.itemconfigure(self.__themark,fill="red")
                    else:
                        self.canvas.itemconfigure(self.__themark,fill="blue")

    def set_trajectory(self,which,datapoints,boxlen):
        blhalf=boxlen/2
        cf=self.size/boxlen
        self.clear_trajectory(which)
        for datapoint in datapoints:
            self.trajectory[which].append(datapoint)
        if len(self.trajectory[which])>1:
            for i in xrange(len(self.trajectory[which])-1):
                x1=self.trajectory[which][i][0]
                y1=self.trajectory[which][i][1]
                x2=self.trajectory[which][i+1][0]
                y2=self.trajectory[which][i+1][1]
                if abs(x1-x2)<blhalf and abs(y1-y2)<blhalf:
                    self.trajobj[which].append(self.canvas.create_line(x1*cf,y1*cf,x2*cf,y2*cf,fill=self.trajcol[which]))

    def add_trajectory(self,which,datapoint,boxlen):
        blhalf=boxlen/2
        cf=self.size/boxlen
        self.trajectory[which].append(datapoint)
        if len(self.trajectory[which])>1:
            i=len(self.trajectory[which])-2
            x1=self.trajectory[which][i][0]
            y1=self.trajectory[which][i][1]
            x2=self.trajectory[which][i+1][0]
            y2=self.trajectory[which][i+1][1]
            if abs(x1-x2)<blhalf and abs(y1-y2)<blhalf:
                self.trajobj[which].append(self.canvas.create_line(x1*cf,y1*cf,x2*cf,y2*cf,fill=self.trajcol[which]))

    def clear_trajectory(self,which):
        for x in self.trajobj[which]:
            self.canvas.delete(x)
        self.trajobj[which]=[]        
        self.trajectory[which]=[]

# Figure out how to make graph boundaries (hackish)
def ground(lo,hi,maxticks=10):
    lo=float(lo)
    hi=float(hi)
    if hi-lo < 2e-6:
        hi+=1e-6
        lo-=1e-6
    
    dist=hi-lo
    p=math.floor(math.log(dist)/math.log(maxticks))
    e=maxticks**p

    xlo=math.floor(lo/e)*e
    xhi=math.ceil(hi/e)*e
    
    nstep=int((xhi-xlo)/e)+1

    dec=0
    if p<0:
        dec=-int(p)
    x=max(abs(xlo),abs(xhi))
    oval=int(math.ceil(math.log10(x)))+2+dec

    steplist=[]
    x=xlo
    for i in range(nstep):
        fstring='%'+`oval`+'.'+`dec`+'f'
        steplist.append(fstring % x)
        x+=e

    return xlo,xhi,steplist

class Graph(Frame):
    def __init__(self,master=None,width=300,height=100,title=''):
        Frame.__init__(self,master)
        self.pack(expand=1,fill=BOTH)
        
        self.width=width
        self.height=height
        self.title=title
        
        self.canvas = Canvas(self,
            bd=0, highlightthickness=0,
            width=self.width,
            height=self.height)

        self.canvas.pack(expand=1, fill=BOTH, side=LEFT)

        # Reset the canvas
        self.canvas.xview("moveto", 0)
        self.canvas.yview("moveto", 0)

        self.__objects=[]

    def graph(self,datasets):
        # Destroy old data
        for x in self.__objects:
            self.canvas.delete(x)
        self.__objects=[]

        gxmin=80
        gymin=20
        gxmax=self.width-20
        gymax=self.height-20
        shiftlen=5

        # Draw lines around graph
        self.__objects.append(self.canvas.create_line(gxmin,gymin,gxmin,gymax))
        self.__objects.append(self.canvas.create_line(gxmin,gymax,gxmax,gymax))
        self.__objects.append(self.canvas.create_line(gxmax,gymax,gxmax,gymin))
        self.__objects.append(self.canvas.create_line(gxmax,gymin,gxmin,gymin))

        # Draw title
        self.__objects.append(self.canvas.create_text((gxmax+gxmin)/2,gymin,text=self.title,anchor='s'))

        showg=0
        dp=None
        for datapairs in datasets:
            if len(datapairs)>2:
                showg=1
                dp=datapairs

        if showg:
            maxx=dp[0][0]
            minx=dp[0][0]
            maxy=dp[0][1]
            miny=dp[0][1]
            for datapairs in datasets:
                for x in datapairs:
                    maxx=max(maxx,x[0])
                    minx=min(minx,x[0])
                    maxy=max(maxy,x[1])
                    miny=min(miny,x[1])

            # Compute ticks/labels

            (minx,maxx,stepsx)=ground(minx,maxx,5)
            (miny,maxy,stepsy)=ground(miny,maxy,6)

            # Draw ticks/labels

            nstepsx=len(stepsx)
            nstepsy=len(stepsy)

            ticklen=5

            cnv=(gxmax-gxmin)/(nstepsx-1)
            for i in xrange(nstepsx):
                a='n'
                #if i==0:
                #    a='nw'
                #if i==nstepsx-1:
                #    a='ne'
                self.__objects.append(self.canvas.create_text(gxmin+i*cnv,
                                                              gymax+shiftlen,text=stepsx[i],anchor=a))
                self.__objects.append(self.canvas.create_line(gxmin+i*cnv,gymax,gxmin+i*cnv,gymax-ticklen))

            cnv=(gymax-gymin)/(nstepsy-1)
            for i in xrange(nstepsy):
                a='e'
                #if i==0:
                #    a='se'
                #if i==nstepsy-1:
                #    a='ne'
                self.__objects.append(self.canvas.create_text(gxmin-shiftlen,
                                                              gymax-i*cnv,
                                                              text=stepsy[i],anchor=a))
                self.__objects.append(self.canvas.create_line(gxmin,gymax-i*cnv,gxmin+ticklen,gymax-i*cnv))



            xscale=(gxmax-gxmin)/(maxx-minx)
            yscale=(gymax-gymin)/(maxy-miny)

            colors="red","green","yellow"
            for j in xrange(len(datasets)):
                k=len(datasets)-j-1
                datapairs=datasets[k]
                if len(datapairs)>2:
                    xy=[]
                    for i in xrange(len(datapairs)-1):
                        xy.extend([gxmin+(datapairs[i][0]-minx)*xscale,
                                  gymax-(datapairs[i][1]-miny)*yscale,
                                  gxmin+(datapairs[i+1][0]-minx)*xscale,
                                  gymax-(datapairs[i+1][1]-miny)*yscale])
                    self.__objects.append(self.canvas.create_line(xy,fill=colors[k]))


class App(Frame):
    def __init__(self,master=None):
        Frame.__init__(self,master)
        self.pack(expand=1,fill=BOTH)

        # Some defaults
        # Size of the system (iatom**2 atoms)
        iatom=5
        # Physical size of the system
        self.boxlen=1.1*iatom

        # Timestep
        self.timestep=5e-3
        self.targetke=10.
        self.ke_scaling=IntVar()
        self.ke_scaling.set(0)
        self.__nkescale=10

        self.__mp=IntVar()
        self.__mp.set(0)

        self.__traj=IntVar()
        self.__traj.set(0)

        self.__show_comparison=IntVar()
        self.__show_comparison.set(0)

        self.__integrator_choice=StringVar()


        self.maincontainer=Frame(self,borderwidth=2,relief=GROOVE)
        self.maincontainer.pack(expand=1,fill=BOTH,side=TOP)

        self.buttons=Frame(self)
        self.buttons.pack(side=BOTTOM,anchor="w")
        
        self.framescont=Frame(self.maincontainer)
        self.framescont.pack(side=LEFT,anchor="n")
        self.framescont2=Frame(self.maincontainer)
        self.framescont2.pack(side=LEFT,anchor="n")
        
        self.atomdraw=AtomDraw(self.framescont)
        self.atomdraw.pack(side=TOP,anchor="n")
        #self.atomdraw=AtomDraw(Tk())
        #self.atomdraw.master.title("Disks")

        self.graphframe=Frame(self.framescont2)
        self.graphframe.pack(side=TOP)

        self.container=Frame(self.framescont2)
        self.container.pack(expand=1,fill=BOTH,side=TOP)

        self.tools=Frame(self.framescont)
        self.tools.pack(side=TOP)
        self.toolsleft=Frame(self.tools)
        self.toolsleft.pack(side=LEFT)
        self.toolsright=Frame(self.tools)
        self.toolsright.pack(side=LEFT)

        self.tools2=Frame(self.framescont,borderwidth=2,relief=GROOVE)
        self.tools2.pack(side=LEFT,expand=1,fill=BOTH,anchor="n")
        self.tools2left=Frame(self.tools2)
        self.tools2left.pack(expand=1,fill=BOTH,side=LEFT,anchor="sw")
        self.tools2right=Frame(self.tools2)
        self.tools2right.pack(expand=1,fill=BOTH,side=LEFT,anchor="se")
        
        #self.graphframe=Frame(Tk())
        #self.graphframe.pack()
        #self.graphframe.master.title("Energies")

        self.ke_graph=Graph(self.graphframe,width=400,height=150,title="Kinetic Energy")
        self.ke_graph.pack(side=TOP)

        self.pe_graph=Graph(self.graphframe,width=400,height=150,title="Potential Energy")
        self.pe_graph.pack(side=TOP)

        self.te_graph=Graph(self.graphframe,width=400,height=150,title="Total Energy")
        self.te_graph.pack(side=TOP)

        self.container.left=Frame(self.container)
        self.container.left.pack(side=LEFT,expand=1,fill=BOTH)
        self.container.right=Frame(self.container)
        self.container.right.pack(side=RIGHT,expand=1,fill=BOTH)

        Label(self.container.left,text="Time").pack(anchor="w")
        self.container.ltime=Label(self.container.right,text="")
        self.container.ltime.pack(anchor="e",expand=0,fill="x")

        Label(self.container.left,text="Potential energy").pack(anchor="w")
        self.container.lpe=Label(self.container.right,text="")
        self.container.lpe.pack(anchor="e",expand=0,fill="x")

        Label(self.container.left,text="Kinetic energy").pack(anchor="w")
        self.container.lke=Label(self.container.right,text="")
        self.container.lke.pack(anchor="e",expand=0,fill="x")

        Label(self.container.left,text="Total energy").pack(anchor="w")
        self.container.lte=Label(self.container.right,text="")
        self.container.lte.pack(anchor="e",expand=0,fill="x")

        mb=Menubutton(self.toolsleft,textvariable=self.__integrator_choice,relief=RAISED)
        mb.menu=Menu(mb,tearoff=0)
        mb["menu"]=mb.menu
        mb.menu.add_radiobutton(label="Velocity Verlet",variable=self.__integrator_choice,indicatoron=0)
        mb.menu.add_radiobutton(label="Gear PC",variable=self.__integrator_choice,indicatoron=0)
        mb.menu.add_radiobutton(label="SPRK4",variable=self.__integrator_choice,indicatoron=0)
        mb.menu.add_radiobutton(label="SPRK6",variable=self.__integrator_choice,indicatoron=0)
        self.__integrator_choice.set("Velocity Verlet")
        mb.pack(anchor="e",expand=0,fill="x")

        self.container.bts=Button(self.toolsleft,text="",command=self.__modify_timestep)
        self.container.bts.pack(anchor="e",expand=0,fill="x")

        self.container.btke=Button(self.toolsleft,text="",command=self.__modify_tke)
        self.container.btke.pack(anchor="e",expand=0,fill="x")

        Checkbutton(self.toolsright,text="Scale kinetic energy",variable=self.ke_scaling).pack(side=TOP,anchor="w")
        Checkbutton(self.toolsright,text="Mark particle",variable=self.__mp,command=self.__updatemp).pack(side=TOP,anchor="w") 
        Checkbutton(self.toolsright,text="Show trajectory",variable=self.__traj,command=self.__updatetraj).pack(side=TOP,anchor="w") 
        Checkbutton(self.toolsright,text="Show comparison/full",variable=self.__show_comparison,command=self.__updatecompare).pack(side=TOP,anchor="w") 

        self.quitbutton=Button(self.tools2left,text="Run",command=self.__runit)
        self.quitbutton.pack(side=TOP,fill="x")
        self.quitbutton=Button(self.tools2left,text="Stop",command=self.__stopit)
        self.quitbutton.pack(side=TOP,fill="x")
        self.quitbutton=Button(self,text="Quit",command=self.quit)
        self.quitbutton.pack(side=LEFT,anchor="w")
        Label(self,text="Get Demosim from http://www.spaangberg.se/daniels/demosim/").pack(side=LEFT,fill="x",expand=1)

        Button(self.tools2right,text="Clear/Set Mark",command=self.__clear).pack(fill="x")
        Button(self.tools2right,text="Save Since Marked",command=self.__save).pack(fill="x")
        Button(self.tools2right,text="Open",command=self.__load).pack(fill="x")
        Button(self.tools2right,text="Reload",command=self.__reload).pack(fill="x")

        # Set up iatom x iatom list in a regular "lattice"
        myatoms=[]
        for i in xrange(iatom):
            for j in xrange(iatom):
                myatoms.append(atom(vec([i*self.boxlen/iatom,j*self.boxlen/iatom]), # POS
                                    vec([10*(random.random()-0.5),10*(random.random()-0.5)]))) #VEL

        self.__amark=iatom*iatom/2

        # Setup the system
        self.__mysystem=chemsystem(self.boxlen,myatoms,self.timestep)

        # Compute initial energy/forces
        self.__mysystem.compute_force()

        self.__curstep=0
        self.__curtime=0
        self.__stopped=0

        self.__atomdraw_step=2
        self.__graphdraw_step=10
        self.__graphstore_step=2
        self.__maxgraphlen=200

        self.__graphdata_lpe=[]
        self.__graphdata_lke=[]
        self.__graphdata_lte=[]

        self.__save_graphdata_lpe=[]
        self.__save_graphdata_lke=[]
        self.__save_graphdata_lte=[]

        self.__save_posvel=[]
        self.__marktraj=[]

        self.__previous_file=None

        self.__compare_traj=[]
        self.__compare_graphdata_lpe=[]
        self.__compare_graphdata_lke=[]
        self.__compare_graphdata_lte=[]
        
        self.__update_values()
        self.__update_btke()
        self.__update_bts()
        self.__draw_atoms()
        self.__update_graphs()

        self.__clear()

    def __updatemp(self):
        self.atomdraw.set_mark(self.__mp.get())

    def __modify_timestep(self):
        rval=tkSimpleDialog.askfloat(title="Modify timestep",prompt="timestep",initialvalue=self.timestep)
        if rval:
            self.timestep=rval
            self.__mysystem.timestep=self.timestep
            self.__update_bts()

    def __update_bts(self):
        self.container.bts.config(text="Timestep: %12.4g" % self.timestep)

        
    def __modify_tke(self):
        rval=tkSimpleDialog.askfloat(title="New target kinetic energy",prompt="target kinetic energy",initialvalue=self.targetke)
        if rval:
            self.targetke=rval
            self.__update_btke()

    def __update_btke(self):
        self.container.btke.config(text="Target KE: %12.4f" % self.targetke)

    def __store_graphdata(self):
        if self.__curstep % self.__graphstore_step == 0:
            self.__graphdata_lpe.append((self.__curtime,self.__mysystem.potential_energy))
            self.__graphdata_lke.append((self.__curtime,self.__mysystem.kinetic_energy))
            self.__graphdata_lte.append((self.__curtime,self.__mysystem.potential_energy+self.__mysystem.kinetic_energy))
            self.__save_graphdata_lpe.append((self.__curtime,self.__mysystem.potential_energy))
            self.__save_graphdata_lke.append((self.__curtime,self.__mysystem.kinetic_energy))
            self.__save_graphdata_lte.append((self.__curtime,self.__mysystem.potential_energy+self.__mysystem.kinetic_energy))
            if (len(self.__graphdata_lpe)>(self.__maxgraphlen/self.__graphstore_step)):
                self.__graphdata_lpe=self.__graphdata_lpe[2:]
            if (len(self.__graphdata_lke)>(self.__maxgraphlen/self.__graphstore_step)):
                self.__graphdata_lke=self.__graphdata_lke[2:]
            if (len(self.__graphdata_lte)>(self.__maxgraphlen/self.__graphstore_step)):
                self.__graphdata_lte=self.__graphdata_lte[2:]

    def __coarse(self,data,gran):
        l=len(data)/gran
        newdata=data
        if l>1:
            newdata=[]
            for i in range(len(data)/l):
                newdata.append(data[i*l])
        return newdata

    def __update_graphs(self):
        if self.__show_comparison.get():
            self.te_graph.graph([self.__coarse(self.__save_graphdata_lte,100),self.__compare_graphdata_lte])
            self.ke_graph.graph([self.__coarse(self.__save_graphdata_lke,100),self.__compare_graphdata_lke])
            self.pe_graph.graph([self.__coarse(self.__save_graphdata_lpe,100),self.__compare_graphdata_lpe])            
        else:
            self.te_graph.graph([self.__graphdata_lte])
            self.ke_graph.graph([self.__graphdata_lke])
            self.pe_graph.graph([self.__graphdata_lpe])

    def __update_values(self):
        self.container.ltime.config(text="%12.4f" % self.__curtime)
        self.container.lpe.config(text="%12.4f" % self.__mysystem.potential_energy)
        self.container.lke.config(text="%12.4f" % self.__mysystem.kinetic_energy)
        self.container.lte.config(text="%12.4f" % (self.__mysystem.potential_energy+self.__mysystem.kinetic_energy))
        
    def __draw_atoms(self):
        self.atomdraw.draw_atoms(self.__mysystem.atoms,self.boxlen,self.__amark)

    def __store_marktraj(self):
        p=self.__mysystem.atoms[self.__amark].pos
        self.__marktraj.append((p.x,p.y))
        if self.__traj.get():
            self.atomdraw.add_trajectory(0,(p.x,p.y),self.boxlen)

    def __updatetraj(self):
        if self.__traj.get():
            if self.__show_comparison.get():
                self.atomdraw.set_trajectory(1,self.__coarse(self.__compare_traj,400),self.boxlen)
            self.atomdraw.set_trajectory(0,self.__marktraj,self.boxlen)
        else:
            self.atomdraw.clear_trajectory(0)
            self.atomdraw.clear_trajectory(1)

    def __updatecompare(self):
        if self.__show_comparison.get():
            if self.__traj.get():
                self.atomdraw.set_trajectory(1,self.__compare_traj,self.boxlen)
            else:
                self.atomdraw.clear_trajectory(1)
        else:
            self.atomdraw.clear_trajectory(1)
        self.__update_graphs()

        

    # Clear trajectory.
    # Clear energy data
    # Store current position of all atoms.
    # Store the current time and timestep and other parameters
    def __clear(self):
        self.__marktraj=[]
        self.atomdraw.clear_trajectory(0)
        self.__graphdata_lpe=[]
        self.__graphdata_lke=[]
        self.__graphdata_lte=[]
        self.__save_graphdata_lpe=[]
        self.__save_graphdata_lke=[]
        self.__save_graphdata_lte=[]
        self.__update_graphs()

        self.__save_posvel=[]
        for atom in self.__mysystem.atoms:
            self.__save_posvel.append((atom.pos.x,
                                       atom.pos.y,
                                       atom.vel.x,
                                       atom.vel.y,
                                       atom.higher_deriv[0].x,
                                       atom.higher_deriv[0].y,
                                       atom.higher_deriv[1].x,
                                       atom.higher_deriv[1].y,
                                       atom.higher_deriv[2].x,
                                       atom.higher_deriv[2].y))
        self.__save_curstep=self.__curstep
        self.__save_curtime=self.__curtime
        self.__save_timestep=self.timestep
        self.__save_targetke=self.targetke
        self.__save_ke_scaling=self.ke_scaling.get()
        self.__save_mp=self.__mp.get()
        self.__save_traj=self.__traj.get()

    def __save(self):
        rval=tkFileDialog.asksaveasfilename(title="Save",filetypes=[("Demosim file",".demosim")])
        if rval:
            f=open(rval,"w")
            f.write(`len(self.__marktraj)`+"\n")
            for x in self.__marktraj:
                f.write(`x[0]`+" "+`x[1]`+"\n")
            f.write(`len(self.__save_graphdata_lpe)`+"\n")
            for i in xrange(len(self.__save_graphdata_lpe)):
                f.write(`self.__save_graphdata_lpe[i][0]`+" "+
                        `self.__save_graphdata_lpe[i][1]`+" "+
                        `self.__save_graphdata_lke[i][1]`+" "+
                        `self.__save_graphdata_lte[i][1]`+"\n")
            f.write(`len(self.__save_posvel)`+"\n")
            for x in self.__save_posvel:
                f.write(("%18.14f" % x[0])+" "+
                        ("%18.14f" % x[1])+" "+
                        ("%18.14f" % x[2])+" "+
                        ("%18.14f" % x[3])+" "+
                        ("%18.14f" % x[4])+" "+
                        ("%18.14f" % x[5])+" "+
                        ("%18.14f" % x[6])+" "+
                        ("%18.14f" % x[7])+" "+
                        ("%18.14f" % x[8])+" "+
                        ("%18.14f" % x[9])+"\n")
            f.write(`self.__save_curstep`+"\n")
            f.write(`self.__save_curtime`+"\n")
            f.write(`self.__save_timestep`+"\n")
            f.write(`self.__save_targetke`+"\n")
            f.write(`self.__save_ke_scaling`+"\n")
            f.write(`self.__save_mp`+"\n")
            f.write(`self.__save_traj`+"\n")
            f.write(self.__integrator_choice.get()+"\n")
            f.close()


    def __load(self):
        rval=tkFileDialog.askopenfilename(title="Open",filetypes=[("Demosim file",".demosim")])
        if rval:
            self.__previous_file=rval
            self.__reload()

    def __reload(self):
        if self.__previous_file:
            f=open(self.__previous_file,"r")
            self.__compare_traj=[]
            l=int(f.readline())
            for i in xrange(l):
                s=string.split(f.readline())
                self.__compare_traj.append((float(s[0]),float(s[1])))
            self.__compare_graphdata_lpe=[]
            self.__compare_graphdata_lke=[]
            self.__compare_graphdata_lte=[]
            l=int(f.readline())
            for i in xrange(l):
                s=string.split(f.readline())
                self.__compare_graphdata_lpe.append((float(s[0]),float(s[1])))
                self.__compare_graphdata_lke.append((float(s[0]),float(s[2])))
                self.__compare_graphdata_lte.append((float(s[0]),float(s[3])))
            self.__compare_graphdata_lpe=self.__coarse(self.__compare_graphdata_lpe,100)
            self.__compare_graphdata_lke=self.__coarse(self.__compare_graphdata_lke,100)
            self.__compare_graphdata_lte=self.__coarse(self.__compare_graphdata_lte,100)
            l=int(f.readline())
            a=self.__mysystem.atoms
            for i in xrange(l):
                s=string.split(f.readline())
                a[i].clear()
                a[i].set_pos(vec([float(s[0]),float(s[1])]))
                a[i].set_vel(vec([float(s[2]),float(s[3])]))
                a[i].higher_deriv[0]=vec([float(s[4]),float(s[5])])
                a[i].higher_deriv[1]=vec([float(s[6]),float(s[7])])
                a[i].higher_deriv[2]=vec([float(s[8]),float(s[9])])
            self.__curstep=int(f.readline())
            self.__curtime=float(f.readline())
            self.timestep=float(f.readline())
            self.__mysystem.timestep=self.timestep
            self.targetke=float(f.readline())
            self.ke_scaling.set(int(f.readline()))
            self.__mp.set(int(f.readline()))
            self.__traj.set(int(f.readline()))
            l=f.readline()
            self.__integrator_choice.set(l[:len(l)-1])
            f.close()
            self.__mysystem.compute_force()
            self.__clear()
            self.__update_bts()
            self.__update_btke()
            self.__updatetraj()
            self.__draw_atoms()

    def __stopit(self):
        self.__stopped=1

    def __runit(self):
        self.__stopped=0
        self.__run()
        
    def __run(self):
        self.__curstep+=1
        self.__curtime+=self.timestep
        # Simulate the system

        if self.__integrator_choice.get()=="Velocity Verlet":
            self.__mysystem.velocity_verlet()
        elif self.__integrator_choice.get()=="SPRK4":
            self.__mysystem.sprk4()
        elif self.__integrator_choice.get()=="SPRK6":
            self.__mysystem.sprk6()
        else:
            self.__mysystem.gear()
        
        self.__store_marktraj()
        if self.ke_scaling.get():
            if self.__curstep%self.__nkescale==0:
                self.__mysystem.remove_linear_momentum()
                self.__mysystem.scale_ke(self.targetke)
        self.__update_values()
        self.__store_graphdata()
        if self.__curstep%self.__atomdraw_step==0:
            self.__draw_atoms()
        if self.__curstep%self.__graphdraw_step==0:
            self.__update_graphs()
        if not self.__stopped:
            self.after(1,func=self.__run)

root=Tk()
app=App(root)
app.master.title("Demosim")
app.mainloop()
