import pickle
from math import log, floor, exp

import probcolors
reload(probcolors)
from probcolors import *

import GridMap
reload(GridMap)
from GridMap import *

import KBest
reload(KBest)
from KBest import *

from math import pi, cos

# Use this as log(0)
beliefEpsilon = -400.0

class GridStateEstimator(GridMap):
    # Estimate the pose of the robot.  Store a grid of probability
    # values, and update them based on observations and transitions.
    # We'll actually store the log values, because they're easier to
    # work with without losing precision.

    # If pose has a value, start the robot knowing it's in that pose;
    # otherwise start with a uniform belief state

    def __init__(self, theMap, xmin, xmax, ymin, ymax, gridN,\
                     readingsPath, obsModel, pose = None, numBestPoses = 5):

        # initialize the basic map
        GridMap.__init__(self, theMap, xmin, xmax, ymin, ymax, gridN)
        self.numBestPoses = numBestPoses
        self.obsModel = obsModel
        
        # Load the cached nominal sonar readings from a .dat file
        f = open(readingsPath, "r")
        self.idealReadings = pickle.load(f)
        f.close()

        # Allocate and initialize the belief state
        if pose == None:
            # Uniform value is log(1/(n**3))
            uniformVal = -3.0*log(self.gridN)
            self.belief = make3DArray(self.gridN, self.gridN, \
                                          self.gridN, uniformVal)
            self.bestPoseIndices = None
        else:
            # This doesn't quite normalize to 1, but it will get fixed
            # right away
            self.belief = make3DArray(self.gridN, self.gridN, \
                                          self.gridN, beliefEpsilon)
            (px,py,pth) = pose
            self.setBeliefAtPose((px, py, pth), 0.0)
            self.bestPoseIndices = self.poseToIndices(pose)

    def initPose(self, odometry):
        self.oldPose = odometry

    def setBeliefAtPose(self, pose, val):
        (ix, iy, ith) = self.poseToIndices(pose)
        self.belief[ix][iy][ith] = val

    def getBeliefAtIndices(self, (ix, iy, ith)):
        return self.belief[ix][iy][ith]

    def getBeliefAtPose(self, pose):
        (ix, iy, ith) = self.poseToIndices(pose)
        return self.belief[ix][iy][ith]

    def mostLikelyPose(self):
        return self.indicesToPose(self.bestPoseIndices)

    # Draw the belief state, summing over orientations for each
    # grid cell
    def drawBelief(self, window):
        size = int(floor(0.9*window.windowWidth/float(self.gridN)))-2
        uniformProb = (1.0/self.gridN)**2
        window.clear()
        x = self.xmin
        for xIndex in range(self.gridN):
            y = self.ymin
            for yIndex in range(self.gridN):
                if not self.free((xIndex,yIndex)):
                    window.drawRobot(x, y, x, y,
                                     color = 'gray',
                                     size = size)
                else:
                    weight = sum([math.exp(th) for th \
                                  in self.belief[xIndex][yIndex]])
                    window.drawRobot(x, y, x, y,
                                     color = probToPyColor(weight, uniformProb),
                                     size = size)
                y = y + self.ystep
            x = x + self.xstep
        if self.bestPoseIndices:
            bestColor = "orange"
            (bx,by,btheta) = self.indicesToPose(self.bestPoseIndices)
            window.drawRobotWithNose(bx, by, btheta, color = bestColor,\
                                         size = max(size/2,1))

    # Draw the observation probability, maxing over orientations for each
    # grid cell
    def drawObsP(self, window):
        def avg(nums):
            return sum(nums) / float(len(nums))
        size = int(floor(0.9*window.windowWidth/float(self.gridN)))-2
        window.clear()
        # Horribly slow
        #maxP = max([max([math.exp(p) for p in row]) for row in self.logObsP])
        #minP = min([min([math.exp(p) for p in row]) for row in self.logObsP])
        #avgP = avg([avg([math.exp(p) for p in row]) for row in self.logObsP])
        # Derived by looking at the exact but slow values from above
        maxP = 10
        minP = 0
        avgP = 0.1
        Prange = maxP - minP
        #print "min, max, avg obs: ", minP, maxP, avgP
        avgV = (avgP - minP) / Prange

        x = self.xmin
        for xIndex in range(self.gridN):
            y = self.ymin
            for yIndex in range(self.gridN):
                weight = (math.exp(self.logObsP[xIndex][yIndex]) - minP) /Prange
                if not self.free((xIndex,yIndex)):
                    window.drawRobot(x, y, x, y,
                                     color = 'gray',
                                     size = size)
                else:
                    window.drawRobot(x, y, x, y,
                                     color = probToPyColor(weight, avgV),
                                     size = size)

