#include "sensor.hpp"
#include "math.h"
#include "sensor_registers.hpp"

Sensor::Sensor(uint8_t address, TwoWire *wire)
{
	_address = address;
	_wire = wire;

	// Set default values
	accel_cal_x = 0;
	accel_cal_y = 0;
	accel_cal_z = 0;

	gyro_cal_x = 0;
	gyro_cal_y = 0;
	gyro_cal_z = 0;

	is_calibrated = false;

	_throttle_enabled = true;
	_throttleTime = DEFAULT_THROTTLE_TIME;

	_lastTime = 0;
	_lastMicros = 0;

	_status = SENSOR_OK;

	_accel_sensitivity = 0;
	_accel_to_g_force = 1.0 / 16384.0;

	_gyro_sensitivity = 0;
	_ang_rate_to_dps = 1.0 / 131.0;

	_temperature = 0;

	// Reset the sensor
	reset();
}

bool Sensor::begin()
{
	_wire->begin();

	if (isConnected())
	{
		_wire->beginTransmission(_address);
		_wire->write(SENSOR_PWR_MGMT_1);
		_wire->write(SENSOR_WAKEUP);

		return (_wire->endTransmission() == 0);
	}

	return false;
}

bool Sensor::isConnected()
{
	_wire->beginTransmission(_address);
	return (_wire->endTransmission() == 0);
}

void Sensor::reset()
{
	setThrottleTime(DEFAULT_THROTTLE_TIME);

	_accel_raw_x = 0;
	_accel_raw_y = 0;
	_accel_raw_z = 0;

	_gyro_raw_x = 0;
	_gyro_raw_y = 0;
	_gyro_raw_z = 0;

	_accel_x = 0;
	_accel_y = 0;
	_accel_z = 0;

	_pitch = 0;
	_roll = 0;
	_yaw = 0;
}

void Sensor::setThrottleEnabled(bool throttle = true)
{
	_throttle_enabled = throttle;
};

bool Sensor::getThrottleEnabled()
{
	return _throttle_enabled;
};

void Sensor::setThrottleTime(uint16_t time_ms)
{
	_throttleTime = time_ms;
};

uint16_t Sensor::getThrottleTime()
{
	return _throttleTime;
};

