#include "fusion/fusion_engine.h"

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

#define TIMEBASE_HZ 1000000u
#define IMU_HZ 52u
#define DURATION_S 25.0
#define PI 3.14159265358979323846264338327950288
#define GPS_EVENT_CAPACITY 300u

typedef struct {
	double position_m[3];
	double velocity_mps[3];
	double yaw_rad;
	double accel_nav_mps2[3];
	double yaw_rate_rps;
} truth_t;

typedef struct {
	fusion_position_measurement_t measurement;
	uint64_t arrival_ticks;
} gps_event_t;

typedef struct {
	unsigned int gps_period_ms;
	double horizontal_rmse_m;
	double peak_outage_error_m;
	double final_horizontal_error_m;
	double final_vertical_error_m;
	double final_speed_mps;
	double outage_sigma_growth_m;
	double final_horizontal_sigma_m;
	double maximum_position_step_m;
	uint64_t gps_accepted;
	uint64_t gps_rejected;
	uint64_t replay_count;
	uint64_t imu_received;
	unsigned int false_stationary_samples;
	bool moving_detected;
	bool recovery_gps_accepted;
	bool local_state_finite;
	bool covariance_finite_positive;
	bool global_origin_available;
} metrics_t;

static bool trace_enabled;

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

static void command_at(double time_s, double accel[3], double* yaw_rate)
{
	memset(accel, 0, 3u * sizeof(double));
	*yaw_rate = 0.0;
	if((time_s >= 3.0) && (time_s < 5.0))
	{
		accel[0] = 1.0;
	}
	else if((time_s >= 8.0) && (time_s < 10.0))
	{
		accel[0] = -1.0;
	}
	else if((time_s >= 12.0) && (time_s < 14.0))
	{
		accel[1] = 0.5;
		*yaw_rate = PI / 4.0;
	}
	else if((time_s >= 18.0) && (time_s < 20.0))
	{
		accel[1] = -0.5;
	}
}

static truth_t truth_at(double time_s)
{
	static const double boundaries[] = {
		0.0, 3.0, 5.0, 8.0, 10.0, 12.0, 14.0, 18.0, 20.0, DURATION_S
	};
	truth_t truth;
	memset(&truth, 0, sizeof(truth));
	for(unsigned int segment = 0u;
	    segment + 1u < sizeof(boundaries) / sizeof(boundaries[0]); segment++)
	{
		double start = boundaries[segment];
		if(time_s <= start)
		{
			break;
		}
		double end = fmin(time_s, boundaries[segment + 1u]);
		double dt = end - start;
		double accel[3];
		double yaw_rate;
		command_at(start + 0.5 * dt, accel, &yaw_rate);
		for(unsigned int axis = 0u; axis < 3u; axis++)
		{
			truth.position_m[axis] += truth.velocity_mps[axis] * dt +
			                           0.5 * accel[axis] * dt * dt;
			truth.velocity_mps[axis] += accel[axis] * dt;
		}
		truth.yaw_rad += yaw_rate * dt;
		if(end >= time_s)
		{
			break;
		}
	}
	command_at(time_s, truth.accel_nav_mps2, &truth.yaw_rate_rps);
	return truth;
}

static fusion_imu_measurement_t make_imu(uint64_t ticks, double gravity)
{
	double time_s = (double)ticks / (double)TIMEBASE_HZ;
	truth_t truth = truth_at(time_s);
	double cosine = cos(truth.yaw_rad);
	double sine = sin(truth.yaw_rad);
	double specific_force_nav[3] = {
		truth.accel_nav_mps2[0], truth.accel_nav_mps2[1],
		truth.accel_nav_mps2[2] + gravity
	};
	fusion_imu_measurement_t imu;
	memset(&imu, 0, sizeof(imu));
	imu.timestamp_ticks = ticks;
	imu.accel_mps2[0] = cosine * specific_force_nav[0] +
	                    sine * specific_force_nav[1];
	imu.accel_mps2[1] = -sine * specific_force_nav[0] +
	                     cosine * specific_force_nav[1];
	imu.accel_mps2[2] = specific_force_nav[2];
	imu.gyro_rps[2] = truth.yaw_rate_rps;
	/* Constant offsets are learned during stationary startup.  Small
	 * deterministic broadband components keep the test reproducible. */
	static const double accel_bias[3] = {0.030, -0.020, 0.050};
	static const double gyro_bias[3] = {0.0008, -0.0005, 0.0006};
	for(unsigned int axis = 0u; axis < 3u; axis++)
	{
		double phase = time_s * (11.0 + 2.0 * axis) + (double)axis;
		imu.accel_mps2[axis] += accel_bias[axis] + 0.008 * sin(phase);
		imu.gyro_rps[axis] += gyro_bias[axis] + 0.00025 * cos(1.7 * phase);
	}
	imu.temperature_c = 25.0 + 0.4 * sin(0.1 * time_s);
	imu.sequence = (uint32_t)((ticks * IMU_HZ + TIMEBASE_HZ / 2u) /
	                          TIMEBASE_HZ);
	imu.vehicle_at_rest =
		(hypot(truth.velocity_mps[0], truth.velocity_mps[1]) < 0.05) &&
		(hypot(truth.accel_nav_mps2[0], truth.accel_nav_mps2[1]) < 0.05) &&
		(fabs(truth.yaw_rate_rps) < 0.02);
	return imu;
}

