###
### Bayes Filter Demo
### ~~~~~~~~~~~~~~~~~
###
### (C) 2009 Rodrigo Ventura, ISR / IST
### Email: rodrigo.ventura@isr.ist.utl.pt
###

import sys
import math
import random
import Tkinter


def p_normal(x, u, sigma2):
    return math.exp(-(x-u)**2/(2*sigma2)) / math.sqrt(2*math.pi*sigma2)


class Viewer(Tkinter.Tk):
    VERSION = "1.1 (5-Oct-2009)" ### <<< VERSION STRING
    N = 20
    ## Interface functions
    def __init__(self):
        Tkinter.Tk.__init__(self)
        self.title("Bayes filter demo "+self.VERSION)
        self.geometry("%sx300"%(110+25*self.N))
        Tkinter.Label(self, text="(C) 2009 Rodrigo Ventura, ISR / IST", font="verdana 9").pack(side="bottom", expand=True, anchor="sw")
        self.canvas = c = Tkinter.Canvas(self)
        c.pack(side="bottom", expand=True, fill="both")
        self.var_meas = Tkinter.IntVar()
        Tkinter.Button(self, text="Init", command=self.init).pack(side="left")
        Tkinter.Button(self, text="<", command=self.do_left).pack(side="left")
        Tkinter.Button(self, text="=", command=self.do_stay).pack(side="left")
        Tkinter.Button(self, text=">", command=self.do_right).pack(side="left")
        # filter menu
        mb = Tkinter.Menubutton(self, text="Filter")
        m = Tkinter.Menu(mb)
        m.add_command(label="Discrete", command=lambda : self.set_filter("discrete"))
        m.add_command(label="Kalman", command=lambda : self.set_filter("kalman"))
        mb.configure(menu=m)
        mb.pack(side="left")
        #
        Tkinter.Checkbutton(self, text="Measurements", variable=self.var_meas).pack(side="left")
        Tkinter.Button(self, text="About", command=self.do_about).pack(side="right")
        c.bind("<Button-1>", self.mouse_press)
        # map slots
        for i in xrange(self.N):
            self.canvas.create_line(50+i*25, 200, 70+i*25, 200)
        self.m = Kalman(self)
    def set_filter(self, name):
        self.m.cleanup()
        if name=="discrete":
            self.m = Discrete(self)
        elif name=="kalman":
            self.m = Kalman(self)
        self.m.init()
        self.m.update_pdf()
    def update_robot(self):
        self.canvas.delete("robot")
        self.canvas.create_oval(50+self.x*25, 190, 70+self.x*25, 170, tags="robot")
    def show_measure(self, z):
        self.canvas.delete("meas")
        if z:
            self.canvas.create_text(50, 210, text="z = %.2f"%z, \
                                    fill="blue", anchor="nw", tags="meas")
            u = 60 + z*25
            self.canvas.create_line(u, 150, u, 160, fill="blue", tags="meas")
    def do_left(self):
        if self.x>0:
            if random.random()>0.2:
                self.x -= 1
        self.do_step(-1)
    def do_stay(self):
        self.do_step(0)
    def do_right(self):
        if self.x<self.N-1:
            if random.random()>0.2:
                self.x += 1
        self.do_step(1)
    def do_step(self, u):
        self.update_robot()
        z = self.measure() if self.var_meas.get() else None
        self.show_measure(z)
        self.m.do_filter(u, z)
        self.m.update_pdf()
    def mouse_press(self, ev):
        x, y = ev.x, ev.y
        if y>=170 and y<=190:
            i = (x-50)//25
            if i>=0 and i<self.N:
                self.x = i
                self.update_robot()
    def do_about(self):
        a = Tkinter.Tk()
        a.title("About...")
        a.geometry("300x150")
        Tkinter.Label(a, text="Bayes filtering demo\nversion "+self.VERSION+"\n(C) 2009 Rodrigo Ventura\n<rodrigo.ventura@isr.ist.utl.pt>\nISR / IST, Lisbon, Portugal").pack(expand=True)
    ## Core functions
    def init(self):
        self.m.init()
        self.m.update_pdf()
        self.x = random.randint(0, self.N-1)
        self.update_robot()
        self.show_measure(None)
    def measure(self):
        """Simulate a measurment"""
        return random.gauss(self.x, 1)
    def main(self):
        self.init()
        self.mainloop()


