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

#include "fusion_engine.h"
#include "imu_sensor.h"

#define QUALITY_PI 3.14159265358979323846264338327950288
#define QUALITY_TIMEBASE_HZ 1000000u
#define QUALITY_IMU_HZ 200u
#define QUALITY_DURATION_S 30u

typedef struct {
	uint32_t state;
} quality_rng_t;

typedef struct {
	double position_m[3];
	double velocity_mps[3];
	double yaw_rad;
} quality_truth_t;

typedef struct {
	const char* name;
	double horizontal_rmse_m;
	double final_horizontal_error_m;
	double final_speed_mps;
	double outage_peak_error_m;
	double first_accel_detection_s;
	double rotation_detection_s;
	double final_stationary_reentry_s;
	uint64_t zero_updates;
	uint64_t accepted_gps;
	unsigned int dynamic_zero_update_count;
	unsigned int false_moving_before_motion;
	bool state_finite;
} quality_metrics_t;

static int failures;

#define CHECK(condition, message) do { \
	if(!(condition)) { \
		fprintf(stderr, "FAIL: %s\n", message); \
		failures++; \
	} \
} while(0)

static double square(double value)
{
	return value * value;
}

static double norm2(double first, double second)
{
	return sqrt(first * first + second * second);
}

static uint32_t rng_next(quality_rng_t* rng)
{
	uint32_t value = rng->state;
	value ^= value << 13;
	value ^= value >> 17;
	value ^= value << 5;
	rng->state = value;
	return value;
}

static double rng_uniform(quality_rng_t* rng)
{
	return ((double)rng_next(rng) + 0.5) / 4294967296.0;
}

static double rng_gaussian(quality_rng_t* rng)
{
	double first = fmax(rng_uniform(rng), 1.0e-12);
	double second = rng_uniform(rng);
	return sqrt(-2.0 * log(first)) * cos(2.0 * QUALITY_PI * second);
}

static void trajectory_input(double time_s,
	                         double acceleration_nav_mps2[3],
	                         double* yaw_rate_rps)
{
	memset(acceleration_nav_mps2, 0, 3u * sizeof(double));
	*yaw_rate_rps = 0.0;
	if((time_s >= 3.0) && (time_s < 4.0))
	{
		acceleration_nav_mps2[0] = 2.0;
	}
	else if((time_s >= 6.0) && (time_s < 7.0))
	{
		acceleration_nav_mps2[0] = -2.0;
	}
	else if((time_s >= 12.0) && (time_s < 13.0))
	{
		*yaw_rate_rps = 0.5 * QUALITY_PI;
	}
	else if((time_s >= 17.0) && (time_s < 17.5))
	{
		acceleration_nav_mps2[1] = -3.0;
	}
	else if((time_s >= 19.0) && (time_s < 19.5))
	{
		acceleration_nav_mps2[1] = 3.0;
	}
}

static bool truth_is_dynamic(const quality_truth_t* truth,
	                         const double acceleration_nav_mps2[3],
	                         double yaw_rate_rps)
{
	return (fabs(acceleration_nav_mps2[0]) > 0.4) ||
	       (fabs(acceleration_nav_mps2[1]) > 0.4) ||
	       (fabs(yaw_rate_rps) > 0.02) ||
	       (norm2(truth->velocity_mps[0], truth->velocity_mps[1]) > 0.20);
}

static void truth_step(quality_truth_t* truth, double dt_s,
	                   const double acceleration_nav_mps2[3],
	                   double yaw_rate_rps)
{
	for(unsigned int axis = 0u ; axis < 3u ; axis++)
	{
		truth->position_m[axis] += truth->velocity_mps[axis] * dt_s +
		                            0.5 * acceleration_nav_mps2[axis] * dt_s * dt_s;
		truth->velocity_mps[axis] += acceleration_nav_mps2[axis] * dt_s;
	}
	truth->yaw_rad += yaw_rate_rps * dt_s;
}

static void publish_position(fusion_engine_t* engine, uint64_t tick,
	                         const quality_truth_t* truth,
	                         quality_rng_t* gps_rng)
{
	fusion_position_measurement_t gps;
	memset(&gps, 0, sizeof(gps));
	gps.timestamp_ticks = tick;
	gps.position_enu_m[0] = truth->position_m[0] + 1.5 * rng_gaussian(gps_rng);
	gps.position_enu_m[1] = truth->position_m[1] + 1.5 * rng_gaussian(gps_rng);
	gps.position_enu_m[2] = truth->position_m[2] + 2.5 * rng_gaussian(gps_rng);
	gps.covariance_enu_m2[0] = square(1.5);
	gps.covariance_enu_m2[4] = square(1.5);
	gps.covariance_enu_m2[8] = square(2.5);
	gps.source_quality = 1.0;
	gps.source_valid = true;
	eskf_update_result_t result = ESKF_UPDATE_INVALID;
	CHECK(fusion_engine_push_position(engine, &gps, &result, NULL),
	      "quality-profile GPS is processed");
	CHECK((result == ESKF_UPDATE_ACCEPTED) ||
	      (result == ESKF_UPDATE_REJECTED_NIS),
	      "quality-profile GPS returns a defined update result");
}