static fusion_position_measurement_t make_gps(uint64_t ticks)
{
	double time_s = (double)ticks / (double)TIMEBASE_HZ;
	truth_t truth = truth_at(time_s);
	fusion_position_measurement_t gps;
	memset(&gps, 0, sizeof(gps));
	gps.timestamp_ticks = ticks;
	gps.position_enu_m[0] = truth.position_m[0] + 0.45 * sin(7.1 * time_s);
	gps.position_enu_m[1] = truth.position_m[1] + 0.45 * cos(5.3 * time_s);
	gps.position_enu_m[2] = truth.position_m[2] + 0.70 * sin(3.7 * time_s);
	gps.covariance_enu_m2[0] = square(1.2);
	gps.covariance_enu_m2[4] = square(1.2);
	gps.covariance_enu_m2[8] = square(2.0);
	gps.source_quality = 1.0;
	gps.source_valid = true;
	return gps;
}

static bool covariance_diagonal_valid(const eskf_t* filter)
{
	for(unsigned int index = 0u; index < ESKF_ERROR_STATE_DIM; index++)
	{
		if(!isfinite(filter->covariance[index][index]) ||
		   !(filter->covariance[index][index] > 0.0))
		{
			return false;
		}
	}
	return true;
}

static metrics_t run_scenario(unsigned int gps_period_ms)
{
	static fusion_engine_t engine;
	memset(&engine, 0, sizeof(engine));
	fusion_engine_config_t config;
	fusion_engine_default_config(&config, TIMEBASE_HZ);
	config.calibration_min_samples = 100u;
	metrics_t metrics;
	memset(&metrics, 0, sizeof(metrics));
	metrics.gps_period_ms = gps_period_ms;
	if(!fusion_engine_init(&engine, &config))
	{
		return metrics;
	}
	if(gps_period_ms != 0u)
	{
		fusion_position_measurement_t origin = make_gps(0u);
		origin.position_enu_m[0] = 0.0;
		origin.position_enu_m[1] = 0.0;
		origin.position_enu_m[2] = 0.0;
		(void)fusion_engine_push_position(&engine, &origin, NULL, NULL);
	}

	gps_event_t events[GPS_EVENT_CAPACITY];
	unsigned int event_head = 0u;
	unsigned int event_count = 0u;
	uint64_t next_gps_ticks = 3037000u;
	uint64_t gps_period_ticks = (uint64_t)gps_period_ms * 1000u;
	double error_square_sum = 0.0;
	uint64_t error_samples = 0u;
	double sigma_at_outage_start = NAN;
	double sigma_before_recovery = NAN;
	uint64_t accepted_before_recovery = 0u;
	double previous_position[3] = {0.0, 0.0, 0.0};
	bool have_previous_position = false;
	uint64_t sample_count = (uint64_t)(DURATION_S * (double)IMU_HZ);

	for(uint64_t sample = 0u; sample <= sample_count; sample++)
	{
		uint64_t ticks = (sample * TIMEBASE_HZ + IMU_HZ / 2u) / IMU_HZ;
		double time_s = (double)ticks / (double)TIMEBASE_HZ;
		fusion_imu_measurement_t imu = make_imu(ticks, config.eskf.gravity_mps2);
		if(gps_period_ms == 0u)
		{
			/* Preserve the pure unaided observability case: no GNSS and no
			 * external platform-rest constraint. */
			imu.vehicle_at_rest = false;
		}
		if(trace_enabled)
		{
			printf("IMU ts=%llu seq=%u accel=%.6f,%.6f,%.6f gyro=%.6f,%.6f,%.6f\n",
			       (unsigned long long)imu.timestamp_ticks, imu.sequence,
			       imu.accel_mps2[0], imu.accel_mps2[1], imu.accel_mps2[2],
			       imu.gyro_rps[0], imu.gyro_rps[1], imu.gyro_rps[2]);
		}
		if(!fusion_engine_push_imu(&engine, &imu))
		{
			break;
		}

		if((gps_period_ms != 0u) && (ticks >= next_gps_ticks))
		{
			double gps_time_s = (double)next_gps_ticks / (double)TIMEBASE_HZ;
			if(!((gps_time_s >= 10.0) && (gps_time_s < 18.0)) &&
			   (event_count < GPS_EVENT_CAPACITY))
			{
				events[event_count].measurement = make_gps(next_gps_ticks);
				events[event_count].arrival_ticks = next_gps_ticks +
					((gps_period_ms == 1000u) ? 310000u : 73000u);
				event_count++;
			}
			next_gps_ticks += gps_period_ticks;
		}

		while((event_head < event_count) &&
		      (events[event_head].arrival_ticks <= ticks))
		{
			eskf_update_result_t result = ESKF_UPDATE_INVALID;
			double nis = NAN;
			(void)fusion_engine_push_position(
				&engine, &events[event_head].measurement, &result, &nis);
			if(trace_enabled)
			{
				printf("GPS ts=%llu arrival=%llu result=%u nis=%.6f pos=%.3f,%.3f,%.3f\n",
				       (unsigned long long)events[event_head].measurement.timestamp_ticks,
				       (unsigned long long)ticks, (unsigned int)result, nis,
				       events[event_head].measurement.position_enu_m[0],
				       events[event_head].measurement.position_enu_m[1],
				       events[event_head].measurement.position_enu_m[2]);
			}
			event_head++;
		}

		if(engine.state == FUSION_ENGINE_RUNNING)
		{
			truth_t truth = truth_at(time_s);
			double east_error = engine.current.filter.position_m[0] - truth.position_m[0];
			double north_error = engine.current.filter.position_m[1] - truth.position_m[1];
			double horizontal_error = hypot(east_error, north_error);
			if(have_previous_position)
			{
				double position_step = hypot(
					engine.current.filter.position_m[0] - previous_position[0],
					engine.current.filter.position_m[1] - previous_position[1]);
				metrics.maximum_position_step_m =
					fmax(metrics.maximum_position_step_m, position_step);
			}
			memcpy(previous_position, engine.current.filter.position_m,
			       sizeof(previous_position));
			have_previous_position = true;
			if(time_s >= 2.5)
			{
				error_square_sum += square(horizontal_error);
				error_samples++;
			}
			if((time_s >= 10.0) && (time_s < 18.0))
			{
				metrics.peak_outage_error_m =
					fmax(metrics.peak_outage_error_m, horizontal_error);
			}
			if((hypot(truth.velocity_mps[0], truth.velocity_mps[1]) > 0.2) &&
			   (engine.current.motion.state == FUSION_MOTION_STATIONARY))
			{
				metrics.false_stationary_samples++;
			}
			if(engine.current.motion.state == FUSION_MOTION_MOVING)
			{
				metrics.moving_detected = true;
			}
			if(isnan(sigma_at_outage_start) && (time_s >= 10.0))
			{
				sigma_at_outage_start = sqrt(engine.current.filter.covariance[0][0]);
				accepted_before_recovery = engine.diagnostics.position_accepted;
			}
			if(isnan(sigma_before_recovery) && (time_s >= 18.0))
			{
				sigma_before_recovery = sqrt(engine.current.filter.covariance[0][0]);
			}
		}
	}

	truth_t final_truth = truth_at(DURATION_S);
	metrics.horizontal_rmse_m = (error_samples != 0u)
		? sqrt(error_square_sum / (double)error_samples) : INFINITY;
	metrics.final_horizontal_error_m = hypot(
		engine.current.filter.position_m[0] - final_truth.position_m[0],
		engine.current.filter.position_m[1] - final_truth.position_m[1]);
	metrics.final_vertical_error_m =
		fabs(engine.current.filter.position_m[2] - final_truth.position_m[2]);
	metrics.final_speed_mps = eskf_speed_mps(&engine.current.filter);
	metrics.outage_sigma_growth_m =
		(isfinite(sigma_at_outage_start) && isfinite(sigma_before_recovery))
		? sigma_before_recovery - sigma_at_outage_start : 0.0;
	metrics.final_horizontal_sigma_m = sqrt(
		engine.current.filter.covariance[0][0] +
		engine.current.filter.covariance[1][1]);
	metrics.gps_accepted = engine.diagnostics.position_accepted;
	metrics.gps_rejected = engine.diagnostics.position_nis_rejected;
	metrics.replay_count = engine.diagnostics.replay_count;
	metrics.imu_received = engine.diagnostics.imu_received;
	metrics.recovery_gps_accepted = (gps_period_ms == 0u) ||
		(engine.diagnostics.position_accepted > accepted_before_recovery);
	metrics.local_state_finite = (engine.state == FUSION_ENGINE_RUNNING) &&
		eskf_state_is_finite(&engine.current.filter);
	metrics.covariance_finite_positive =
		covariance_diagonal_valid(&engine.current.filter);
	metrics.global_origin_available = engine.have_initial_position;
	return metrics;
}

