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

import sys
import math
import random
import Tkinter

class Main(Tkinter.Tk):
    VERSION = "1.0 (1-Oct-2009)" ### <<< VERSION STRING
    N = 10
    ## Interface functions
    def __init__(self):
        Tkinter.Tk.__init__(self)
        self.title("Bayes filter demo "+self.VERSION)
        self.geometry("360x300")
        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")
        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)
        for i in xrange(self.N):
            self.canvas.create_line(50+i*25, 200, 70+i*25, 200)
    def update_pdf(self):
        self.canvas.delete("pdf")
        for i in xrange(self.N):
            self.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.canvas.create_line(u, 130, u, 140, fill="#888", tags="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 = 50 + (z+0.5)*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.bayes_filter(u, z)
        self.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 filter 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.px = self.N*[1.0/self.N]
        self.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 bayes_filter(self, u, z):
        """Bayes 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 self.var_meas.get():
            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 math.exp(-(z-x)**2/(2*sigma2)) / math.sqrt(2*math.pi*sigma2)
    def main(self):
        self.init()
        self.mainloop()


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

# EOF