int16_t Sensor::read()
{
	auto time_now = millis();

	if (_throttle_enabled)
	{
		if ((time_now - _lastTime) < _throttleTime)
		{
			return SENSOR_THROTTLED;
		}
	}

	_lastTime = time_now;

	_wire->beginTransmission(_address);
	_wire->write(SENSOR_ACCEL_XOUT_H);

	if (_wire->endTransmission() != 0)
	{
		_status = SENSOR_ERR_WRITE;
		return _status;
	}

	int8_t response_length = _wire->requestFrom(_address, (uint8_t)14);

	if (response_length != 14)
	{
		_status = SENSOR_ERR_READ;
		return _status;
	}

	// ACCELEROMETER
	_accel_raw_x = _readTwoBytes(); // ACCEL_XOUT_H  ACCEL_XOUT_L
	_accel_raw_y = _readTwoBytes(); // ACCEL_YOUT_H  ACCEL_YOUT_L
	_accel_raw_z = _readTwoBytes(); // ACCEL_ZOUT_H  ACCEL_ZOUT_L

	// TEMPERATURE
	_temperature = _readTwoBytes(); // TEMP_OUT_H    TEMP_OUT_L

	// GYROSCOPE
	_gyro_raw_x = _readTwoBytes(); // GYRO_XOUT_H   GYRO_XOUT_L
	_gyro_raw_y = _readTwoBytes(); // GYRO_YOUT_H   GYRO_YOUT_L
	_gyro_raw_z = _readTwoBytes(); // GYRO_ZOUT_H   GYRO_ZOUT_L

	// Duration interval
	time_now = micros();
	float duration = (time_now - _lastMicros) * 1e-6; // Duration in seconds.
	_lastMicros = time_now;

	// Convert raw acceleration to g:s (g-force)
	_accel_raw_x *= _accel_to_g_force;
	_accel_raw_y *= _accel_to_g_force;
	_accel_raw_z *= _accel_to_g_force;

	// Error correct raw acceleration
	_accel_raw_x += accel_cal_x;
	_accel_raw_y += accel_cal_y;
	_accel_raw_z += accel_cal_z;

	// Prepare for Pitch Roll Yaw
	float accel_x_pow_two = pow(_accel_raw_x, 2);
	float accel_y_pow_two = pow(_accel_raw_y, 2);
	float accel_z_pow_two = pow(_accel_raw_z, 2);

	_accel_x = atan(_accel_raw_y / sqrt(accel_x_pow_two + accel_z_pow_two)) * RAD_TO_DEGREES;
	_accel_y = atan(-1.0 * _accel_raw_x / sqrt(accel_y_pow_two + accel_z_pow_two)) * RAD_TO_DEGREES;
	_accel_z = atan(_accel_raw_z / sqrt(accel_x_pow_two + accel_y_pow_two)) * RAD_TO_DEGREES;

	// Convert temperature to celsius
	_temperature = _temperature * 0.00294117647 + 36.53;

	// Convert raw Gyro to degrees/s
	_gyro_raw_x *= _ang_rate_to_dps;
	_gyro_raw_y *= _ang_rate_to_dps;
	_gyro_raw_z *= _ang_rate_to_dps;

	// Error correct raw gyro measurements.
	_gyro_raw_x += gyro_cal_x;
	_gyro_raw_y += gyro_cal_y;
	_gyro_raw_z += gyro_cal_z;

	_gyro_x += _gyro_raw_x * duration;
	_gyro_y += _gyro_raw_y * duration;
	_gyro_z += _gyro_raw_z * duration;

	_yaw = _gyro_z;
	_pitch = 0.96 * _gyro_y + 0.04 * _accel_y;
	_roll = 0.96 * _gyro_x + 0.04 * _accel_x;

	return SENSOR_OK;
}

bool Sensor::setAccelSensitivity(uint8_t sensitivity)
{
	_accel_sensitivity = sensitivity;

	if (_accel_sensitivity > 3)
		_accel_sensitivity = 3;

	uint8_t val = getRegister(SENSOR_ACCEL_CONFIG);

	if (_status != SENSOR_OK)
	{
		return false;
	}

	// No need to write same value
	if (((val >> 3) & 3) != _accel_sensitivity)
	{
		val &= 0xE7;
		val |= (_accel_sensitivity << 3);

		if (setRegister(SENSOR_ACCEL_CONFIG, val) != SENSOR_OK)
		{
			return false;
		}
	}

	// Calculate conversion factor.
	_accel_to_g_force = (1 << _accel_sensitivity) * RAW_TO_G_FACTOR;

	return true;
}

uint8_t Sensor::getAccelSensitivity()
{
	uint8_t val = getRegister(SENSOR_ACCEL_CONFIG);

	if (_status != SENSOR_OK)
	{
		return _status; // return and propagate error (best thing to do)
	}

	_accel_sensitivity = (val >> 3) & 3;

	return _accel_sensitivity;
}

bool Sensor::setGyroSensitivity(uint8_t sensitivity)
{
	_gyro_sensitivity = sensitivity;

	if (_gyro_sensitivity > 3)
		_gyro_sensitivity = 3;

	uint8_t val = getRegister(SENSOR_GYRO_CONFIG);

	if (_status != 0)
	{
		return false;
	}

	// No need to write same value
	if (((val >> 3) & 3) != _gyro_sensitivity)
	{
		val &= 0xE7;
		val |= (_gyro_sensitivity << 3);

		if (setRegister(SENSOR_GYRO_CONFIG, val) != SENSOR_OK)
		{
			return false;
		}
	}

	_ang_rate_to_dps = (1 << _gyro_sensitivity) * RAW_TO_DPS_FACTOR;

	return true;
}

uint8_t Sensor::getGyroSensitivity()
{
	uint8_t val = getRegister(SENSOR_GYRO_CONFIG);

	if (_status != SENSOR_OK)
	{
		return _status;
	}

	_gyro_sensitivity = (val >> 3) & 3;

	return _gyro_sensitivity;
}

