#include "common.h"
#include "bus/i2c.h"

#include <avr/io.h>
#include <math.h>

// I2C status flags
#define TW_START       0x08  // Start condition transmitted
#define TW_REP_START   0x10  // Repeated start condition transmitted
#define TW_MT_SLA_ACK  0x18  // SLA+W transmitted, ACK received
#define TW_MT_SLA_NACK 0x20  // SLA+W transmitted, NACK received
#define TW_MT_DATA_ACK 0x28  // Data transmitted, ACK received
#define TW_MR_SLA_ACK  0x40  // SLA+R transmitted, ACK received
#define TW_MR_SLA_NACK 0x48  // SLA+R transmitted, NACK received
#define TW_MR_DATA_ACK 0x50  // Data received, ACK returned

// NTC conversion constants
#define NTC01 1.009249522e-3
#define NTC02 2.378405444e-4
#define NTC03 2.019202697e-7

// TODO: Error handling and recovery besides watchdog timer.
// TODO: Add more documentation from the atmel data sheet.
// TODO: ADS1115 continuous mode instead of single-shot?
// TODO: Implement TWI_vect ISR? This may not actually be
// much better than the blocking approach.

static void  I2C_AHT20_Reset(void);
static bool  I2C_AHT20_IsCalibrated(void);
static void  I2C_AHT20_Calibrate(void);
static void  I2C_AHT20_Trigger(void);

static int   I2C_ADS1115_WriteRegister(byte reg, word data);
static word  I2C_ADS1115_ReadRegister(byte reg);
static bool  I2C_ADS1115_IsReady(void);

int I2C_Init(void)
{
	// Set SCL bit rate to 400kHz

	TWSR = 0x00;  // TWI status register
	TWBR = 0x0C;  // TWI bit rate register

	return 0;
}

int I2C_Start(byte addr, byte mode)
{
	unsigned int status;

	// Send start condition

	TWCR = (1 << TWEN)    // Enable TWI
	     | (1 << TWINT)   // Clear interrupt flag
	     | (1 << TWSTA);  // Send start condition

	// Wait until start condition sent

	while ((TWCR & (1 << TWINT)) == 0);
	if ((TWSR & 0xf8) != TW_START)
		return -1;

	// Send slave address

	addr = addr << 1;     // Lowest bit for mode
	TWDR = addr + mode;   // Mode 0=R, 1=W
	TWCR = (1 << TWEN)    // Enable TWI
	     | (1 << TWINT);  // Clear interrupt flag

	// Wait until slave address sent

	while ((TWCR & (1 << TWINT)) == 0);
	status = TWSR & 0xF8;

	if ((mode == 0 && status != TW_MT_SLA_ACK) ||
	    (mode == 1 && status != TW_MR_SLA_ACK))
		return -2;

	return 0;
}

int I2C_Write(byte data)
{
	TWDR = data;
	TWCR = (1 << TWEN)    // Enable TWI
	     | (1 << TWINT);  // Clear interrupt flag

	while ((TWCR & (1 << TWINT)) == 0);

	if ((TWSR & 0xF8) != TW_MT_DATA_ACK)
		return -1;

	return 0;
}

byte I2C_Read_ACK(void)
{
	// Read data and acknowledge

	TWCR = (1 << TWEN)   // Enable TWI
	     | (1 << TWINT)  // Clear interrupt flag
	     | (1 << TWEA);  // Send acknowledgment

	// Wait until data read

	while ((TWCR & (1 << TWINT)) == 0);

	return TWDR;
}

byte I2C_Read_NACK(void)
{
	// Read data, expect last byte

	TWCR = (1 << TWEN)    // Enable TWI
	     | (1 << TWINT);  // Clear interrupt flag

	// Wait until data read

	while ((TWCR & (1 << TWINT)) == 0);

	return TWDR;
}

int I2C_Wait_ACK(void)
{
	unsigned int status;

	while ((TWCR & (1 << TWINT)) == 0);
	status = TWSR & 0xF8;

	if (status != TW_MT_DATA_ACK &&
	    status != TW_MR_DATA_ACK)
		return -1;

	return 0;
}

