package mpu6050

import (
	"github.com/go-daq/smbus"
)

const (
	Gravity       = 9.80665
	AccelScale2G  = 16384.0
	AccelScale4G  = 8192.0
	AccelScale8G  = 4096.0
	AccelScale16G = 2048.0
	GyroScale250  = 131.0
	GyroScale500  = 65.5
	GyroScale1000 = 32.8
	GyroScale2000 = 16.4

	AccelRange2G  = 0x00
	AccelRange4G  = 0x08
	AccelRange8G  = 0x10
	AccelRange16G = 0x18

	GyroRange250  = 0x00
	GyroRange500  = 0x08
	GyroRange1000 = 0x10
	GyroRange2000 = 0x18

	PwrMgmt1    = 0x6B
	AccelConfig = 0x1C
	GyroConfig  = 0x1B
	TempOut0    = 0x41
	AccelXOut0  = 0x3B
	AccelYOut0  = 0x3D
	AccelZOut0  = 0x3F
	GyroXOut0   = 0x43
	GyroYOut0   = 0x45
	GyroZOut0   = 0x47
)

type MPU6050 struct {
	conn *smbus.Conn
	addr uint8
}

func New(bus int, addr uint8) (*MPU6050, error) {
	c, err := smbus.Open(bus, addr)
	if err != nil {
		return nil, err
	}
	if err := c.WriteReg(addr, PwrMgmt1, 0x00); err != nil {
		c.Close()
		return nil, err
	}
	return &MPU6050{conn: c, addr: addr}, nil
}

func (m *MPU6050) Close() error {
	return m.conn.Close()
}

func (m *MPU6050) readWord(reg byte) (int16, error) {
	hi, err := m.conn.ReadReg(m.addr, reg)
	if err != nil {
		return 0, err
	}
	lo, err := m.conn.ReadReg(m.addr, reg+1)
	if err != nil {
		return 0, err
	}

	val := int16(int16(hi)<<8 | int16(lo))

	return val, nil
}

func (m *MPU6050) GetTemp() (float64, error) {
	raw, err := m.readWord(TempOut0)
	if err != nil {
		return 0, err
	}
	return float64(raw)/340.0 + 36.53, nil
}

func (m *MPU6050) GetAccel(inG bool) (x, y, z float64, err error) {
	rx, err := m.readWord(AccelXOut0)
	if err != nil {
		return
	}
	ry, err := m.readWord(AccelYOut0)
	if err != nil {
		return
	}
	rz, err := m.readWord(AccelZOut0)
	if err != nil {
		return
	}

	cfg, err := m.conn.ReadReg(m.addr, AccelConfig)
	if err != nil {
		return
	}

	var scale float64
	switch cfg & 0x18 {
	case AccelRange4G:
		scale = AccelScale4G
	case AccelRange8G:
		scale = AccelScale8G
	case AccelRange16G:
		scale = AccelScale16G
	default:
		scale = AccelScale2G
	}

	x, y, z = float64(rx)/scale, float64(ry)/scale, float64(rz)/scale
	if !inG {
		x *= Gravity
		y *= Gravity
		z *= Gravity
	}
	return
}

func (m *MPU6050) GetGyro() (x, y, z float64, err error) {
	rx, err := m.readWord(GyroXOut0)
	if err != nil {
		return
	}
	ry, err := m.readWord(GyroYOut0)
	if err != nil {
		return
	}
	rz, err := m.readWord(GyroZOut0)
	if err != nil {
		return
	}

	cfg, err := m.conn.ReadReg(m.addr, GyroConfig)
	if err != nil {
		return
	}

	var scale float64
	switch cfg & 0x18 {
	case GyroRange500:
		scale = GyroScale500
	case GyroRange1000:
		scale = GyroScale1000
	case GyroRange2000:
		scale = GyroScale2000
	default:
		scale = GyroScale250
	}

	x, y, z = float64(rx)/scale, float64(ry)/scale, float64(rz)/scale
	return
}
