import GridStateEstimator
reload(GridStateEstimator)
from GridStateEstimator import *

import Map
reload(Map)
from Map import *

import AvoidWanderTB
reload(AvoidWanderTB)
from AvoidWanderTB import *

import Sequence
reload(Sequence)
from Sequence import *

import sonarModel
reload(sonarModel)
from sonarModel import *

from robotPoses import *

import SoaR
import math

## Change this to be your directory
labDirectory = "/Users/tlp/Documents/Teaching/6.01/fall07/ps14/ps14-code/"
dataDirectory = labDirectory+"distAngle/"

# pick one of the lists of robot poses
robot_poses = robot_2465
#robot_poses = robot_50

def setup():
    robot.updateNo = 0
    gridN = 20

    # Only one of the following calls to GridStateEstimator should be
    # uncommented.
    # To change granularity, remember to change both the numeric argument to
    # GridStateEstimator and the file name.

    # Uncomment these lines to use the She world
#     (xmin, xmax, ymin, ymax) = (0.0, 4.0, 0.0, 4.0)
#     m = GridStateEstimator(Map(sheBoxes), xmin, xmax, ymin, ymax, gridN,\
#                            dataDirectory + "she20.dat",
#                            numBestPoses = 5)

    # Uncomment these lines to use the boxWorld world
    (xmin, xmax, ymin, ymax) = (0.0, 3.05, 0.0, 3.05)
    m = GridStateEstimator(Map(boxWorldBoxes), xmin, xmax, ymin, ymax, gridN,\
                          dataDirectory + "boxWorld20.dat", obsProb,
                          numBestPoses = 5)
    robot.m = m

    # Uncomment these lines to use Shannon's world
    #(xmin, xmax, ymin, ymax) = (0.0, 3.05, 0.0, 3.05)
    #m = GridStateEstimator(Map(shannonBoxes), xmin, xmax, ymin, ymax, gridN,\
    #                       dataDirectory + "shannon20.dat",
    #                       numBestPoses = 5)

    # Uncomment these lines to use the emptySquare world
    #(xmin, xmax, ymin, ymax) = (0.0, 3.05, 0.0, 3.05)
    #m = GridStateEstimator(Map(emptySquareBoxes), xmin, xmax, ymin, ymax, gridN,\
    #                      dataDirectory + "emptySquare20.dat",
    #                      numBestPoses = 5)

    # Make two 300 x 300 pixel windows, with a free margin in each 
    # direction
    (xrange, yrange) = (xmax-xmin, ymax-ymin)
    (wxmin, wxmax, wymin, wymax) = (xmin - 0.05*xrange, xmax + 0.05*xrange,
                                    ymin - 0.05*yrange, ymax +
                                    0.05*yrange)
    robot.beliefWindow = DrawingWindow(300, 300, wxmin, wxmax, wymin, wymax, \
                                 "Belief")
    robot.percWindow = DrawingWindow(300, 300, wxmin, wxmax, wymin, wymax,"P(O | S)")
    # Draw initial belief and initialize pose for state estimator
    m.drawBelief(robot.beliefWindow)
    m.initPose(pose())

    # Just move for 0.1 seconds, no matter how long it takes us to
    # recompute the next action.
    def discreteMotor(trans, rot):
        discreteMotorOutput(trans, rot, 0.1)

    # for drawing true pose
    size = int(floor(0.5*robot.beliefWindow.windowWidth/float(gridN)))-2

    # Returns a function that updates the belief state every n times
    # it's called.
    def makeBeliefUpdateEveryN(n):
        updateCount = [n]
        def beliefUpdateEveryN():
            if updateCount[0] == n:
                updateCount[0] = 0
                m.update(sonarDistances(), pose())
                m.drawObsP(robot.percWindow)
                m.drawBelief(robot.beliefWindow)
                print 'Best pose:', m.kBestPoses[0]
                if robot.updateNo < len(robot_poses):
                    ## Get the true pose indices for this update
                    p = robot_poses[robot.updateNo]
                    (x,y,th) = m.indicesToPose(p)
                    ## Draw the true pose in purple on the belief window
                    robot.beliefWindow.drawRobotWithNose(x, y, th,
                                                         color = "purple",
                                                         size=size)
                    ## Get the probability for this observation (sonarDistances)
                    ## at the true pose, this will printi the individual sonar
                    ## probabilities, as well as the total printed here.
                    obsProb = m.prob(p, sonarDistances())
                    print "P(o | truePose) = ", math.exp(obsProb)
                    robot.updateNo += 1
                    print "true indices=", p
                    ## Compute and print your quality metric here...
            updateCount[0] += 1
        return beliefUpdateEveryN

    robot.driver = TBParallelWithFun(AvoidWanderTB(sonarDistances, 
                                                   discreteMotor),
                                     makeBeliefUpdateEveryN(10))
    robot.driver.start()
    readLog(labDirectory+"robot_2.465.log")
#    readLog(labDirectory+"robot_5.0.log")


def step():
    robot.driver.step()

def destroy():
    robot.percWindow.destroy()
    robot.beliefWindow.destroy()