int I2C_Stop(void)
{
	// Send stop condition

	TWCR = (1 << TWEN)    // Enable TWI
	     | (1 << TWINT)   // Clear interrupt flag
	     | (1 << TWSTO);  // Send stop condition

	// Wait until stop condition sent

	while (TWCR & (1 << TWSTO));

	return 0;
}

void I2C_Reset(void)
{
	byte TWBRold = TWBR;
	byte TWARold = TWAR;

	// TODO: Handle timeouts and reset TWI if necessary,
	// fatal error after multiple attempts and then bring
	// system to a safe stop.

	// TWI_Disable();
	// TWI_Init();

	TWAR = TWARold;
	TWBR = TWBRold;;
}

int I2C_SetChannel(int channel)
{
	byte crb;

	// Excerpts taken from the PCA9546A datasheet:
	// https://www.ti.com/lit/ds/symlink/pca9546a.pdf

	// Following a start condition, the bus master
	// must output the address of the slave it is
	// accessing.

	//  7  6  5  4   |   3   2   1   0
	//  1  1  1  0   |  A2  A1  A0  RW
	//  FIXED_____   |  SELECTABLE____

	I2C_Start(0x70, 0);
	I2C_Wait_ACK();

	// Following the successful acknowledgment of
	// the slave address, the bus master sends a byte
	// which is stored in the control register.

	//  7  6  5  4   |   3   2   1   0
	//  X  X  X  X   |  B3  B2  B1  B0
	//  UNUSED____   |  CHANNEL_______

	crb = (1 << channel);

	// crb = 0x01; //  Channel 0  |  00000001
	//       0x02  //  Channel 1  |  00000010
	//       0x04  //  Channel 2  |  00000100
	//       0x08  //  Channel 3  |  00001000

	I2C_Write(crb);
	I2C_Wait_ACK();

	// When a channel is selected, the channel
	// becomes active after a stop condition has been
	// placed on the I2C bus. A stop condition always
	// must occur right after the acknowledge cycle.

	I2C_Stop();

	return 0;
}

int I2C_AHT20_Init(void)
{
	I2C_AHT20_Reset();

	// Wait 40ms after power-on

	Sleep(40);

	Info("Initializing AHT20 sensor...");

	if (!I2C_AHT20_IsCalibrated()) {
		I2C_AHT20_Calibrate();
		if (!I2C_AHT20_IsCalibrated()) {
			Info("Error: Calibration failed.");
			return -1;
		}
	}

	return 0;
}

int I2C_AHT20_Read(float *temp, float *rhum)
{
	byte data[6], crc;
	unsigned long hraw, traw;

	I2C_AHT20_Trigger();

	I2C_Start(0x38, 1);
	I2C_Wait_ACK();

	data[0] = I2C_Read_ACK();
	data[1] = I2C_Read_ACK();
	data[2] = I2C_Read_ACK();
	data[3] = I2C_Read_ACK();
	data[4] = I2C_Read_ACK();
	data[5] = I2C_Read_ACK();
	crc = I2C_Read_NACK();
	I2C_Stop();

	// After receiving six bytes, the next byte is the
	// CRC check data, the user can read it as needed,
	// if the receiving end needs CRC check, then send
	// it after receiving the sixth byte ACK response,
	// otherwise NACK is sent out, CRC initial value is
	// 0xFF. The CRC8 check polynomial is:
	// CRC[7:0]= 1 + (x^4) + (x^5) + (x^8)

	// TODO: Implement CRC8 check
	UNUSED(crc);

	hraw = (unsigned long) data[1] << 12;
	hraw |= (unsigned short) data[2] << 4;
	hraw |= data[3] >> 4;

	traw = (unsigned long) (data[3] & 0x0F) << 16;
	traw |= (unsigned short) data[4] << 8;
	traw |= data[5];

	*rhum = (float) hraw / 1048576.0f * 100.0f;
	*temp = (float) traw / 1048576.0f * 200.0f - 50.0f;

	return 0;
}

static void I2C_AHT20_Reset(void)
{
	// The command 0xBA is used to restart the sensor
	// system without turning the power off and on
	// again. After receiving this command, the sensor
	// system begins to re-initialize and restore the
	// default setting state. The time required for
	// soft reset does not exceed 20 ms.

	I2C_Start(0x38, 0);
	I2C_Wait_ACK();

	I2C_Write(0xBA);
	I2C_Wait_ACK();
	I2C_Stop();

	Sleep(20);
}

