import numpy as np
from kalman_helpers import ObservationKind
from ekf_sym import EKF_sym
from laika.raw_gnss import GNSSMeasurement
def parse_prr(m):
	sat_pos_vel_i = np.concatenate((m[GNSSMeasurement.SAT_POS], m[GNSSMeasurement.SAT_VEL]))
	R_i = np.atleast_2d(m[GNSSMeasurement.PRR_STD] ** 2)
	z_i = m[GNSSMeasurement.PRR]
	return z_i, R_i, sat_pos_vel_i
def parse_pr(m):
	pseudorange = m[GNSSMeasurement.PR]
	pseudorange_stdev = m[GNSSMeasurement.PR_STD]
	sat_pos_freq_i = np.concatenate((m[GNSSMeasurement.SAT_POS], np.array([m[GNSSMeasurement.GLONASS_FREQ]])))
	z_i = np.atleast_1d(pseudorange)
	R_i = np.atleast_2d(pseudorange_stdev ** 2)
	return z_i, R_i, sat_pos_freq_i
class States(object):
	ECEF_POS = slice(0, 3)
	ECEF_VELOCITY = slice(3, 6)
	CLOCK_BIAS = slice(6, 7)
	CLOCK_DRIFT = slice(7, 8)
	CLOCK_ACCELERATION = slice(8, 9)
	GLONASS_BIAS = slice(9, 10)
	GLONASS_FREQ_SLOPE = slice(10, 11)
class GNSSKalman(object):
	def __init__(self, N=0, max_tracks=3000):
		x_initial = np.array([-2712700.6008, -4281600.6679, 3859300.1830, 0, 0, 0, 0, 0, 0, 0, 0])
		P_initial = np.diag([10000 ** 2, 10000 ** 2, 10000 ** 2, 10 ** 2, 10 ** 2, 10 ** 2, (2000000) ** 2, (100) ** 2, (0.5) ** 2, (10) ** 2, (1) ** 2])
		Q = np.diag([0.3 ** 2, 0.3 ** 2, 0.3 ** 2, 3 ** 2, 3 ** 2, 3 ** 2, (0.1) ** 2, (0) ** 2, (0.01) ** 2, 0.1 ** 2, (0.01) ** 2])
		self.dim_state = x_initial.shape[0]
		maha_test_kinds = []
		name = "gnss"
		self.filter = EKF_sym(name, Q, x_initial, P_initial, self.dim_state, self.dim_state, maha_test_kinds=maha_test_kinds)
	@property
	def x(self):
		return self.filter.state()
	@property
	def P(self):
		return self.filter.covs()
	def predict(self, t):
		return self.filter.predict(t)
	def rts_smooth(self, estimates):
		return self.filter.rts_smooth(estimates, norm_quats=False)
	def init_state(self, state, covs_diag=None, covs=None, filter_time=None):
		if covs_diag is not None:
			P = np.diag(covs_diag)
		elif covs is not None:
			P = covs
		else:
			P = self.filter.covs()
		self.filter.init_state(state, P, filter_time)
	def predict_and_observe(self, t, kind, data):
		if len(data) > 0:
			data = np.atleast_2d(data)
		if kind == ObservationKind.PSEUDORANGE_GPS or kind == ObservationKind.PSEUDORANGE_GLONASS:
			r = self.predict_and_update_pseudorange(data, t, kind)
		elif kind == ObservationKind.PSEUDORANGE_RATE_GPS or kind == ObservationKind.PSEUDORANGE_RATE_GLONASS:
			r = self.predict_and_update_pseudorange_rate(data, t, kind)
		return r
	def predict_and_update_pseudorange(self, meas, t, kind):
		R = np.zeros((len(meas), 1, 1))
		sat_pos_freq = np.zeros((len(meas), 4))
		z = np.zeros((len(meas), 1))
		for i, m in enumerate(meas):
			z_i, R_i, sat_pos_freq_i = parse_pr(m)
			sat_pos_freq[i, :] = sat_pos_freq_i
			z[i, :] = z_i
			R[i, :, :] = R_i
		return self.filter.predict_and_update_batch(t, kind, z, R, sat_pos_freq)
	def predict_and_update_pseudorange_rate(self, meas, t, kind):
		R = np.zeros((len(meas), 1, 1))
		z = np.zeros((len(meas), 1))
		sat_pos_vel = np.zeros((len(meas), 6))
		for i, m in enumerate(meas):
			z_i, R_i, sat_pos_vel_i = parse_prr(m)
			sat_pos_vel[i] = sat_pos_vel_i
			R[i, :, :] = R_i
			z[i, :] = z_i
		return self.filter.predict_and_update_batch(t, kind, z, R, sat_pos_vel)
if __name__ == "__main__":
	GNSSKalman()