static quality_metrics_t run_profile(imu_sensor_backend_t backend,
	                                 uint32_t seed)
{
	static fusion_engine_t engine;
	fusion_imu_quality_profile_t profile;
	CHECK(imu_sensor_quality_profile_for_backend(backend, &profile),
	      "sensor fusion quality profile exists");
	fusion_engine_config_t config;
	fusion_engine_default_config(&config, QUALITY_TIMEBASE_HZ);
	CHECK(fusion_engine_apply_imu_quality(&config, &profile),
	      "sensor quality profile applies to ESKF config");
	config.calibration_duration_s = 2.0;
	config.calibration_min_samples = QUALITY_IMU_HZ;
	CHECK(fusion_engine_init(&engine, &config),
	      "quality-profile fusion engine initializes");

	fusion_position_measurement_t origin;
	memset(&origin, 0, sizeof(origin));
	origin.covariance_enu_m2[0] = square(1.5);
	origin.covariance_enu_m2[4] = square(1.5);
	origin.covariance_enu_m2[8] = square(2.5);
	origin.source_quality = 1.0;
	origin.source_valid = true;
	CHECK(fusion_engine_push_position(&engine, &origin, NULL, NULL),
	      "quality-profile initial GPS is retained");

	quality_rng_t imu_rng = {0x1735a9cdu ^ seed};
	quality_rng_t gps_rng = {0x9b6d42e1u ^ seed};
	quality_truth_t truth = {{0.0, 0.0, 0.0}, {0.0, 0.0, 0.0}, 0.0};
	double accel_bias[3] = {0.030, -0.020, 0.050};
	double gyro_bias[3] = {0.0008, -0.0005, 0.0007};
	double quality_factor = (backend == IMU_SENSOR_ST_LSM6DSOX) ? 1.35 : 1.0;
	for(unsigned int axis = 0u ; axis < 3u ; axis++)
	{
		accel_bias[axis] *= quality_factor;
		gyro_bias[axis] *= quality_factor;
	}

	const uint64_t imu_period_ticks = QUALITY_TIMEBASE_HZ / QUALITY_IMU_HZ;
	const double dt_s = 1.0 / (double)QUALITY_IMU_HZ;
	const double bandwidth_sqrt_hz = sqrt(0.5 * (double)QUALITY_IMU_HZ);
	quality_metrics_t metrics;
	memset(&metrics, 0, sizeof(metrics));
	metrics.name = profile.name;
	metrics.first_accel_detection_s = NAN;
	metrics.rotation_detection_s = NAN;
	metrics.final_stationary_reentry_s = NAN;
	double squared_horizontal_error_sum = 0.0;
	uint64_t error_samples = 0u;
	uint64_t next_gps_tick = 3000000u;
	uint64_t previous_zero_updates = 0u;

	for(uint64_t tick = 0u ;
	    tick <= (uint64_t)QUALITY_DURATION_S * QUALITY_TIMEBASE_HZ ;
	    tick += imu_period_ticks)
	{
		double acceleration_nav[3] = {0.0, 0.0, 0.0};
		double yaw_rate = 0.0;
		double time_s = (double)tick / (double)QUALITY_TIMEBASE_HZ;
		if(tick != 0u)
		{
			double midpoint_s = time_s - 0.5 * dt_s;
			trajectory_input(midpoint_s, acceleration_nav, &yaw_rate);
			truth_step(&truth, dt_s, acceleration_nav, yaw_rate);
		}

		/* Inject the profile's bias random walk into the synthetic sensor. */
		for(unsigned int axis = 0u ; axis < 3u ; axis++)
		{
			accel_bias[axis] += profile.accel_bias_random_walk_mps2_sqrt_s[axis] *
			                    sqrt(dt_s) * rng_gaussian(&imu_rng);
			gyro_bias[axis] += profile.gyro_bias_random_walk_rps_sqrt_s[axis] *
			                   sqrt(dt_s) * rng_gaussian(&imu_rng);
		}

		double midpoint_yaw = truth.yaw_rad - 0.5 * yaw_rate * dt_s;
		double cosine = cos(midpoint_yaw);
		double sine = sin(midpoint_yaw);
		fusion_imu_measurement_t imu;
		memset(&imu, 0, sizeof(imu));
		imu.timestamp_ticks = tick;
		imu.accel_mps2[0] = cosine * acceleration_nav[0] +
		                    sine * acceleration_nav[1];
		imu.accel_mps2[1] = -sine * acceleration_nav[0] +
		                     cosine * acceleration_nav[1];
		imu.accel_mps2[2] = config.eskf.gravity_mps2;
		imu.gyro_rps[2] = yaw_rate;
		for(unsigned int axis = 0u ; axis < 3u ; axis++)
		{
			double accel_sigma =
				3.0 * profile.nominal_accel_noise_density_mps2_sqrt_hz[axis] *
				bandwidth_sqrt_hz;
			double gyro_sigma =
				3.0 * profile.nominal_gyro_noise_density_rps_sqrt_hz[axis] *
				bandwidth_sqrt_hz;
			imu.accel_mps2[axis] += accel_bias[axis] +
			                         accel_sigma * rng_gaussian(&imu_rng);
			imu.gyro_rps[axis] += gyro_bias[axis] +
			                      gyro_sigma * rng_gaussian(&imu_rng);
		}
		imu.temperature_c = 25.0 + 0.05 * time_s;
		imu.sequence = (uint32_t)(tick / imu_period_ticks);
		CHECK(fusion_engine_push_imu(&engine, &imu),
		      "quality-profile IMU is accepted");

		if((tick >= next_gps_tick) &&
		   !((time_s >= 16.0) && (time_s < 25.0)))
		{
			publish_position(&engine, tick, &truth, &gps_rng);
			next_gps_tick += QUALITY_TIMEBASE_HZ;
		}
		else if((tick >= next_gps_tick) && (time_s >= 25.0))
		{
			/* Skip the 16--25 s events, then resume at the current epoch. */
			publish_position(&engine, tick, &truth, &gps_rng);
			next_gps_tick = tick + QUALITY_TIMEBASE_HZ;
		}

		if(engine.state != FUSION_ENGINE_RUNNING)
		{
			continue;
		}
		bool dynamic = truth_is_dynamic(&truth,
		                                acceleration_nav, yaw_rate);
		uint64_t zero_updates = engine.diagnostics.zero_velocity_update_count;
		if(dynamic && (zero_updates > previous_zero_updates))
		{
			metrics.dynamic_zero_update_count +=
				(unsigned int)(zero_updates - previous_zero_updates);
		}
		previous_zero_updates = zero_updates;
		if((time_s > 2.2) && (time_s < 3.0) &&
		   (engine.current.motion.state != FUSION_MOTION_STATIONARY))
		{
			metrics.false_moving_before_motion++;
		}
		if(isnan(metrics.first_accel_detection_s) && (time_s >= 3.0) &&
		   (engine.current.motion.state == FUSION_MOTION_MOVING))
		{
			metrics.first_accel_detection_s = time_s;
		}
		if(isnan(metrics.rotation_detection_s) && (time_s >= 12.0) &&
		   (engine.current.motion.state == FUSION_MOTION_MOVING))
		{
			metrics.rotation_detection_s = time_s;
		}
		if(isnan(metrics.final_stationary_reentry_s) && (time_s >= 19.5) &&
		   (engine.current.motion.state == FUSION_MOTION_STATIONARY))
		{
			metrics.final_stationary_reentry_s = time_s;
		}

		double horizontal_error = norm2(
			engine.current.filter.position_m[0] - truth.position_m[0],
			engine.current.filter.position_m[1] - truth.position_m[1]);
		if(time_s >= 3.0)
		{
			squared_horizontal_error_sum += square(horizontal_error);
			error_samples++;
		}
		if((time_s >= 16.0) && (time_s < 25.0))
		{
			metrics.outage_peak_error_m =
				fmax(metrics.outage_peak_error_m, horizontal_error);
		}
	}

	metrics.horizontal_rmse_m =
		(error_samples != 0u)
		? sqrt(squared_horizontal_error_sum / (double)error_samples)
		: NAN;
	metrics.final_horizontal_error_m = norm2(
		engine.current.filter.position_m[0] - truth.position_m[0],
		engine.current.filter.position_m[1] - truth.position_m[1]);
	metrics.final_speed_mps = eskf_speed_mps(&engine.current.filter);
	metrics.zero_updates = engine.diagnostics.zero_velocity_update_count;
	metrics.accepted_gps = engine.diagnostics.position_accepted;
	metrics.state_finite = eskf_state_is_finite(&engine.current.filter);
	return metrics;
}

