from mpmath import *
import rand

# Matrices are labelled in the notation from
# http://en.wikipedia.org/wiki/Kalman_filter#Underlying_dynamic_system_model

# If using a implementing a non-linear motion or observation model, then for those
# functions which return a matrix, you'll need to find a linear approximation to the
# true model.  This linear approximation is usually found by using the tangent line to
# the non-linear function.  The slope of a tangent line to a function is the derivative of
# the function at that point.
# See http://en.wikipedia.org/wiki/Kalman_filter#Extended_Kalman_filter or, for example,
# http://www.acsu.buffalo.edu/~terejanu/files/tutorialEKF.pdf

# For a system with more than one dimension, this means:
# <complete this comment>

def getStateDimensions():
	return 2

def getObservationDimensions():
	return 1

def motionMatrixF(state, action):
	"Return the F matrix (using wikipedia notation)"
	return matrix([[1.0, 1.0],[0.0,1.0]])

def motionMatrixB(state, action):
	"Return the B matrix (using wikipedia notation)"
	return matrix([0.0,1.0])

def motionNoiseCoVar(state, action):
	"Return the Q matrix (using wikipedia notation)"
	return eye(2) + matrix([[0,0],[0,1]]) * action[0]

def expectedMotionResult(state, action):
	"Return the expected state after executing an action in a state"
	# for a linear filter: F.x + B.u
	return motionMatrixF(state, action) * state + motionMatrixB(state, action) * action

def motionSample(state, action):
	"Return a sample from the motion model if the agent started in the supplied state and performed the supplied action"
	return rand.mv_normal_sample(expectedMotionResult(state, action), motionNoiseCoVar(state, action))

def observationMatrix(state, obs):
	"Return the H matrix (using wikipedia notation)"
	return matrix([[1,0]])

def observationNoiseCoVar(state, obs):
	"Return the R matrix (using wikipedia notation)"
	return eye(1)

def expectedObservation(state, obs):
	"Return the mean observation for a given state"
	# Note, an observation is given.  If there is a discrete component to the observation, then match the supplied observation
	return observationMatrix(state, obs)*state

def observationProbability(state, obs):
	"Return the probability, or probability density as appropriate, of the supplied observation assuming we're in the supplied state"
	return rand.mv_normal_val(expectedObservation(state, obs), observationNoiseCoVar(state, obs), obs)

def observationSample(state):
	"Return a sample observation for this state."
	return rand.mv_normal_sample(expectedObservation(state, None), observationNoiseCoVar(state, None))

def normaliseStateAngles(state):
	"Make sure any angles in the state are between -pi and +pi"
	for i in []:	# list angles in the state here
		while state[i] > math.pi:
			state[i] = state[i] - 2*math.pi
		while state[i] < -math.pi:
			state[i] = state[i] + 2*math.pi
	return state

def normaliseObsAngles(obs):
	"Make sure any angles in the observation are between -pi and +pi"
	for i in []:	# list angles in the observation here
		while obs[i] > math.pi:
			obs[i] = obs[i] - 2*math.pi
		while obs[i] < -math.pi:
			obs[i] = obs[i] + 2*math.pi
	return obs