#                 elif weight > 0.995:
#                     window.drawRobot(x, y, x, y, color = color1,
#                                      size = size)
#                 elif weight > 0.99:
#                     window.drawRobot(x, y, x, y, color = color2,
#                                      size = size)
#                 elif weight > 0.95:
#                     window.drawRobot(x, y, x, y, color = color3,
#                                      size = size)
#                 else:
#                     window.drawRobot(x, y, x, y, color = color4,
#                                      size = size)
                y = y + self.ystep
            x = x + self.xstep

    def update(self, sonarDistances, odometry):
        globalDeltaPose = diffPose(odometry, self.oldPose)
        # Delta pose relative to old pose
        localDeltaPose = globalDeltaToLocal(self.oldPose, globalDeltaPose[0:2])
        localDeltaPose.append(globalDeltaPose[2])
        self.transition(localDeltaPose)
        self.reweight(sonarDistances)
        self.oldPose = odometry

    def entropy(self):
        return self.lastEntropy
        #return sumArray3D(self.belief, lambda lp: lp * math.exp(lp))

    def prob(self, (ix, iy, ith), sonarDistances):
        idealReadings = self.idealReadings[ix][iy][ith]
        # Figure out sensor likelihood at this pose;
        # Sum over sensors
        ps = 0.0
        for i in range(8):
            ps += log(self.obsModel(sonarDistances[i], idealReadings[i], verbose=True))
        return ps

    def reweight(self, sonarDistances):
        self.logObsP = make2DArray(self.gridN, self.gridN, beliefEpsilon)
        # Exponent in log-sum-exp trick
        b = 15
        # For each grid location, compute P(sonars | pose), using the
        # nominal pose of the grid location.  Use this as the weight.
        # Keep the sum and renormalize at the end.  Return most likely state.
        totalWeight = 0.0
        totalProb = 0.0
        sumLogProbs = 0.0
        unnormalizedEntropy = 0.0
        kBest = KBest(self.numBestPoses)
        for ith in range(self.gridN):
            for ix in range(self.gridN):
                for iy in range(self.gridN):
                    if self.free((ix,iy)):
                        idealReadings = self.idealReadings[ix][iy][ith]
                        # Figure out sensor likelihood at this pose;
                        # Sum over sensors
                        ps = 0.0
                        for i in range(8):
                            ps += log(self.obsModel(sonarDistances[i], idealReadings[i]))
                        # Add likelihood into grid
                        self.belief[ix][iy][ith] += ps
                        self.logObsP[ix][iy] = max(ps,self.logObsP[ix][iy])
                        kBest.insert((ix, iy, ith), self.belief[ix][iy][ith])
                    else:
                        self.belief[ix][iy][ith] = beliefEpsilon
                    # Increment total weight in map, and entropy stuff
                    logProb = self.belief[ix][iy][ith]
                    shiftedProb = math.exp(logProb + b)
                    prob = math.exp(logProb)
                    totalWeight += shiftedProb
                    totalProb += prob
                    unnormalizedEntropy += prob * logProb
                    sumLogProbs += logProb
        normalizer = math.log(totalWeight) - b
        mapArray3D(self.belief, lambda x: x - normalizer)
        self.bestPoseIndices = kBest.best()[1]
        self.kBestPoses = [(logp - normalizer, pose) for (logp, pose) in \
                           kBest.kBest()]
        self.lastEntropy = (unnormalizedEntropy/totalProb) - math.log(totalProb)

    def transition(self, (dx, dy, dtheta)):
        newBelief = make3DArray(self.gridN, self.gridN, self.gridN, \
                                 (1.0/float(self.gridN))**10)
        dtheta = fixAngle02Pi(dtheta)
        th = self.thetaMin
        for ith in range(self.gridN):
            # for this angle, compute offsets in world units
            (ox,oy) = localToGlobal((0,0,th),(dx,dy))
            # convert offsets to grid units
            oxGrid = ox/self.xstep
            oyGrid = oy/self.ystep
            othGrid = dtheta/self.thetaStep
            # Integral offsets
            oxi = int(oxGrid)
            oyi = int(oyGrid)
            othi = int(othGrid)
            # Fractional offsets
            # What percentage lands in the nominal spot in each dimension?
            oxfa = abs(oxGrid - oxi)
            oyfa = abs(oyGrid - oyi)
            othfa = abs(othGrid - othi)
            # Use signs to change indices into new grid;  use abs
            # in multiplier.
            # oxs is the sign of the difference between the integral
            # offset and the actual offset.  For example, if oxGrid is
            # 1.1 and oxi is 1.0, then oxs is positive.  That means
            # that .1 (oxfa) of the probability mass goes in the cell
            # that is one bigger in the x dimension.
            oxs = sign(oxGrid - oxi)
            oys = sign(oyGrid - oyi)
            oths = sign(othGrid - othi)

            for ix in range(self.gridN):
                for iy in range(self.gridN):
                    # Need an abbreviation
                    n = self.gridN
                    # No noise.  Add mass from here to the new grid.
                    # Deal with obstacles.  Slightly wrong.  Assumes
                    # that if the mass would move into a square with
                    # an obstacle, then it stays where it was.
                    # Broken when steps are big compared to grid size.
                    if self.free((ix+oxi,iy+oyi)):
                        newBelief[ix+oxi][iy+oyi][(ith+othi)%n] += \
                                (1-oxfa)*(1-oyfa)*(1-othfa)*\
                                math.exp(self.belief[ix][iy][ith])
                    else:
                        newBelief[ix][iy][ith] += \
                                (1-oxfa)*(1-oyfa)*(1-othfa)*\
                                math.exp(self.belief[ix][iy][ith])
                    if self.free((ix+oxi,iy+oyi)):
                        newBelief[ix+oxi][iy+oyi][(ith+othi+oths)%n] += \
                                (1-oxfa)*(1-oyfa)*(othfa)* \
                                math.exp(self.belief[ix][iy][ith])
                    else:
                        newBelief[ix][iy][ith] += \
                                (1-oxfa)*(1-oyfa)*(othfa)*\
                                math.exp(self.belief[ix][iy][ith])
                    if self.free((ix+oxi,iy+oyi+oys)):
                        newBelief[ix+oxi][iy+oyi+oys][(ith+othi)%n] += \
                                (1-oxfa)*(oyfa)*(1-othfa)*\
                                math.exp(self.belief[ix][iy][ith])
                    else:
                        newBelief[ix][iy][ith] += \
                            (1-oxfa)*(oyfa)*(1-othfa)*\
                            math.exp(self.belief[ix][iy][ith])
                    if self.free((ix+oxi,iy+oyi+oys)):
                        newBelief[ix+oxi][iy+oyi+oys][(ith+othi+oths)%n] += \
                                (1-oxfa)*(oyfa)*(othfa)*\
                                math.exp(self.belief[ix][iy][ith])
                    else:
                        newBelief[ix][iy][ith] += \
                               (1-oxfa)*(oyfa)*(othfa)*\
                               math.exp(self.belief[ix][iy][ith])
                    if self.free((ix+oxi+oxs,iy+oyi)):
                        newBelief[ix+oxi+oxs][iy+oyi][(ith+othi)%n] += \
                                 (oxfa)*(1-oyfa)*(1-othfa)*\
                                 math.exp(self.belief[ix][iy][ith])
                    else:
                        newBelief[ix][iy][ith] += \
                               (oxfa)*(1-oyfa)*(1-othfa)*\
                               math.exp(self.belief[ix][iy][ith])
                    if self.free((ix+oxi+oxs,iy+oyi)):
                        newBelief[ix+oxi+oxs][iy+oyi][(ith+othi+oths)%n] += \
                                (oxfa)*(1-oyfa)*(othfa)*\
                                math.exp(self.belief[ix][iy][ith])
                    else:
                        newBelief[ix][iy][ith] += \
                                 (oxfa)*(1-oyfa)*(othfa)*\
                                 math.exp(self.belief[ix][iy][ith])
                    if self.free((ix+oxi+oxs,iy+oyi+oys)):
                        newBelief[ix+oxi+oxs][iy+oyi+oys][(ith+othi)%n] += \
                                (oxfa)*(oyfa)*(1-othfa)*\
                                math.exp(self.belief[ix][iy][ith])
                    else:
                        newBelief[ix][iy][ith] += \
                                (oxfa)*(oyfa)*(1-othfa)*\
                                math.exp(self.belief[ix][iy][ith])
                    if self.free((ix+oxi+oxs,iy+oyi+oys)):
                        newBelief[ix+oxi+oxs][iy+oyi+oys][(ith+othi+oths)%n]+=\
                                (oxfa)*(oyfa)*(othfa)*\
                                math.exp(self.belief[ix][iy][ith])
                    else:
                        newBelief[ix][iy][ith] += \
                                (oxfa)*(oyfa)*(othfa)*\
                                math.exp(self.belief[ix][iy][ith])
            th = th + self.thetaStep
        mapArray3D(newBelief, math.log)
        self.belief = newBelief




    
