# Based on https://github.com/badgeteam/badgePython/blob/9163a8683c1ed6029b7059eb72c7fd3d431493f0/python_modules/campzone2019/mpu6050.py

# See: https://badge.team/docs/badges/mch2022/software-development/api/bsp/

# https://docs.micropython.org/en/latest/library/machine.I2C.html
# https://www.bosch-sensortec.com/media/boschsensortec/downloads/application_notes_1/bst-bno055-an007.pdf

# https://github.com/BoschSensortec/BNO055_driver/blob/master/bno055.c

import machine,utime,ustruct

i2c = machine.I2C(0)
sensor = 0x28

def i2c_write_u8(sensor, reg, val):
    return i2c.writeto_mem(sensor, reg, val.to_bytes(1, 'little'))

def i2c_read_u8(sensor, reg):
    return int.from_bytes(i2c.readfrom_mem(sensor, reg, 1), 'little')

def i2c_read_s16le(sensor, reg):
    return ustruct.unpack("<h", i2c.readfrom_mem(sensor, reg, 2))[0]

def has_sensor():
    return sensor is not None

# /* Page id register definition */
BNO055_PAGE_ID_ADDR = 0X07

# /* PAGE0 REGISTER DEFINITION START*/
BNO055_CHIP_ID_ADDR = 0x00
BNO055_ACCEL_REV_ID_ADDR = 0x01
BNO055_MAG_REV_ID_ADDR = 0x02
BNO055_GYRO_REV_ID_ADDR = 0x03
BNO055_SW_REV_ID_LSB_ADDR = 0x04
BNO055_SW_REV_ID_MSB_ADDR = 0x05
BNO055_BL_REV_ID_ADDR = 0X06

# /* Accel data register */
BNO055_ACCEL_DATA_X_LSB_ADDR = 0X08
BNO055_ACCEL_DATA_X_MSB_ADDR = 0X09
BNO055_ACCEL_DATA_Y_LSB_ADDR = 0X0A
BNO055_ACCEL_DATA_Y_MSB_ADDR = 0X0B
BNO055_ACCEL_DATA_Z_LSB_ADDR = 0X0C
BNO055_ACCEL_DATA_Z_MSB_ADDR = 0X0D

# /* Gravity data registers */
BNO055_GRAVITY_DATA_X_LSB_ADDR = 0X2E
BNO055_GRAVITY_DATA_X_MSB_ADDR = 0X2F
BNO055_GRAVITY_DATA_Y_LSB_ADDR = 0X30
BNO055_GRAVITY_DATA_Y_MSB_ADDR = 0X31
BNO055_GRAVITY_DATA_Z_LSB_ADDR = 0X32
BNO055_GRAVITY_DATA_Z_MSB_ADDR = 0X33

# /* Temperature data register */
BNO055_TEMP_ADDR = 0X34

# /* Status registers */
BNO055_CALIB_STAT_ADDR = 0X35
BNO055_SELFTEST_RESULT_ADDR = 0X36
BNO055_INTR_STAT_ADDR = 0X37

BNO055_SYS_CLK_STAT_ADDR = 0X38
BNO055_SYS_STAT_ADDR = 0X39
BNO055_SYS_ERR_ADDR = 0X3A

# /* Unit selection register */
BNO055_UNIT_SEL_ADDR = 0X3B

# /* Mode registers */
BNO055_OPR_MODE_ADDR = 0X3D
BNO055_PWR_MODE_ADDR = 0X3E

BNO055_SYS_TRIGGER_ADDR = 0X3F
BNO055_TEMP_SOURCE_ADDR = 0X40

# /* Axis remap registers */
BNO055_AXIS_MAP_CONFIG_ADDR = 0X41
BNO055_AXIS_MAP_SIGN_ADDR = 0X42

# /** Operation mode settings **/
OPERATION_MODE_ACCONLY = 0x01


OPERATION_MODE_CONFIG = 0X00

POWER_MODE_NORMAL = 0X00

BNO055_ID = 0xA0



