/* wmp.c */

#include "wmp.h"
#include "i2c.h"
#include "uart.h"
#include "dcm.h"
#include "fisqrt.h"
#include "stick.h"
#include "watchdog.h"
#include "status.h"
#include "abs.h"

#define WMP_ZERO_COUNT 100


#define ACCEL_ZERO_X 520
#define ACCEL_ZERO_Y 516
#define ACCEL_ZERO_Z 514

/*
516, 510, 710
-4, -6, 196
16, 36, 38416   = 38468
sqrt(38468) = 196.1326...
... somehow once we scale by gravity we get almost exactly 0.05.
*/

#define ACCEL_SCALE 0.05

/* Nunchuck pass-through mode */
unsigned char wmp_init_command[2] = {0xfe, 0x05};

i2c_result wmp_result;
unsigned int wmp_generation;

struct i2c_transaction wmp_init_transaction = {
	(0x53 << 1) + 0, /* write */
	2,
	wmp_init_command,
	&wmp_result,
	NULL
};

unsigned char wmp_read_cal_command[1] = {0x20};

struct i2c_transaction wmp_read_cal_transaction2;

struct i2c_transaction wmp_read_cal_transaction = {
	(0x53 << 1) + 0, /* write */
	1,
	wmp_read_cal_command,
	&wmp_result,
	&wmp_read_cal_transaction2
};

struct i2c_transaction wmp_read_cal_transaction2 = {
	(0x53 << 1) + 1, /* read */
	0x20,
	wmp_calibration_data,
	&wmp_result,
	NULL
};

unsigned char wmp_sample_command[1] = {0x00};

unsigned char wmp_sample_data[6];

struct i2c_transaction wmp_sample_transaction2;

struct i2c_transaction wmp_sample_transaction = {
	(0x52 << 1) + 0, /* write */
	1,
	wmp_sample_command,
	&wmp_result,
	&wmp_sample_transaction2
};

struct i2c_transaction wmp_sample_transaction2 = {
	(0x52 << 1) + 1, /* read */
	6,
	wmp_sample_data,
	&wmp_result,
	NULL
};


bool wmp_init(void)
{
	if (!i2c_start_transaction(&wmp_init_transaction))
		return FALSE;
	while (i2c_busy()) ;
	return (wmp_result == I2C_SUCCESS);
}

unsigned char wmp_calibration_data[0x20];

bool wmp_read_calibration_data(void)
{
	if (!i2c_start_transaction(&wmp_read_cal_transaction))
		return FALSE;
	while (i2c_busy());
	return (wmp_result == I2C_SUCCESS);
}

unsigned int wmp_yaw;
unsigned int wmp_pitch;
unsigned int wmp_roll;

unsigned int wmp_yaw_zero;
unsigned int wmp_pitch_zero;
unsigned int wmp_roll_zero;

bool wmp_yaw_fast;
bool wmp_pitch_fast;
bool wmp_roll_fast;

int accel_x;
int accel_y;
int accel_z;

bool wmp_update;
bool wmp_zero;
unsigned int wmp_discard;

#define TWO_PI 6.28318531f
#define DEG_TO_RAD (TWO_PI/360.0f)

/* There's considerable debate about these values, and they may vary
 * between different models of the Wii Motion Plus. It would be nice
 * to be able to use the calibration data stored on the device itself
 * but we don't know the format yet.
 */
#define SLOW_YAW_STEP	(20 / DEG_TO_RAD)
#define SLOW_PITCH_STEP (20 / DEG_TO_RAD)
#define SLOW_ROLL_STEP  (20 / DEG_TO_RAD)
#define FAST_YAW_STEP	(4 / DEG_TO_RAD)
#define FAST_PITCH_STEP (4 / DEG_TO_RAD)
#define FAST_ROLL_STEP  (4 / DEG_TO_RAD)

/* The gyro has to stay within this limit in each axis in order to arm */
#define GYRO_RATE_THRESHOLD (0.01 / DEG_TO_RAD)

bool wmp_sample(void)
{
	if (!i2c_start_transaction(&wmp_sample_transaction))
		return FALSE;

	while (i2c_busy());

	if (wmp_result != I2C_SUCCESS)
		return FALSE;

	wmp_result = I2C_IN_PROGRESS;

	wmp_yaw   = ((wmp_sample_data[3]>>2)<<8) + wmp_sample_data[0];
	wmp_pitch = ((wmp_sample_data[4]>>2)<<8) + wmp_sample_data[1];
	wmp_roll  = ((wmp_sample_data[5]>>2)<<8) + wmp_sample_data[2];

	/* XXX We don't take into account the fast/slow mode flag here */
	wmp_yaw_fast = !(wmp_sample_data[3] & 0x2);
	wmp_pitch_fast = !(wmp_sample_data[3] & 0x1);
	wmp_roll_fast = !(wmp_sample_data[4] & 0x2);

	return TRUE;
}

bool wmp_start_sample(void)
{
	return i2c_start_transaction(&wmp_sample_transaction);
}