float Sensor::getAccelX() { return _accel_raw_x; };
float Sensor::getAccelY() { return _accel_raw_y; };
float Sensor::getAccelZ() { return _accel_raw_z; };

float Sensor::getAngleX() { return _accel_x; };
float Sensor::getAngleY() { return _accel_y; };

float Sensor::getTemperature() { return _temperature; };

float Sensor::getGyroX() { return _gyro_raw_x; };
float Sensor::getGyroY() { return _gyro_raw_y; };
float Sensor::getGyroZ() { return _gyro_raw_z; };

float Sensor::getPitch() { return _pitch; };
float Sensor::getRoll() { return _roll; };

uint32_t Sensor::lastTime() { return _lastTime; };

uint8_t Sensor::setRegister(uint8_t reg, uint8_t value)
{
	_wire->beginTransmission(_address);
	_wire->write(reg);
	_wire->write(value);

	if (_wire->endTransmission() != 0)
	{
		_status = SENSOR_ERR_WRITE;
		return _status;
	}

	return SENSOR_OK;
}

uint8_t Sensor::getRegister(uint8_t reg)
{
	_wire->beginTransmission(_address);
	_wire->write(reg);

	if (_wire->endTransmission() != 0)
	{
		_status = SENSOR_ERR_WRITE;
		return _status;
	}

	uint8_t response_length = _wire->requestFrom(_address, (uint8_t)1);
	if (response_length != 1)
	{
		_status = SENSOR_ERR_READ;
		return _status;
	}

	uint8_t val = _wire->read();

	return val;
}

int16_t Sensor::getStatus()
{
	auto error = _status;
	_status = SENSOR_OK;

	return error;
};

int16_t Sensor::_readTwoBytes()
{
	int16_t response = _wire->read();
	response <<= 8;
	response |= _wire->read();

	return response;
}

bool Sensor::autoCalibrate()
{
	bool done = false;

	float accel_x, accel_y, accel_z;
	float gyro_x, gyro_y, gyro_z;

	auto cal_start_time = millis();

	while (!done)
	{
		if (millis() - cal_start_time >= CALIBRATION_TIMEOUT)
		{
			return false;
		}

		delay(100);

		accel_x = 0;
		accel_y = 0;
		accel_z = 0;

		gyro_x = 0;
		gyro_y = 0;
		gyro_z = 0;

		for (int i = 0; i < 20; i++)
		{
			read();

			accel_x -= getAccelX();
			accel_y -= getAccelY();
			accel_z -= getAccelZ();

			gyro_x -= getGyroX();
			gyro_y -= getGyroY();
			gyro_z -= getGyroZ();
		}

		const float adjust_val = 0.05;

		accel_x *= adjust_val;
		accel_y *= adjust_val;
		accel_z *= adjust_val;

		gyro_x *= adjust_val;
		gyro_y *= adjust_val;
		gyro_z *= adjust_val;

		if (
			accel_x < ACCEL_CAL_X_MAX && accel_x > ACCEL_CAL_X_MIN &&
			accel_y < ACCEL_CAL_Y_MAX && accel_y > ACCEL_CAL_Y_MIN &&
			accel_z < ACCEL_CAL_Z_MAX && accel_z > ACCEL_CAL_Z_MIN &&
			gyro_x < GYRO_CAL_X_MAX && gyro_x > GYRO_CAL_X_MIN &&
			gyro_y < GYRO_CAL_Y_MAX && gyro_y > GYRO_CAL_Y_MIN &&
			gyro_z < GYRO_CAL_Z_MAX && gyro_z > GYRO_CAL_Z_MIN)
		{
			done = true;
		}

		// Adjust calibration values
		accel_cal_x += accel_x;
		accel_cal_y += accel_y;
		accel_cal_z += accel_z;

		gyro_cal_x += gyro_x;
		gyro_cal_y += gyro_y;
		gyro_cal_z += gyro_z;
	}

	is_calibrated = true;

	return true;
}