import GridStateEstimator
reload(GridStateEstimator)
from GridStateEstimator import *

import Map
reload(Map)
from Map import *

import XYThGridPlanner
reload(XYThGridPlanner)
from XYThGridPlanner import *

import XYThDriver
reload(XYThDriver)
from XYThDriver import *

import Sequence
reload(Sequence)
from Sequence import *

import sonarModel
reload(sonarModel)
from sonarModel import *

import SoaR

## Change this to be your directory
dataDirectory = "/afs/athena.mit.edu/user/t/l/tlp/Desktop/6.01/lab14/distAngle/"

def cheatPose():
    return app().output.abspose.get()

def setup():
    SoaR.Output.Simulator.SONAR_ANGLE = 'bubble'
    SoaR.Output.Simulator.MAX_RELIABLE_SONAR_DIST = 1.5
    SoaR.Output.Simulator.MAX_SONAR_DIST = 1.75
    SoaR.Output.Simulator.SONAR_MAX_REPORT = 5.0

    # 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, 20,\
    #                       dataDirectory + "she20.dat", obsProb)

    (xmin, xmax, ymin, ymax) = (0.0, 3.05, 0.0, 3.05)
    m = GridStateEstimator(Map(boxWorldBoxes), xmin, xmax, ymin, ymax, 20,\
                           dataDirectory + "boxWorld20.dat", obsProb,
                           numBestPoses = 5)

    boxDim = 4.0/20.0
    # Make three 300 x 300 pixel windows, with coordinates ranging from
    # -0.5 to 4.5 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.planWindow = DrawingWindow(300, 300, wxmin, wxmax, wymin, wymax, "Plan")
    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())

    # Goal grid pose and indices
    goalPose = (0.6, 0.6, math.pi / 2)
    goalIndicesAndDir = m.poseToIndicesAndDir(goalPose)

    # Remember the last plan we drew, so we don't redraw too much
    robot.lastDrawnPlan = None
    # Remember whether the last action was to move, so we don't update
    # the belief state too much
    robot.numContiguousStopActions = 0

    planner = XYThGridPlanner(m)

    # 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)

    def beliefUpdate():
        m.update(sonarDistances(), pose())
        m.drawObsP(robot.percWindow)
        m.drawBelief(robot.beliefWindow)

    updateInterval = 30
    
    # Procedure that will plan a path and execute the first step
    def dynamicReplanner():        
        # Update belief if we've just moved or if we've been sitting
        # for a long time
        if robot.numContiguousStopActions % updateInterval == 0:
            beliefUpdate()

        # Get current robot pose.  
        currentPose = m.mostLikelyPose()
        # Get grid indices for robot's x,y location and
        # 0, 1, 2, 3 for the closest compass direction
        currentIndicesAndDir = m.poseToIndicesAndDir(currentPose)
        # Make a plan and draw it into the window
        plan = planner.makePlan(currentIndicesAndDir, goalIndicesAndDir)
        if plan and (plan != robot.lastDrawnPlan):
            planner.drawPlan(plan, robot.planWindow)
            robot.lastDrawnPlan = plan
            samePlan = True
        else:
            samePlan = False
        if plan and len(plan) == 1:
            print "Woo hoo", goalIndicesAndDir
        if plan and len(plan) > 1:
            if not samePlan:
                robot.numContiguousStopActions = 0
            # Return a terminating behavior to drive to the first subgoal
            # Subgoal is in planning world coordinates
            subgoal = m.indicesAndDirToPose(plan[1])
            # Compute the vector to the goal, in robocentric coordinates
            desiredDeltaXYTh = globalPoseToLocalPose(m.mostLikelyPose(),
                                                     subgoal)
            # Now, compute a subgoal in the odometry frame, for the driver
            odometryGoal = localPoseToGlobalPose(pose(), desiredDeltaXYTh)
            print "Subgoal Indices: ", plan[1]
            print "Current Indices: ", currentIndicesAndDir
            print "World coords: ", subgoal
            print "Robot coords: ", desiredDeltaXYTh
            print "Odo coords: ", odometryGoal
            print "Current pose: ", pose()
            return XYThDriverOA(odometryGoal, pose, discreteMotor, \
                                  sonarDistances, boxDim/2.0)
        else:
            robot.numContiguousStopActions += 1
            return Stop(discreteMotor)
    robot.driver = TBStream(dynamicReplanner)
    robot.driver.start()

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

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