IN_IDLE = 0

# motion of coupled oscillators on a string
# starting displacement and fixed/free ends
# can be selected
# total energy displayed can be used as integration error control

### 03/2006
### Wolfgang.Urban@schule.at

from Tkinter import *
from math import *

STEPSIZE = pi/100.0

class Vibrate:

    ###### Beschleunigungen berechnen
    def f(self,t,y):
        n = self.n
        D = self.D
        a = [0.0]*n
        for i in range(1,n-1):
            a[i] = D*(y[i-1]-2*y[i]+y[i+1])
        if self.free_left: a[0] = D*(y[1]-y[0])
        if self.free_right: a[n-1] = D*(y[n-2]-y[n-1])
        return a
        
    ##########################################################################

    def __init__(self,n):
        self.n = n
        self.root = Tk()
	self.root.title("HIB Oscillations")
        self.cmdframe(self.root)
        self.sizex,self.sizey = 630,int(630/(2*pi/2)*2.6)
        self.c = Canvas(self.root,width=self.sizex+10,height=self.sizey+10)
        self.c.pack(side=LEFT)

        f = Frame(self.root,border=3,relief = RAISED)
        Label(f,text=" ",font=("Arial",5)).pack(side=TOP)
        self.nframe(f)
        Label(f,text=" ",font=("Arial",4)).pack(side=TOP)
        self.endframe(f)
        Label(f,text=" ",font=("Arial",4)).pack(side=TOP)
        self.waveframe(f)
        Label(f,text=" ",font=("Arial",5)).pack(side=TOP)
        f.pack()

        self.n = 0
        self.newwave()

        self.BUSY = False
        if not IN_IDLE: self.root.mainloop()

    # prepare canvas with n oscillators at y=0
    def makeoscillators(self,n):
        self.BUSY = True
        dx = self.sizex/(n-1)      # so viele Distanzen haben Platz
        self.x0 = 6+(self.sizex-dx*(n-1))/2 # Start x
        self.y0 = self.sizey/2+5            # Mitte y

        h = dx*(n-1)/(2*pi)    # height/width = 2pi
        self.sx = dx
        self.sy = h
                
        self.c.delete("all")
        self.c.create_line(self.x0,self.y0,self.x0+(n-1)*self.sx,self.y0,fill="#0000ff")
        self.c.create_line(self.x0,self.y0-self.sy,self.x0+(n-1)*self.sx,self.y0-self.sy,fill="#0000ff")
        self.c.create_line(self.x0,self.y0+self.sy,self.x0+(n-1)*self.sx,self.y0+self.sy,fill="#0000ff")
        self.osc = []
        self.size = s = 2
        for i in range(n):
            x = self.x0+i*self.sx
            y = self.y0-0*self.sy
            r = self.c.create_rectangle(x-s,y-s,x+s,y+s,outline="#000000")
            self.osc.append(r)
        self.c.update()
        self.BUSY = False

    def refresh(self):
        s = self.size
        for i in range(self.n):
            x = self.x0+i*self.sx
            y = self.y0-self.s[i]*self.sy
            self.c.coords(self.osc[i],x-s,y-s,x+s,y+s)
        self.c.update()

    # show moving oscillators
    def play(self):
        self.RUNNING = True
        s = self.s
        v = self.v
        t = 0.0
        h = STEPSIZE
        while self.RUNNING:
            # calculate next positions
            for i in range(3): 
                t,s,v = rkn4_0(t,s,v,self.f,h)
                if not self.RUNNING: return
            # and show it
            self.s = s
            self.refresh()
            ekin,epot,e = energy(s,v,self.D)
            self.root.title("Ekin + Epot = %6.4f + %6.4f = %6.4f "%(ekin,epot,e))

    ##########################################################################
    ## GUI            
    ## (not using lambdas) 
    ##########################################################################

    def cmdframe(self,parent):
        cf = Frame(parent,relief=RIDGE,bd=3)
        f = ("Arial",16) 
        Button(cf,text="Quit",font=f,command=self.cmd_quit).pack(side=LEFT,expand=1,fill=X)
        Button(cf,text="Stop",font=f,command=self.cmd_stop).pack(side=LEFT,expand=1,fill=X)
        Button(cf,text="Run",font=f,command=self.cmd_start).pack(side=LEFT,expand=1,fill=X)
        cf.pack(side=BOTTOM,expand=1,fill=X)

    def cmd_start(self):
        self.newwave()
        self.play()
        
    def cmd_stop(self):
        self.RUNNING = False

    def cmd_quit(self):
        self.cmd_stop()
        self.root.after(200)
        self.root.destroy()

    def nframe(self,parent):
        nf = Frame(parent,relief=RIDGE,bd=3)
        f = ("Arial",10)
        self.number = StringVar()
        self.number.set("51")
        self.constant = StringVar()
        self.constant.set("0.5")
        Label(nf,font=f,text="number of\noscillators").grid(row=0,column=0,padx=5)
        Entry(nf,width=4,textvariable=self.number,font=f).grid(row=0,column=1)
        Label(nf,font=f,text="force\ncoupling").grid(row=1,column=0)
        Entry(nf,width=4,textvariable=self.constant,font=f).grid(row=1,column=1)
        
        nf.pack(side=TOP)

    def endframe(self,parent):
        ef = Frame(parent,relief=RIDGE,bd=3)
        f = ("Arial",10) 
        self.leftend = IntVar()
        self.rightend = IntVar()
        self.leftend.set(2)
        self.rightend.set(2)
        
        Label(ef,font=f,text="left < ending > right").grid(row=0,column=0,columnspan=2)
        Radiobutton(ef,command=self.newwave,font=f,variable=self.leftend,value=1,text="fixed",bd=0).grid(row=1,column=0,sticky="W")
        Radiobutton(ef,command=self.newwave,font=f,variable=self.leftend,value=2,text="free",bd=0).grid(row=2,column=0,sticky="W")

        Radiobutton(ef,command=self.newwave,font=f,variable=self.rightend,value=1,text="fixed",bd=0).grid(row=1,column=1,sticky="W")
        Radiobutton(ef,command=self.newwave,font=f,variable=self.rightend,value=2,text="free",bd=0).grid(row=2,column=1,sticky="W")

        ef.pack(side=TOP)

    def waveframe(self,parent):
        wf = Frame(parent,relief=RIDGE,bd=3)
        f = ("Arial",10) 
        self.waveselect = IntVar()
        self.frequencies = []
        for i in range(5):
            self.frequencies.append(StringVar())
            self.frequencies[i].set("1.0")

        Radiobutton(wf,command=self.newwave,font=f,variable=self.waveselect,value=1,text="pluck",bd=0).grid(row=0,column=0,sticky="W")
        Radiobutton(wf,command=self.newwave,font=f,variable=self.waveselect,value=2,text="pluck1",bd=0).grid(row=2,column=0,sticky="W")
        Radiobutton(wf,command=self.newwave,font=f,variable=self.waveselect,value=3,text="pluck twice",bd=0).grid(row=1,column=0,sticky="W")
        Radiobutton(wf,command=self.newwave,font=f,variable=self.waveselect,value=4,text="kick1",bd=0).grid(row=3,column=0,sticky="W")
        Radiobutton(wf,command=self.newwave,font=f,variable=self.waveselect,value=5,text="sin",bd=0).grid(row=4,column=0,sticky="W")
        e1 = Entry(wf,font=f,width=5,textvariable=self.frequencies[0])
        e1.grid(row=4,column=1)
        Radiobutton(wf,command=self.newwave,font=f,variable=self.waveselect,value=6,text="cos",bd=0).grid(row=5,column=0,sticky="W")
        e2 = Entry(wf,font=f,width=5,textvariable=self.frequencies[1])
        e2.grid(row=5,column=1)
        Radiobutton(wf,command=self.newwave,font=f,variable=self.waveselect,value=7,text="rect",bd=0).grid(row=6,column=0,sticky="W")
        e3 = Entry(wf,font=f,width=5,textvariable=self.frequencies[2])
        e3.grid(row=6,column=1)
        Radiobutton(wf,command=self.newwave,font=f,variable=self.waveselect,value=8,text="saw",bd=0).grid(row=7,column=0,sticky="W")
        e4 = Entry(wf,font=f,width=5,textvariable=self.frequencies[3])
        e4.grid(row=7,column=1)
        Radiobutton(wf,command=self.newwave,font=f,variable=self.waveselect,value=9,text="tri",bd=0).grid(row=8,column=0,sticky="W")
        e5 = Entry(wf,font=f,width=5,textvariable=self.frequencies[4])
        e5.grid(row=8,column=1)
        for e in (e1,e2,e3,e4,e5): e.bind('<Return>',self.newwave1)
        wf.pack(side=TOP)        
        self.waveselect.set(1)

    def newwave1(self,key):
        self.newwave()
        
    def newwave(self):
        self.RUNNING = False
        self.root.after(100)

        new_n = int(self.number.get())
        if new_n<15: new_n = 15
        if new_n != self.n:
            self.n = new_n
            self.makeoscillators(self.n)
            
        self.D = float(self.constant.get())
        if self.leftend.get()==1: self.free_left = False
        else: self.free_left = True
        if self.rightend.get()==1: self.free_right = False
        else: self.free_right = True
        

        w = self.waveselect.get()
        if w==1:
            self.s,self.v = pluck(self.n)
        elif w==2:
            self.s,self.v = pluck1(self.n)
        elif w==3:
            self.s,self.v = pluck2(self.n)
        elif w==4:
            self.s,self.v = kick(self.n)
        elif w==5:
            self.s,self.v = sine(self.n,float(self.frequencies[0].get()))
        elif w==6:
            self.s,self.v = cosine(self.n,float(self.frequencies[1].get()))
        elif w==7:
            self.s,self.v = rectangle(self.n,float(self.frequencies[2].get()))
        elif w==8:
            self.s,self.v = saw(self.n,float(self.frequencies[3].get()))
        elif w==9:
            self.s,self.v = triangle(self.n,float(self.frequencies[4].get()))

        self.refresh()
        