static bool I2C_AHT20_IsCalibrated(void)
{
	byte status;

	// Before reading the temperature and humidity
	// values, first check whether the calibration
	// enable bit Bit [3] of the status word is 1 (you
	// can get a byte of status word by sending 0x71).

	I2C_Start(0x38, 0);
	I2C_Wait_ACK();

	I2C_Write(0x71);
	I2C_Wait_ACK();
	I2C_Stop();

	I2C_Start(0x38, 1);
	I2C_Wait_ACK();

	status = I2C_Read_NACK();
	I2C_Stop();

	// Note: The calibration status check only needs to
	// be checked at power-on. No operation is required
	// during the normal acquisition process.

	return status & BIT(3);
}

static void I2C_AHT20_Calibrate(void)
{
	// To calibrate you need to send 0xBE command (for
	// initialization), this command parameter has two
	// bytes, the first byte is 0x08, the second byte
	// is 0x00, and then wait for 10ms.

	I2C_Start(0x38, 0);
	I2C_Wait_ACK();

	I2C_Write(0xBE);
	I2C_Wait_ACK();
	I2C_Write(0x08);
	I2C_Wait_ACK();
	I2C_Write(0x00);
	I2C_Wait_ACK();
	I2C_Stop();
}

static void I2C_AHT20_Trigger(void)
{
	byte status;

	// Send the 0xAC command directly (trigger
	// measurement). The parameter of this command has
	// two bytes, the first byte is 0x33 and the second
	// byte is 0x00.

	I2C_Start(0x38, 0);
	I2C_Wait_ACK();

	I2C_Write(0xAC);
	I2C_Wait_ACK();
	I2C_Write(0x33);
	I2C_Wait_ACK();
	I2C_Write(0x00);
	I2C_Wait_ACK();
	I2C_Stop();

	// Wait for 80ms to wait for the measurement to be
	// completed. If the read status word Bit [7] is 0,
	// it indicates that the measurement is completed,
	// and six bytes can be read in a row. Otherwise
	// continue to wait.

	I2C_Start(0x38, 1);
	I2C_Wait_ACK();

	do {
		Sleep(80);
		status = I2C_Read_ACK();
	} while ((status & BIT(7)) == 1);
	I2C_Stop();
}

static int I2C_ADS1115_WriteRegister(byte reg, word data)
{
	// The ADS111x have one address pin, ADDR, that
	// configures the I2C address of the device. This
	// pin can be connected to GND, VDD, SDA, or SCL,
	// allowing for four different addresses to be
	// selected with one pin:

	// GND 1001000
	// VDD 1001001
	// SDA 1001010
	// SCL 1001011

	// In slave receive mode, the first byte transmitted
	// from the master to the slave consists of the 7-bit
	// device address followed by a low R/W bit.

	I2C_Start(0x48, 0); // ADDR pin to GND
	I2C_Wait_ACK();

	// The next byte transmitted by the master is the
	// Address Pointer register. The ADS111x then
	// acknowledge receipt of the Address Pointer
	// register byte.

	I2C_Write(reg);
        I2C_Wait_ACK();

	// The next two bytes are written to the address
	// given by the register address pointer bits P[1:0].
	// The ADS111x acknowledge each byte sent. Register
	// bytes are sent with the most significant byte
	// first, followed by the least significant byte.

	I2C_Write(data >> 8);
        I2C_Wait_ACK();
	I2C_Write(data & 0xFF);
        I2C_Wait_ACK();

	I2C_Stop();

	return 0;
}

static word I2C_ADS1115_ReadRegister(byte reg)
{
	byte data[2];

	// To access a specific register from the ADS1115,
	// the master must first write an appropriate value
	// to register address pointer bits P[1:0] in the
	// Address Pointer register.

	// The Address Pointer register is written to
	// directly after the slave address byte, low R/W bit,
	// and a successful slave acknowledgment.

	I2C_Start(0x48, 0);
	I2C_Wait_ACK();
	I2C_Write(reg);

	// After the Address Pointer register
	// is written, the slave acknowledges, and the master
	// issues a STOP or a repeated START condition.

        I2C_Wait_ACK();
	I2C_Stop();
	I2C_Start(0x48, 1);
	I2C_Wait_ACK();

	// The ADS111x provide 16 bits of data in binary
	// two's complement format. A positive full-scale (+FS)
	// input produces an output code of 7FFFh and a negative
	// full-scale (–FS) input produces an output code of
	// 8000h. The output clips at these codes for signals
	// that exceed full-scale. XXX: ???

	data[0] = I2C_Read_ACK();
	data[1] = I2C_Read_NACK();

	I2C_Stop();

	return data[0] << 8 | data[1];
}