static void check_metrics(const quality_metrics_t* metrics,
	                      bool noisier_st_profile)
{
	CHECK(metrics->state_finite, "quality-profile ESKF state remains finite");
	CHECK(metrics->false_moving_before_motion == 0u,
	      "stationary sensor noise does not create false motion");
	CHECK(metrics->dynamic_zero_update_count == 0u,
	      "ZUPT and zero-gyro are disabled throughout true motion");
	CHECK(isfinite(metrics->first_accel_detection_s) &&
	      (metrics->first_accel_detection_s <= 3.010),
	      "sudden acceleration is detected within two samples");
	CHECK(isfinite(metrics->rotation_detection_s) &&
	      (metrics->rotation_detection_s <= 12.010),
	      "sudden rotation is detected within two samples");
	CHECK(isfinite(metrics->final_stationary_reentry_s) &&
	      (metrics->final_stationary_reentry_s >= 20.5) &&
	      (metrics->final_stationary_reentry_s <=
	       (noisier_st_profile ? 29.5 : 25.0)),
	      "stationary re-entry waits for motion hysteresis and then recovers");
	CHECK(metrics->zero_updates > (noisier_st_profile ? 100u : 150u),
	      "stationary intervals apply rate-limited zero updates");
	CHECK(metrics->accepted_gps >= 15u,
	      "GPS updates remain usable with the sensor profile");
	CHECK(metrics->horizontal_rmse_m < (noisier_st_profile ? 5.0 : 2.0),
	      "sensor-profile horizontal RMSE remains bounded");
	CHECK(metrics->outage_peak_error_m < (noisier_st_profile ? 12.0 : 3.0),
	      "nine-second GPS outage drift remains bounded");
	CHECK(metrics->final_horizontal_error_m < 2.0,
	      "post-outage final horizontal error remains bounded");
	CHECK(metrics->final_speed_mps < 0.05,
	      "final stationary speed is near zero");
}

