module System.Orientation.BNO055 where

import           Control.Concurrent (threadDelay)
import           Control.Monad
import           Data.Bits
import           Data.ByteString (ByteString, pack)
import qualified Data.ByteString as ByteString
import           Data.Int
import           Data.Serialize.Get
import           Data.Word
import           System.IO.I2C
import           System.Orientation.BNO055.Internal
import           Text.Printf

-- | The BNO055 has two pages of registers, so a register address may point to two different registers depending on which page is selected.
data RegisterPage = Page0 | Page1

-- | The BNO055 can have slave address 0x28 or 0x29.
--
-- On the Adafruit board without modification, it takes address 0x28.
data BNO055Address = Address0x28 | Address0x29

-- | Serialize the register page to the value for the page register
registerPage :: RegisterPage -> Word8
registerPage Page0 = bno055PageZero
registerPage Page1 = bno055PageOne

-- | Serialize the address to an I2C address
addressBNO055 :: BNO055Address -> I2CAddress
addressBNO055 Address0x28 = 0x28
addressBNO055 Address0x29 = 0x29

-- | Read a sequence of registers from the BNO055
readRegisters
  :: I2Cfd         -- ^ A handle to the I2C bus to which the BNO055 is connected
  -> BNO055Address -- ^ The address of the BNO055 on the bus
  -> RegisterPage  -- ^ The register page to read from
  -> Word8         -- ^ The first register to read
  -> Word8         -- ^ The number of registers to read
  -> IO ByteString
readRegisters fd address page register count =
  let i2cAddress = addressBNO055 address
      pageByte = registerPage page
  in do
    runI2CTransaction fd
      $ writeToSlave i2cAddress (pack [bno055PageIdAddr, pageByte])
    runI2CTransaction fd
      $  writeToSlave i2cAddress (pack [register])
      *> readFromSlave i2cAddress (fromIntegral count)

-- | Write a register on the BNO055
writeRegister
  :: I2Cfd         -- ^ A handle to the I2C bus to wich the BNO055 is connected
  -> BNO055Address -- ^ The address of the BNO055 on the bus
  -> RegisterPage  -- ^ The register page to write to
  -> Word8         -- ^ The register to write
  -> Word8         -- ^ The byte to write to the register
  -> IO ()
writeRegister fd address page register value =
  let i2cAddress = addressBNO055 address
      pageByte = registerPage page
  in do
    runI2CTransaction fd
      $ writeToSlave i2cAddress (pack [bno055PageIdAddr, pageByte])
    runI2CTransaction fd
      $ writeToSlave i2cAddress (pack [register, value])

-- | An enumeration of calibration statuses
data CalibrationStatus
  = NotCalibrated
  | PoorlyCalibrated
  | FairlyCalibrated
  | Calibrated
  deriving (Eq, Ord, Show, Enum, Bounded)

-- | A record of calibration statuses for the accelerometer, gyroscope, magnometer, and entire system
data BNO055CalibrationStatus
  = BNO055CalibrationStatus
  { systemCalibrationStatus :: CalibrationStatus
  , gyroscopeCalibrationStatus :: CalibrationStatus
  , accelerometerCalibrationStatus :: CalibrationStatus
  , magnometerCalibrationStatus :: CalibrationStatus
  } deriving (Eq, Show)

-- | Deserialize the calibration record read from the device
getBNO055CalibrationStatus :: Get BNO055CalibrationStatus
getBNO055CalibrationStatus = calibrationStatusRecordFromWord <$> getWord8
  where
    calibrationStatusFromWord :: Int -> Int -> Word8 -> CalibrationStatus
    calibrationStatusFromWord lsb msb word8 = toEnum $ (if testBit word8 msb then bit 1 else zeroBits) .|. (if testBit word8 lsb then bit 0 else zeroBits)
    calibrationStatusRecordFromWord :: Word8 -> BNO055CalibrationStatus
    calibrationStatusRecordFromWord word8
      = BNO055CalibrationStatus
        (calibrationStatusFromWord 7 6 word8)
        (calibrationStatusFromWord 5 4 word8)
        (calibrationStatusFromWord 3 2 word8)
        (calibrationStatusFromWord 1 0 word8)

-- | Read the calibration status record from the device
readCalibrationStatus
  :: I2Cfd                      -- ^ The bus to which the BNO055 is connected
  -> BNO055Address              -- ^ The address of the BNO055 on the bus
  -> IO BNO055CalibrationStatus
readCalibrationStatus fd address =
  either error id <$> runGet getBNO055CalibrationStatus <$> readRegisters fd address Page0 bno055CalibStatAddr 1

-- | Deserialize an orientation quaternion read from the device
getBNO055Quaternion :: Get (Double, Double, Double, Double)
getBNO055Quaternion =
      (,,,)
  <$> (toQuaternionComponent <$> getInt16le) -- w
  <*> (toQuaternionComponent <$> getInt16le) -- x
  <*> (toQuaternionComponent <$> getInt16le) -- y
  <*> (toQuaternionComponent <$> getInt16le) -- z
  where
    toQuaternionComponent :: Int16 -> Double
    toQuaternionComponent component = fromIntegral component / fromIntegral (2^14) -- Magic number divisor from BNO055 data sheet
  
-- | Read the fused orientation as a quaternion from the device
readQuaternion
  :: I2Cfd                               -- ^ The bus to which the BNO055 is connected
  -> BNO055Address                       -- ^ The address of the BNO055 on the bus
  -> IO (Double, Double, Double, Double) -- ^ (w, x, y, z)
readQuaternion fd address =
  either error id <$> runGet getBNO055Quaternion <$> readRegisters fd address Page0 bno055QuaternionDataWLsbAddr 8

-- | An enumeration of the operation modes of the BNO055 (incomplete)
data BNO055Mode
  = BNO055ConfigMode
  | BNO055NdofMode
  deriving (Eq, Show)

-- | Serialize the operation mode to a byte for the mode register
bno055Mode :: BNO055Mode -> Word8
bno055Mode BNO055ConfigMode = bno055OperationModeConfig
bno055Mode BNO055NdofMode = bno055OperationModeNdof

-- | Set the mode of the BNO055
setMode
  :: I2Cfd         -- ^ The bus to which the BNO055 is connected
  -> BNO055Address -- ^ The address of the BNO055 on the bus
  -> BNO055Mode    -- ^ The mode to set the device to
  -> IO ()
setMode fd address mode = writeRegister fd address Page0 bno055OprModeAddr $ bno055Mode mode

-- | An example which streams orientation and calibration status information to the console
spewData :: I2Cfd -> BNO055Address -> IO ()
spewData fd address = do
  setMode fd address BNO055NdofMode
  forever $ do
    calibrationStatus <- readCalibrationStatus fd address
    (w, x, y, z) <- readQuaternion fd address
    printf "(%6.3f, %6.3f, %6.3f, %6.3f)\t\t" w x y z
    printf "s %d g %d a %d m %d\n"
      (fromEnum $ systemCalibrationStatus calibrationStatus)
      (fromEnum $ gyroscopeCalibrationStatus calibrationStatus)
      (fromEnum $ accelerometerCalibrationStatus calibrationStatus)
      (fromEnum $ magnometerCalibrationStatus calibrationStatus)
    threadDelay 10000