static bool I2C_ADS1115_IsReady(void)
{
	word reg;

	Sleep(20);

	reg = I2C_ADS1115_ReadRegister(0x01);

	return reg & BIT(15);
}

word I2C_ADS1115_Read(int channel)
{
	word config;
	word os, mode, rate, gain, mux;

	// Operational status or single-shot conversion
	// start: This bit determines the operational
	// status of the device. OS can only be written
	// when in power-down state and has no effect
	// when a conversion is ongoing.

	// When writing:
	// 0 : No effect
	// 1 : Start single conversion (power-down state)
	// When reading:
	// 0 : Device is performing a conversion
	// 1 : Device is not performing a conversion

	os = 0x8000; // Start single   |  10000000 00000000

	// Device operating mode: This bit controls the
	// operating mode.

	// 0 : Continuous-conversion mode
	// 1 : Single-shot mode or power-down state

	mode = 0x0100; // Single-shot  |  00000001 00000000

	// Data rate: These bits control the data rate setting.

	// 000  |    8 SPS
	// 001  |   16 SPS
	// 010  |   32 SPS
	// 011  |   64 SPS
	// 100  |  128 SPS
	// 101  |  250 SPS
	// 110  |  475 SPS
	// 111  |  860 SPS

	rate = 0x0080; // 128 SPS      |  00000000 10000000

	// Programmable gain amplifier configuration: These
	// bits set the FSR of the programmable gain amplifier.

	// 000  |  ±6.144 V
	// 001  |  ±4.096 V
	// 010  |  ±2.048 V
	// 011  |  ±1.024 V
	// 100  |  ±0.512 V
	// 101  |  ±0.256 V
	// 110  |  ±0.256 V
	// 111  |  ±0.256 V

	gain = 0x0000; // 6.144 V      |  00000100 00000000

	// Input multiplexer configuration: These bits
	// configure the input multiplexer.

	// 000  |  AINP = AIN0, AINN = AIN1
	// 001  |  AINP = AIN0, AINN = AIN3
	// 010  |  AINP = AIN1, AINN = AIN3
	// 011  |  AINP = AIN2, AINN = AIN3
	// 100  |  AINP = AIN0, AINN = GND
	// 101  |  AINP = AIN1, AINN = GND
	// 110  |  AINP = AIN2, AINN = GND
	// 111  |  AINP = AIN3, AINN = GND

	mux = 0;
	if (channel == ADS01)
		mux = 0x4000; //      |  01000000 00000000
	else if (channel == ADS02)
		mux = 0x5000; //      |  01010000 00000000
	else if (channel == ADS03)
		mux = 0x6000; //      |  01100000 00000000
	else if (channel == ADS04)
		mux = 0x7000; //      |  01110000 00000000

	// Single-shot conversion, active low
	config = os | mode | rate | gain | mux;
	I2C_ADS1115_WriteRegister(0x01, config);

	// Wait until conversion finished
	while (!I2C_ADS1115_IsReady());

	return I2C_ADS1115_ReadRegister(0x00);
}

float I2C_ADS1115_ReadThermistor(int channel)
{
	word raw;
	float r1, t;
	float logr2;

	// TODO: Improve readability
	// https://en.wikipedia.org/wiki/Steinhart%E2%80%93Hart_equation

	raw = I2C_ADS1115_Read(channel);

	r1 = (float) raw / 264.8f * 0.05f;
	logr2 = log(10000.0f / (5.0f / r1 - 1.0f));
	t = 1.0f / (NTC01 + NTC02 * logr2 + NTC03 * pow(logr2, 3));

	return t - 273.15f; // Convert fahrenheit to celsius
}
