import math
import numpy as np
from Pose import Pose
from Point import Point

np.set_printoptions(precision=3)
np.set_printoptions(suppress=True)
np.set_printoptions(linewidth=100)

def matmult(*x):
  return reduce(np.dot, x)

def wrap_radians(angle):
 return float(((angle + np.pi) % (2*np.pi)) - np.pi)

class EKF(object):

    def __init__(self, controlSigma, controlTurnSigma, sensorDistSigma):
        self.controlSigma = controlSigma
        self.controlTurnSigma = controlTurnSigma
        self.sensorDistSigma = sensorDistSigma
        self.sensorAngleSigma = np.pi/360

        self.X = np.zeros([3,1])
        self.P = np.eye(3,3)

        self.JG = np.eye(3,3)
        self.JXR = np.eye(2,3)
        self.JZ = np.zeros([2,2])

        self.landmarkCount = 0

        self.lastControl = Pose()

    def motionModel(self, control):
        deltaControl = control - self.lastControl

        if deltaControl.x == 0 and deltaControl.y == 0 and deltaControl.r == 0:
            self.JG = np.eye(3,3)
            return 0

        rot1 = wrap_radians(math.atan2(control.y - self.lastControl.y, control.x - self.lastControl.x) - self.lastControl.r)
        trans = math.sqrt(((self.lastControl.x - control.x)**2) + ((self.lastControl.y - control.y)**2))
        rot2 = wrap_radians(control.r - self.lastControl.r - rot1)

        drot1 = ((self.controlTurnSigma*(rot1**2)) + (self.controlSigma*(trans**2)))
        dtrans = ((self.controlTurnSigma*(rot1**2)) + (self.controlTurnSigma*(rot2**2)) + (self.controlSigma*(trans**2)))
        drot2 = ((self.controlTurnSigma*(rot2**2)) + (self.controlSigma*(trans**2)))

        rot1N = wrap_radians(rot1 - (drot1 * np.random.normal()))
        transN = trans - (dtrans * np.random.normal())
        rot2N = wrap_radians(rot2 - (drot2 * np.random.normal()))

        self.X[0] += transN * math.cos(self.X[2] + rot1N)
        self.X[1] += transN * math.sin(self.X[2] + rot1N)
        self.X[2] += rot1N + rot2N

        self.X[2] = wrap_radians(self.X[2])

        self.JG[0][2] = -trans * math.sin(self.lastControl.r + rot1)
        self.JG[1][2] = trans * math.cos(self.lastControl.r + rot1)

        V = np.array([[-trans * math.sin(self.X[2] + rot1), math.cos(self.X[2] + rot1), 0],
                      [trans * math.cos(self.X[2] + rot1), math.sin(self.X[2] + rot1), 0],
                      [1, 0, 1]])

        M = np.diag([drot1, dtrans, drot2])

        self.lastControl = control

        return matmult(V, M, V.T)

    def sensorModel(self, landmark):
        q = math.pow(landmark.x - self.X[0], 2) + math.pow(landmark.y - self.X[1], 2)

        dist = math.sqrt(q)
        angle = math.atan2(landmark.y - self.X[1], landmark.x - self.X[0]) - self.X[2]

        dist += self.sensorDistSigma * np.random.normal()
        angle += self.sensorAngleSigma * np.random.normal()

        return Point(dist, angle), q

    def inverseSensorModel(self, landmark):
        x = self.X[0] + landmark.dist * math.cos(self.X[2] + math.radians(landmark.angle))
        y = self.X[1] + landmark.dist * math.sin(self.X[2] + math.radians(landmark.angle))

        landmark.pos = Point(x,y)

    def predict(self, poseEstimate):
        Q = self.motionModel(poseEstimate)

        self.getRR()[:] = matmult(self.JG, self.getRR(), self.JG.T) + Q
        if self.landmarkCount > 0:
            self.getRM()[:] = matmult(self.JG, self.getRM())
            self.getMR()[:] = self.getRM().T

    def correct(self, reobservedLandmarks):
        for landmark in reobservedLandmarks:
            inv, JH = self.getH(landmark)

            if inv == None:
                continue

            V = np.identity(2)
            R = np.array([[landmark.dist*self.sensorDistSigma, 0],[0, self.sensorAngleSigma]])

            Z = matmult(JH, self.P, JH.T) + matmult(V, R, V.T)
            K = matmult(self.P, JH.T, np.linalg.inv(Z))

            print 'INV:', inv
            print 'Kalman Gain * Inv'
            print matmult(K, inv)

            self.X = self.X - matmult(K, inv)
            self.X[2] = wrap_radians(self.X[2])

            self.P = matmult((np.eye(self.X.shape[0]) - matmult(K, JH)), self.P)
            # self.P = self.P - matmult(K, matmult(self.P, JH.T).T)

            print 'Updating LANDMARK: ',landmark.id
            print self.P

        for lm in reobservedLandmarks:
            lm.pos = self.getLandmark(lm.id)

            eigenvals, eigenvects = np.linalg.eig(self.getLM(lm.id))

            if eigenvals[0] > 0 and eigenvals[1] > 0:
                lm.error = Point(math.sqrt(eigenvals[0]), math.sqrt(eigenvals[1]))
            elif eigenvals[0] < 0 or eigenvals[1] < 0:
                print 'NEG'

    def grow(self, newLandmarks):
        for landmark in newLandmarks:
            self.addLandmark(landmark)

            la = math.radians(landmark.angle)

            self.JXR[0][2] = -landmark.dist * math.sin(la + self.X[2])
            self.JXR[1][2] = landmark.dist * math.cos(la + self.X[2])

            self.JZ[0][0] = math.cos(la + self.X[2])
            self.JZ[0][1] = -landmark.dist * math.sin(la + self.X[2])
            self.JZ[1][0] = math.sin(la + self.X[2])
            self.JZ[1][1] = landmark.dist * math.cos(la + self.X[2])

            R = np.array([[landmark.dist*self.sensorDistSigma, 0],[0, self.sensorAngleSigma]])

            self.getLM(landmark.id)[:] = matmult(self.JXR, self.getRR(), self.JXR.T) + matmult(self.JZ, R, self.JZ.T)

            self.P[3+(self.landmarkCount*2)-2:, 0:-2] = matmult(self.JXR, self.P[0:3, 0:3+(self.landmarkCount*2)-2])
            self.P[0:-2, 3+(self.landmarkCount*2)-2:] = matmult(self.JXR, self.P[0:3, 0:3+(self.landmarkCount*2)-2]).T

            print 'Adding LANDMARK: ', landmark.id
            print self.P

    def getH(self, landmark):
        storedLandmark = self.getLandmark(landmark.id)

        z = Point(landmark.dist, math.radians(landmark.angle))
        h, q = self.sensorModel(storedLandmark)

        if q <= 0:
            return None, None

        z.y = wrap_radians(z.y)
        h.y = wrap_radians(h.y)

        inv = np.array([[z.x-h.x], [wrap_radians(z.y-h.y)]])

        JH = np.zeros([2, 3 + (self.landmarkCount*2)])
        JH[1][2] = -q

        JH[0][0] = self.X[0] - storedLandmark.x / math.sqrt(q)
        JH[0][1] = self.X[1] - storedLandmark.y / math.sqrt(q)
        JH[1][0] = storedLandmark.y - self.X[1] / q
        JH[1][1] = storedLandmark.x - self.X[0] / q

        JH[0][3+(landmark.id*2)] = -JH[0][0]
        JH[0][4+(landmark.id*2)] = -JH[0][1]
        JH[1][3+(landmark.id*2)] = -JH[1][0]
        JH[1][4+(landmark.id*2)] = -JH[1][1]

        return inv, JH

    def addLandmark(self, landmark):
        self.inverseSensorModel(landmark)

        self.X = np.vstack((self.X, landmark.pos.x))
        self.X = np.vstack((self.X, landmark.pos.y))

        self.P = np.hstack((self.P, np.zeros((len(self.P), 2))))
        self.P = np.vstack((self.P, np.zeros((2, len(self.P[0])))))

        self.getLM(self.landmarkCount)[:] = np.array([[0, 0],
                                                      [0, 0]])
        landmark.id = self.landmarkCount
        self.landmarkCount += 1

    def getLandmark(self, lmID):
        return Point(self.X[3+(lmID*2)], self.X[4+(lmID*2)])

    def getRR(self):
        return self.P[0:3, 0:3]

    def getRM(self):
        return self.P[0:3, 3:3+(self.landmarkCount*2)]

    def getMR(self):
        return self.P[3:3+(self.landmarkCount*2), 0:3]

    def getMM(self):
        return self.P[3:3+(self.landmarkCount*2), 3:3+(self.landmarkCount*2)]

    def getLM(self, lmID):
        return self.P[3+(lmID*2):5+(lmID*2), 3+(lmID*2):5+(lmID*2)]

    def getRM_LM(self, lmID):
        return self.P[0:3, 3+(lmID*2):5+(lmID*2)]

    def getMR_LM(self, lmID):
        return self.P[3+(lmID*2):5+(lmID*2), 0:3]

    def getPose(self):
        return Pose(self.X[0], self.X[1], self.X[2])

    def getPoseError(self):
        eigenvals, eigenvects = np.linalg.eig(self.getRR()[0:2,0:2])
        if eigenvals[0] > 0 and eigenvals[1] > 0:
            return Point(math.sqrt(eigenvals[0]), math.sqrt(eigenvals[1]))
        else:
            return Point(1,1)