#include "fusion/fusion_engine.h"

#include <math.h>
#include <stdbool.h>
#include <stdint.h>
#include <stdio.h>
#include <string.h>

#define TEST_TIMEBASE_HZ 1000000u
#define TEST_IMU_HZ 52u
#define TEST_GPS_PERIOD_TICKS 200000u

static fusion_imu_measurement_t stationary_imu(uint64_t timestamp_ticks,
	                                            double gravity_mps2)
{
	fusion_imu_measurement_t imu;
	memset(&imu, 0, sizeof(imu));
	imu.timestamp_ticks = timestamp_ticks;
	imu.accel_mps2[2] = gravity_mps2;
	imu.temperature_c = 25.0;
	imu.vehicle_at_rest = true;
	return imu;
}

static fusion_position_measurement_t gps_position(
	uint64_t timestamp_ticks,
	double east_m,
	double north_m,
	double up_m)
{
	fusion_position_measurement_t gps;
	memset(&gps, 0, sizeof(gps));
	gps.timestamp_ticks = timestamp_ticks;
	gps.position_enu_m[0] = east_m;
	gps.position_enu_m[1] = north_m;
	gps.position_enu_m[2] = up_m;
	gps.covariance_enu_m2[0] = 4.0;
	gps.covariance_enu_m2[4] = 4.0;
	gps.covariance_enu_m2[8] = 16.0;
	gps.source_quality = 100.0;
	gps.source_valid = true;
	return gps;
}

static bool calibrate_at_origin(fusion_engine_t* engine,
	                            fusion_engine_config_t* config)
{
	fusion_engine_default_config(config, TEST_TIMEBASE_HZ);
	if(!fusion_engine_init(engine, config))
	{
		return false;
	}
	fusion_position_measurement_t origin = gps_position(1u, 0.0, 0.0, 0.0);
	if(!fusion_engine_push_position(engine, &origin, NULL, NULL))
	{
		return false;
	}
	for(uint64_t sample = 0u ; sample <= 3u * TEST_IMU_HZ ; sample++)
	{
		uint64_t ticks = sample * TEST_TIMEBASE_HZ / TEST_IMU_HZ;
		fusion_imu_measurement_t imu = stationary_imu(ticks, config->eskf.gravity_mps2);
		if(!fusion_engine_push_imu(engine, &imu))
		{
			return false;
		}
	}
	return engine->state == FUSION_ENGINE_RUNNING;
}

static int startup_offset_recovers(uint64_t gps_period_ticks)
{
	static fusion_engine_t engine;
	fusion_engine_config_t config;
	if(!calibrate_at_origin(&engine, &config))
	{
		fprintf(stderr, "FAIL: calibration did not complete\n");
		return 1;
	}

	/* Models a poor first LLH origin followed by the stable, correct fixes seen
	 * on UART.  ENU remains relative to the first origin, so the correct fixes
	 * legitimately appear far from the initialized ESKF position. */
	const double expected[3] = {80.0, 40.0, -170.0};
	uint64_t next_gps_ticks = 3200000u;
	for(uint64_t sample = 3u * TEST_IMU_HZ + 1u ;
	    sample <= 12u * TEST_IMU_HZ ; sample++)
	{
		uint64_t ticks = sample * TEST_TIMEBASE_HZ / TEST_IMU_HZ;
		fusion_imu_measurement_t imu = stationary_imu(ticks, config.eskf.gravity_mps2);
		if(!fusion_engine_push_imu(&engine, &imu))
		{
			fprintf(stderr, "FAIL: IMU propagation failed\n");
			return 1;
		}
		while(next_gps_ticks <= ticks)
		{
			fusion_position_measurement_t gps = gps_position(
				next_gps_ticks, expected[0], expected[1], expected[2]);
			(void)fusion_engine_push_position(&engine, &gps, NULL, NULL);
			next_gps_ticks += gps_period_ticks;
		}
	}

	double error = sqrt(
		pow(engine.current.filter.position_m[0] - expected[0], 2.0) +
		pow(engine.current.filter.position_m[1] - expected[1], 2.0) +
		pow(engine.current.filter.position_m[2] - expected[2], 2.0));
	printf("startup offset GPS=%4llu ms: error=%.3f m accepted=%llu "
	       "rejected=%llu recovery=%llu\n",
	       (unsigned long long)(gps_period_ticks / 1000u), error,
	       (unsigned long long)engine.diagnostics.position_accepted,
	       (unsigned long long)engine.diagnostics.position_nis_rejected,
	       (unsigned long long)engine.diagnostics.position_recovery_count);
	if((error > 5.0) || (engine.diagnostics.position_recovery_count != 1u))
	{
		fprintf(stderr, "FAIL: consistent GPS did not recover the wrong startup position\n");
		return 1;
	}
	return 0;
}

static int isolated_outlier_does_not_reset(void)
{
	static fusion_engine_t engine;
	fusion_engine_config_t config;
	if(!calibrate_at_origin(&engine, &config))
	{
		fprintf(stderr, "FAIL: outlier scenario calibration did not complete\n");
		return 1;
	}

	for(uint64_t gps_ticks = 3200000u ; gps_ticks <= 7000000u ;
	    gps_ticks += TEST_GPS_PERIOD_TICKS)
	{
		uint64_t imu_ticks = gps_ticks + 10000u;
		fusion_imu_measurement_t imu = stationary_imu(
			imu_ticks, config.eskf.gravity_mps2);
		if(!fusion_engine_push_imu(&engine, &imu))
		{
			fprintf(stderr, "FAIL: outlier scenario IMU propagation failed\n");
			return 1;
		}
		bool outlier = (gps_ticks == 5000000u);
		fusion_position_measurement_t gps = gps_position(
			gps_ticks, outlier ? 150.0 : 0.0, outlier ? -120.0 : 0.0,
			outlier ? 80.0 : 0.0);
		(void)fusion_engine_push_position(&engine, &gps, NULL, NULL);
	}

	double distance = sqrt(
		pow(engine.current.filter.position_m[0], 2.0) +
		pow(engine.current.filter.position_m[1], 2.0) +
		pow(engine.current.filter.position_m[2], 2.0));
	printf("single outlier: error=%.3f m rejected=%llu recovery=%llu\n",
	       distance,
	       (unsigned long long)engine.diagnostics.position_nis_rejected,
	       (unsigned long long)engine.diagnostics.position_recovery_count);
	if((distance > 1.0) || (engine.diagnostics.position_nis_rejected == 0u) ||
	   (engine.diagnostics.position_recovery_count != 0u))
	{
		fprintf(stderr, "FAIL: isolated GPS outlier moved/reset the ESKF\n");
		return 1;
	}
	return 0;
}

int main(void)
{
	int failures = startup_offset_recovers(100000u);
	failures += startup_offset_recovers(200000u);
	failures += startup_offset_recovers(1000000u);
	failures += isolated_outlier_does_not_reset();
	if(failures == 0)
	{
		puts("GPS reacquisition tests passed");
	}
	return failures == 0 ? 0 : 1;
}