# /** BNO055 power settings */
# typedef enum {
# POWER_MODE_NORMAL = 0X00,
# POWER_MODE_LOWPOWER = 0X01,
# POWER_MODE_SUSPEND = 0X02
# } adafruit_bno055_powermode_t;
#
# /** Operation mode settings **/
# typedef enum {
# OPERATION_MODE_CONFIG = 0X00,
# OPERATION_MODE_ACCONLY = 0X01,
# OPERATION_MODE_MAGONLY = 0X02,
# OPERATION_MODE_GYRONLY = 0X03,
# OPERATION_MODE_ACCMAG = 0X04,
# OPERATION_MODE_ACCGYRO = 0X05,
# OPERATION_MODE_MAGGYRO = 0X06,
# OPERATION_MODE_AMG = 0X07,
# OPERATION_MODE_IMUPLUS = 0X08,
# OPERATION_MODE_COMPASS = 0X09,
# OPERATION_MODE_M4G = 0X0A,
# OPERATION_MODE_NDOF_FMC_OFF = 0X0B,
# OPERATION_MODE_NDOF = 0X0C
# } adafruit_bno055_opmode_t;
#
# /** Remap settings **/
# typedef enum {
# REMAP_CONFIG_P0 = 0x21,
# REMAP_CONFIG_P1 = 0x24, // default
# REMAP_CONFIG_P2 = 0x24,
# REMAP_CONFIG_P3 = 0x21,
# REMAP_CONFIG_P4 = 0x24,
# REMAP_CONFIG_P5 = 0x21,
# REMAP_CONFIG_P6 = 0x21,
# REMAP_CONFIG_P7 = 0x24
# } adafruit_bno055_axis_remap_config_t;
#
# /** Remap Signs **/
# typedef enum {
# REMAP_SIGN_P0 = 0x04,
# REMAP_SIGN_P1 = 0x00, // default
# REMAP_SIGN_P2 = 0x06,
# REMAP_SIGN_P3 = 0x02,
# REMAP_SIGN_P4 = 0x03,
# REMAP_SIGN_P5 = 0x01,
# REMAP_SIGN_P6 = 0x07,
# REMAP_SIGN_P7 = 0x05
# } adafruit_bno055_axis_remap_sign_t;


# init stuffs
def init():
    # Select page 0
    i2c_write_u8(sensor, BNO055_PAGE_ID_ADDR, 0)

    # Make sure we have the right device
    id_ = i2c_read_u8(sensor, BNO055_CHIP_ID_ADDR)
    if (id_ != BNO055_ID):
        raise RuntimeError('Sensor not found')

    # Switch to config mode (just in case since this is the default)
    i2c_write_u8(sensor, BNO055_OPR_MODE_ADDR, OPERATION_MODE_CONFIG)
    utime.sleep_ms(30)

    # Reset
    i2c_write_u8(sensor, BNO055_SYS_TRIGGER_ADDR, 0x20)
    utime.sleep_ms(30)

    while True:
        try:
            if (i2c_read_u8(sensor, BNO055_CHIP_ID_ADDR) == BNO055_ID):
                break
        except:
          utime.sleep_ms(10)

    utime.sleep_ms(50)

    # Set to normal power mode 
    i2c_write_u8(sensor, BNO055_PWR_MODE_ADDR, POWER_MODE_NORMAL)
    utime.sleep_ms(10)

    i2c_write_u8(sensor, BNO055_SYS_TRIGGER_ADDR, 0x0)
    utime.sleep_ms(10)

    # Set the requested operating mode (see section 3.3)
    i2c_write_u8(sensor, BNO055_OPR_MODE_ADDR, OPERATION_MODE_ACCONLY)
    utime.sleep_ms(20)


def get_accel():
    return (
        i2c_read_s16le(sensor, BNO055_ACCEL_DATA_X_LSB_ADDR),
        i2c_read_s16le(sensor, BNO055_ACCEL_DATA_Y_LSB_ADDR),
        i2c_read_s16le(sensor, BNO055_ACCEL_DATA_Z_LSB_ADDR)
    )