void wmp_process_gyro_sample(void)
{
	float yaw, pitch, roll;

	wmp_yaw   = ((wmp_sample_data[3]>>2)<<8) + wmp_sample_data[0];
	wmp_pitch = ((wmp_sample_data[4]>>2)<<8) + wmp_sample_data[1];
	wmp_roll  = ((wmp_sample_data[5]>>2)<<8) + wmp_sample_data[2];

	/* XXX We don't take into account the fast/slow mode flag here */
	wmp_yaw_fast = !(wmp_sample_data[3] & 0x2);
	wmp_pitch_fast = !(wmp_sample_data[3] & 0x1);
	wmp_roll_fast = !(wmp_sample_data[4] & 0x2);

	if (wmp_update) {
		int tmp_yaw = wmp_yaw;
		int tmp_pitch = wmp_pitch;
		int tmp_roll = wmp_roll;

		tmp_yaw -= wmp_yaw_zero;
		tmp_pitch -= wmp_pitch_zero;
		tmp_roll -= wmp_roll_zero;

		if (wmp_yaw_fast)
			yaw = ((float)tmp_yaw) / FAST_YAW_STEP;
		else
			yaw = ((float)tmp_yaw) / SLOW_YAW_STEP;

		if (wmp_pitch_fast)
			pitch = ((float)tmp_pitch) / FAST_PITCH_STEP;
		else
			pitch = ((float)tmp_pitch) / SLOW_PITCH_STEP;

		if (wmp_roll_fast)
			roll = ((float)tmp_roll) / FAST_ROLL_STEP;
		else
			roll = ((float)tmp_roll) / SLOW_ROLL_STEP;

		dcm_update(roll, pitch, yaw);

		if (!status_armed()) {
			if (  (abs(roll)  < GYRO_RATE_THRESHOLD) &&
			      (abs(pitch) < GYRO_RATE_THRESHOLD) &&
			      (abs(yaw)   < GYRO_RATE_THRESHOLD)  ) {
			    status_set_ready(STATUS_MODULE_GYRO_RATE, TRUE);
			} else {
			    status_set_ready(STATUS_MODULE_GYRO_RATE, FALSE);
			}
		}

		wmp_generation++;

#if SEND_DCM
		if ((wmp_generation % 20) == 0)
			dcm_send_packet();
#endif

	} else if (wmp_zero) {
		if (wmp_discard) {
			wmp_discard--;
		} else {
			wmp_yaw_zero += wmp_yaw;
			wmp_pitch_zero += wmp_pitch;
			wmp_roll_zero += wmp_roll;
			wmp_generation++;
			if (wmp_generation >= WMP_ZERO_COUNT) {
				wmp_zero = FALSE;
				wmp_update = TRUE;
				wmp_generation = 0;
				wmp_yaw_zero /= WMP_ZERO_COUNT;
				wmp_pitch_zero /= WMP_ZERO_COUNT;
				wmp_roll_zero /= WMP_ZERO_COUNT;
				putstr("Zero finished\r\n");
				status_set_ready(STATUS_MODULE_GYRO_ZERO, TRUE);
			}
		}
	}
	watchdog_kick(WATCHDOG_GYRO);
}

void wmp_process_accel_sample(void)
{
	float x, y, z;
#if 0
	float invmag;
#endif

	accel_x   = (wmp_sample_data[2]<<2) + ((wmp_sample_data[5]>>3) & 0x02);
	accel_y   = (wmp_sample_data[3]<<2) + ((wmp_sample_data[5]>>4) & 0x02);
	accel_z   = ((wmp_sample_data[4]<<2) & 0x3f8) +
					      ((wmp_sample_data[5]>>5) & 0x06);

	x = (accel_x - ACCEL_ZERO_X) * ACCEL_SCALE;
	y = (accel_y - ACCEL_ZERO_Y) * ACCEL_SCALE;
	z = (accel_z - ACCEL_ZERO_Z) * ACCEL_SCALE;

#if 0
	invmag = fisqrt(x*x + y*y + z*z);

	x = x * invmag;
	y = y * invmag;
	z = z * invmag;
#endif

#if 0
	accel_x = (x * 512.0 + 1000.0);
	accel_y = (y * 512.0 + 1000.0);
	accel_z = (z * 512.0 + 1000.0);
#endif

#if 0
	putstr("(");
	putint(accel_x);
	putstr(", ");
	putint(accel_y);
	putstr(", ");
	putint(accel_z);
	putstr(")\r\n");
#endif

	/* The minus signs are needed because something is upside down.
	 * It might actually be the WMP, but we're defining coordinates based
	 * on that so we'll just fudge it here.
	 */
	dcm_drift_correction(x, -y, -z);
	watchdog_kick(WATCHDOG_ACCEL);
	stick_input();
}

void wmp_event_handler(void)
{
	if (wmp_result != I2C_SUCCESS)
		return;

	wmp_result = I2C_IN_PROGRESS;

	if (wmp_sample_data[5] & 0x02)
		wmp_process_gyro_sample();
	else
		wmp_process_accel_sample();
}

void wmp_start_zero(void)
{
	wmp_zero = TRUE;
	wmp_update = FALSE;
	wmp_discard = 100;
	wmp_generation = 0;
	putstr("Starting zero\r\n");
}