int main(void)
{
	fusion_imu_quality_profile_t tdk_profile;
	fusion_imu_quality_profile_t st_profile;
	CHECK(imu_sensor_quality_profile_for_backend(IMU_SENSOR_TDK_ICM42688,
	                                             &tdk_profile),
	      "TDK profile is valid");
	CHECK(imu_sensor_quality_profile_for_backend(IMU_SENSOR_ST_LSM6DSOX,
	                                             &st_profile),
	      "ST profile is valid");
	CHECK(tdk_profile.nominal_gyro_noise_density_rps_sqrt_hz[0] <
	      st_profile.nominal_gyro_noise_density_rps_sqrt_hz[0],
	      "TDK profile models lower nominal gyro noise than ST");
	CHECK(tdk_profile.process_gyro_noise_density_rps_sqrt_hz[0] <
	      st_profile.process_gyro_noise_density_rps_sqrt_hz[0],
	      "TDK profile gives the ESKF lower gyro process noise");

	quality_metrics_t tdk = run_profile(IMU_SENSOR_TDK_ICM42688, 0x11111111u);
	quality_metrics_t st = run_profile(IMU_SENSOR_ST_LSM6DSOX, 0x11111111u);
	check_metrics(&tdk, false);
	check_metrics(&st, true);
	CHECK(tdk.outage_peak_error_m <= st.outage_peak_error_m + 0.25,
	      "TDK-like profile is no worse than ST-like profile during GPS outage");

	printf("Synthetic IMU quality-profile comparison (same production C ESKF)\n");
	printf("  %-16s RMSE=%5.3f m outage-peak=%5.3f m final=%5.3f m speed=%7.4f m/s GPS=%llu\n",
	       tdk.name, tdk.horizontal_rmse_m, tdk.outage_peak_error_m,
	       tdk.final_horizontal_error_m, tdk.final_speed_mps,
	       (unsigned long long)tdk.accepted_gps);
	printf("  %-16s RMSE=%5.3f m outage-peak=%5.3f m final=%5.3f m speed=%7.4f m/s GPS=%llu\n",
	       st.name, st.horizontal_rmse_m, st.outage_peak_error_m,
	       st.final_horizontal_error_m, st.final_speed_mps,
	       (unsigned long long)st.accepted_gps);
	printf("  detection TDK accel=%.3f s turn=%.3f s stop=%.3f s; ST accel=%.3f s turn=%.3f s stop=%.3f s\n",
	       tdk.first_accel_detection_s, tdk.rotation_detection_s,
	       tdk.final_stationary_reentry_s, st.first_accel_detection_s,
	       st.rotation_detection_s, st.final_stationary_reentry_s);
	printf("  dynamic ZUPT violations TDK=%u ST=%u; zero-updates TDK=%llu ST=%llu\n",
	       tdk.dynamic_zero_update_count, st.dynamic_zero_update_count,
	       (unsigned long long)tdk.zero_updates,
	       (unsigned long long)st.zero_updates);
	if(failures != 0)
	{
		fprintf(stderr, "%d IMU quality-profile check(s) failed\n", failures);
		return 1;
	}
	printf("PASS: both TDK-like and noisier ST-like streams remain bounded.\n");
	return 0;
}