############################################################################
        
# RKN4 y''=f(t,y)
def rkn4_0(t,y,yd,f,h):
    n = len(y)
    ynew = [0.0]*n
    ydnew = [0.0]*n
    
    k1 = f(t,y)
    for i in range(n):
        ynew[i] = y[i]+h/2.0*yd[i]+h*h/8.0*k1[i]
    k2 = f(t+h/2.0,ynew)
    for i in range(n):
        ynew[i] = y[i]+h*yd[i]+h*h/2.0*k2[i]
    k3 = f(t+h,ynew)
    for i in range(n):
        ynew[i] = y[i]+h*yd[i]+h*h/6.0*(k1[i]+2.0*k2[i])
        ydnew[i] = yd[i]+h/6.0*(k1[i]+4.0*k2[i]+k3[i])
    return (t+h,ynew,ydnew)

def energy(s,v,D):
    ekin = epot = 0
    for i in range(len(s)): ekin += v[i]**2
    ekin = ekin/2
    for i in range(1,len(s)): epot += (s[i]-s[i-1])**2
    epot = D*epot/2
    return ekin,epot,ekin+epot


############################################################################
# initial waveforms
############################################################################

def sine(n,frequ=1):
    s = [0.0]*n
    x,dx = 0.0,2*pi/(n-1)
    for i in range(n):
        s[i] = sin(x*frequ)
        x += dx
    return s,[0.0]*n
        
