from model import motionMatrixF, motionMatrixB, motionNoiseCoVar, observationMatrix, observationNoiseCoVar, getStateDimensions, getObservationDimensions, expectedObservation, expectedMotionResult, normaliseStateAngles, normaliseObsAngles
#from robotModel import motionMatrixF, motionMatrixB, motionNoiseCoVar, observationMatrix, observationNoiseCoVar, getStateDimensions, getObservationDimensions, expectedObservation, expectedMotionResult, normaliseStateAngles, normaliseObsAngles
from mpmath import *
import rand

largeNum = 10e5

mean = zeros(getStateDimensions(),1)
Sigma = eye(getStateDimensions())

def setInitialStateUniform():
	global mean
	global Sigma
	
	mean = zeros(getStateDimensions(),1)
	Sigma = eye(getStateDimensions())*largeNum

def setInitialStateGaussian(lmean, lSigma):
	global mean
	global Sigma
	mean = lmean.copy()
	Sigma = lSigma.copy()

def getMean():
	global mean
	global Sigma
	return mean

def getSigma():
	global mean
	global Sigma
	return Sigma

def observationUpdate(obs):
	global mean
	global Sigma
	H = observationMatrix(mean, obs)
	R = observationNoiseCoVar(mean, obs)
	# print "H:", H
	# print "R:", R
	# y = obs - H * mean
	# y = obs - expectedObservation(mean, obs)
	y = normaliseObsAngles(obs - expectedObservation(mean, obs))
	# print "y:", y
	S = H * Sigma * H.T + R
	# print "S:", S
	K = Sigma * H.T * inverse(S)
	# print "K:", K
	mean = mean + K*y
	# print "KH:", (K*H)
	# print "I - KH:", (eye(mean.rows) - K * H)
	Sigma = (eye(mean.rows) - K * H)*Sigma
	mean = normaliseStateAngles(mean)

def motionUpdate(action):
	global mean
	global Sigma
	F = motionMatrixF(mean, action)
	B = motionMatrixB(mean, action)
	Q = motionNoiseCoVar(mean, action)
	# print "F:", F
	# print "Q:", Q
	mean = normaliseStateAngles(expectedMotionResult(mean, action))
	Sigma = F * Sigma * F.T + Q