static void print_metrics(const metrics_t* metrics)
{
	printf("GPS=%4u ms IMU=%llu GPS-ok/rej=%llu/%llu replay=%llu "
	       "rmse-h=%.3f m outage-peak=%.3f m sigma-growth=%.3f m "
	       "final=%.3f/%.3f m sigma-h=%.3f m step-max=%.3f m "
	       "speed=%.4f m/s false-ZUPT=%u origin=%u\n",
	       metrics->gps_period_ms,
	       (unsigned long long)metrics->imu_received,
	       (unsigned long long)metrics->gps_accepted,
	       (unsigned long long)metrics->gps_rejected,
	       (unsigned long long)metrics->replay_count,
	       metrics->horizontal_rmse_m, metrics->peak_outage_error_m,
	       metrics->outage_sigma_growth_m,
	       metrics->final_horizontal_error_m, metrics->final_vertical_error_m,
	       metrics->final_horizontal_sigma_m, metrics->maximum_position_step_m,
	       metrics->final_speed_mps, metrics->false_stationary_samples,
	       metrics->global_origin_available ? 1u : 0u);
}

int main(int argc, char** argv)
{
	trace_enabled = (argc == 2) && (strcmp(argv[1], "--trace") == 0);
	const unsigned int periods_ms[] = {100u, 200u, 1000u, 0u};
	bool passed = true;
	printf("Timestamped 52 Hz IMU, asynchronous GPS, 8 s outage (10..18 s)\n");
	for(unsigned int index = 0u;
	    index < sizeof(periods_ms) / sizeof(periods_ms[0]); index++)
	{
		metrics_t metrics = run_scenario(periods_ms[index]);
		print_metrics(&metrics);
		bool common_ok = metrics.local_state_finite &&
		                 metrics.covariance_finite_positive &&
		                 metrics.moving_detected &&
		                 (metrics.false_stationary_samples == 0u) &&
		                 (metrics.imu_received == 1301u) &&
		                 (metrics.maximum_position_step_m < 1.0);
		bool gps_ok = (periods_ms[index] == 0u)
			? (!metrics.global_origin_available && (metrics.gps_accepted == 0u) &&
			   (metrics.peak_outage_error_m < 1.5) &&
			   (metrics.final_horizontal_error_m < 5.0) &&
			   (metrics.final_horizontal_error_m <
			    3.0 * metrics.final_horizontal_sigma_m) &&
			   (metrics.outage_sigma_growth_m > 0.0))
			: (metrics.global_origin_available && metrics.recovery_gps_accepted &&
			   (metrics.gps_rejected == 0u) && (metrics.replay_count > 0u) &&
			   (metrics.peak_outage_error_m < 2.0) &&
			   (metrics.outage_sigma_growth_m > 0.0) &&
			   (metrics.final_speed_mps < 0.08) &&
			   (metrics.final_horizontal_error_m < 1.5) &&
			   (metrics.final_vertical_error_m < 1.5));
		passed = passed && common_ok && gps_ok;
	}
	if(!passed)
	{
		fprintf(stderr, "FAIL: asynchronous rate/dead-reckoning matrix.\n");
		return 1;
	}
	printf("PASS: all rates use the same IMU prediction path; GPS loss remains "
	       "local dead reckoning and recovery is an external update.\n");
	return 0;
}