def cosine(n,frequ=1):
    s = [0.0]*n
    x,dx = 0.0,2*pi/(n-1)
    for i in range(n):
        s[i] = cos(x*frequ)
        x += dx
    return s,[0.0]*n

# kick center oscillator
def kick(n):
    s = [0.0]*n
    v = [0.0]*n
    v[n/2] = 1
    return s,v

# pluck center oscillator
def pluck1(n):
    s = [0.0]*n
    v = [0.0]*n
    s[n/2] = 1
    return s,v

# pluck 15 center oscillators
def pluck(n):
    s = [0.0]*n
    v = [0.0]*n
    m = n/2
    #d = [1.0,0.9,0.7,0.5,0.3,0.2,0.1,0.05]
    d = [1.0,0.88,0.61,0.32,0.14,0.04,0.01]
    for i in range(len(d)):
        s[m-i] = s[m+i] = d[i] 
    return s,v

def pluck2(n):
    if n<30: return pluck(n)
    s = [0.0]*n
    v = [0.0]*n
    m = n/2
    mm = n/4
    d = [1.0,0.88,0.61,0.32,0.14,0.04,0.01]
    for i in range(len(d)):
        s[m-mm-i] = s[m-mm+i] = d[i] 
        s[m+mm-i] = s[m+mm+i] = -d[i] 
    return s,v

def saw(n,frequ=1):
    w = [0.0]*n
    for i in range(n):
        x = float(i)/(n)*frequ
        x = x-int(x)
        w[i] = -1.0+x*2
    return w,[0.0]*n

def triangle(n,frequ=1):
    w = [0.0]*n
    for i in range(n):
        x = float(i)/(n-1)*frequ
        x = x-int(x)
        if x<0.5: w[i] = -1.0+x*4
        else: w[i] = 1.0-(x-0.5)*4
    return w,[0.0]*n
    
def rectangle(n,frequ=1):
    w = [0.0]*n
    for i in range(n):
        x = float(i)/(n)*frequ
        x = x-int(x)
        if x<0.5: w[i] = -1.0
        else: w[i] = 1.0
    return w,[0.0]*n
    

v = Vibrate(51)