class Discrete():
    ## Interface functions
    def __init__(self, viewer):
        self.v = viewer
        self.N = viewer.N
    def update_pdf(self):
        self.v.canvas.delete("pdf")
        for i in xrange(self.N):
            self.v.canvas.create_rectangle(50+i*25, 120, 70+i*25, 120-100*self.px[i], \
                                           fill="#888", tags="pdf")
        m = 0
        for i in xrange(self.N):
            m += i*self.px[i]
        u = 50 + (m+0.5)*25
        self.v.canvas.create_line(u, 130, u, 140, fill="#888", tags="pdf")
    def cleanup(self):
        self.v.canvas.delete("pdf")
    ## Core functions
    def init(self):
        self.px = self.N*[1.0/self.N]
    def do_filter(self, u, z):
        """Discrete filter algorithm: given control u and measurment z, updates px distribution"""
        bel1 = self.N*[0]
        for j in xrange(self.N):
            for i in xrange(self.N):
                bel1[j] += self.p_trans(i, u, j)*self.px[i]
        if z:
            bel2 = []
            for j in xrange(self.N):
                bel2.append(self.p_meas(z, j)*bel1[j])
        else:
            bel2 = bel1
        eta = 1.0/sum(bel2)
        self.px = map(lambda p: eta*p, bel2)
    def p_trans(self, x1, a, x2):
        """State transition model"""
        if a==0:
            return 1 if x1==x2 else 0
        else:
            y = x1+a
            if y<0 or y>=self.N:
                return 1 if x1==x2 else 0
            else:
                if x2==y:
                    return 0.8
                elif x2==x1:
                    return 0.2
                else:
                    return 0
    def p_meas(self, z, x):
        """Measurment model"""
        sigma2 = 1
        return p_normal(z, x, sigma2)


class Kalman():
    ## Interface functions
    def __init__(self, viewer):
        self.v = viewer
        self.N = viewer.N
        # plot axis arrows
        self.v.canvas.create_line(60, 120, 60+25*self.N, 120, fill="#888", arrow="last", tags="axis")
        self.v.canvas.create_line(60, 120, 60, 20, fill="#888", arrow="last", tags="axis")
    def update_pdf(self):
        self.v.canvas.delete("pdf")
        pts = []
        for u in xrange(25*self.N):
            x = u/25.0
            px = p_normal(x, self.ux, self.sx)
            pts.extend([60+u,120-100*px])
        self.v.canvas.create_line(pts, tags="pdf")
        v = 60 + self.ux*25
        self.v.canvas.create_line(v, 130, v, 140, fill="#888", tags="pdf")
    def cleanup(self):
        self.v.canvas.delete("pdf", "axis")
    ## Core functions
    def init(self):
        self.ux = self.N/2.0
        self.sx = 2.0*self.N
    def do_filter(self, cmd, z):
        """Kalman filter algorithm: given control u and measurment z, updates px distribution"""
        sigma2_state = 0.16
        sigma2_meas  = 1.0
        u = 0.8*cmd # scaling to get state transition mean value
        if u==0:
            u1, s1 = self.ux, self.sx
        else:
            # (1) predict step   
            u1 = self.ux + u
            s1 = self.sx + sigma2_state
        if z is None:
            u2, s2 = u1, s1
        else:
            # (2) update step
            k  = s1/(s1+sigma2_meas)
            u2 = u1 + k*(z-u1)
            s2 = (1-k)*s1
        self.ux, self.sx = u2, s2


# Main entry
if __name__=="__main__":
    Viewer().main()

# EOF