# bma400.py - BMA400 Accelerometer Library for Python

import time
import struct
from typing import Tuple, Optional, Union, List

try:
    import smbus
except ImportError:
    pass  # Allow for non-Raspberry Pi environments


class BMA400:
    """
    Python driver for the Bosch BMA400 Accelerometer.
    
    This library provides access to the BMA400 accelerometer over I2C,
    allowing reading of acceleration data and configuration of sensor parameters.
    
    :param i2c_bus: The I2C bus number (usually 1 on Raspberry Pi 3+)
    :param address: The I2C device address, defaults to 0x14
    """
    
    # Register addresses
    _REG_WHOAMI = 0x00
    _ACC_CONFIG0 = 0x19
    _ACC_CONFIG1 = 0x1A
    _ACC_CONFIG2 = 0x1B
    _ACCEL_DATA = 0x04
    _TEMP_DATA = 0x11
    
    # Constants for conversion
    _ACC_CONVERSION = 9.80665  # Convert to m/s²
    
    # Power mode values
    SLEEP_MODE = 0x00
    LOW_POWER_MODE = 0x01
    NORMAL_MODE = 0x02
    SWITCH_TO_SLEEP = 0x03
    
    # Output data rate values
    ACCEL_12_5HZ = 0x05
    ACCEL_25HZ = 0x06
    ACCEL_50HZ = 0x07
    ACCEL_100HZ = 0x08
    ACCEL_200HZ = 0x09
    ACCEL_400HZ = 0xA4
    ACCEL_800HZ = 0xB8
    
    # Filter bandwidth values
    ACC_FILT_BW0 = 0x00  # 0.48 x ODR
    ACC_FILT_BW1 = 0x01  # 0.24 x ODR
    
    # Oversampling values
    OVERSAMPLING_0 = 0x00
    OVERSAMPLING_1 = 0x01
    OVERSAMPLING_2 = 0x02
    OVERSAMPLING_3 = 0x03
    
    # Acceleration range values
    ACC_RANGE_2 = 0x00   # ±2g
    ACC_RANGE_4 = 0x01   # ±4g
    ACC_RANGE_8 = 0x02   # ±8g
    ACC_RANGE_16 = 0x03  # ±16g
    
    # Range factors for converting raw values
    _acc_range_factor = {
        ACC_RANGE_2: 1024,
        ACC_RANGE_4: 512,
        ACC_RANGE_8: 256,
        ACC_RANGE_16: 128
    }
    
    # Data source register values
    ACC_FILT1 = 0x00
    ACC_FILT2 = 0x01
    ACC_FILT_LP = 0x02
    
    def __init__(self, i2c_bus: int = 1, address: int = 0x14):
        """
        Initialize the BMA400 sensor.
        
        :param i2c_bus: The I2C bus number
        :param address: The I2C device address
        :raises RuntimeError: If the sensor is not found
        """
        self._bus = smbus.SMBus(i2c_bus)
        self._address = address
        
        # Check device ID
        device_id = self._read_byte(BMA400._REG_WHOAMI)
        if device_id != 0x90:
            raise RuntimeError(f"Failed to find BMA400, got device ID: 0x{device_id:02X}")
        
        # Initialize with default settings
        self._power_mode = self.NORMAL_MODE
        self._acc_range = self.ACC_RANGE_2
        self._output_data_rate = self.ACCEL_100HZ
        self._oversampling_rate = self.OVERSAMPLING_3
        self._filter_bandwidth = self.ACC_FILT_BW0
        self._source_data_registers = self.ACC_FILT1
        
        # Apply default settings
        self._write_register_bits(self._ACC_CONFIG0, 0, 2, self._power_mode)
        self._write_register_bits(self._ACC_CONFIG0, 7, 1, self._filter_bandwidth)
        self._write_register_bits(self._ACC_CONFIG1, 0, 4, self._output_data_rate)
        self._write_register_bits(self._ACC_CONFIG1, 4, 2, self._oversampling_rate)
        self._write_register_bits(self._ACC_CONFIG1, 6, 2, self._acc_range)
        self._write_register_bits(self._ACC_CONFIG2, 2, 2, self._source_data_registers)
    
    def _read_byte(self, register: int) -> int:
        """Read a single byte from the specified register"""
        return self._bus.read_byte_data(self._address, register)
    
    def _read_bytes(self, register: int, length: int) -> List[int]:
        """Read multiple bytes from the specified register"""
        return self._bus.read_i2c_block_data(self._address, register, length)
    
    def _write_byte(self, register: int, value: int) -> None:
        """Write a single byte to the specified register"""
        self._bus.write_byte_data(self._address, register, value)
    
    def _write_register_bits(self, register: int, pos: int, length: int, value: int) -> None:
        """
        Write specific bits in a register
        
        :param register: Register address
        :param pos: Position of the LSB (0-indexed)
        :param length: Number of bits to modify
        :param value: Value to write
        """
        current = self._read_byte(register)
        mask = (1 << length) - 1
        current &= ~(mask << pos)  # Clear the bits
        current |= (value & mask) << pos  # Set the bits
        self._write_byte(register, current)
    
    def _read_register_bits(self, register: int, pos: int, length: int) -> int:
        """
        Read specific bits from a register
        
        :param register: Register address
        :param pos: Position of the LSB (0-indexed)
        :param length: Number of bits to read
        :return: Value of the specified bits
        """
        value = self._read_byte(register)
        mask = (1 << length) - 1
        return (value >> pos) & mask
    
    @property
    def power_mode(self) -> str:
        """
        Get the current power mode.
        
        :return: String describing the current power mode
        """
        mode = self._read_register_bits(self._ACC_CONFIG0, 0, 2)
        modes = ["SLEEP_MODE", "LOW_POWER_MODE", "NORMAL_MODE", "SWITCH_TO_SLEEP"]
        return modes[mode]
    
    @power_mode.setter
    def power_mode(self, value: int) -> None:
        """
        Set the power mode.
        
        :param value: Power mode value (use BMA400.XXX_MODE constants)
        :raises ValueError: If the value is invalid
        """
        if value not in [self.SLEEP_MODE, self.LOW_POWER_MODE, self.NORMAL_MODE, self.SWITCH_TO_SLEEP]:
            raise ValueError("Value must be a valid power mode setting")
        self._write_register_bits(self._ACC_CONFIG0, 0, 2, value)
        self._power_mode = value
    
    @property
    def output_data_rate(self) -> str:
        """
        Get the current output data rate.
        
        :return: String describing the current output data rate
        """
        odr = self._read_register_bits(self._ACC_CONFIG1, 0, 4)
        rates = {
            0x05: "ACCEL_12_5HZ",
            0x06: "ACCEL_25HZ",
            0x07: "ACCEL_50HZ",
            0x08: "ACCEL_100HZ",
            0x09: "ACCEL_200HZ",
            0xA4: "ACCEL_400HZ",
            0xB8: "ACCEL_800HZ"
        }
        return rates.get(odr, f"UNKNOWN (0x{odr:02X})")
    
    @output_data_rate.setter
    def output_data_rate(self, value: int) -> None:
        """
        Set the output data rate.
        
        :param value: Output data rate value (use BMA400.ACCEL_XXX constants)
        :raises ValueError: If the value is invalid
        """
        valid_values = [self.ACCEL_12_5HZ, self.ACCEL_25HZ, self.ACCEL_50HZ, 
                        self.ACCEL_100HZ, self.ACCEL_200HZ, self.ACCEL_400HZ, 
                        self.ACCEL_800HZ]
        if value not in valid_values:
            raise ValueError("Value must be a valid output data rate setting")
        self._write_register_bits(self._ACC_CONFIG1, 0, 4, value)
        self._output_data_rate = value
    
    @property
    def oversampling_rate(self) -> str:
        """
        Get the current oversampling rate.
        
        :return: String describing the current oversampling rate
        """
        osr = self._read_register_bits(self._ACC_CONFIG1, 4, 2)
        rates = ["OVERSAMPLING_0", "OVERSAMPLING_1", "OVERSAMPLING_2", "OVERSAMPLING_3"]
        return rates[osr]
    
    @oversampling_rate.setter
    def oversampling_rate(self, value: int) -> None:
        """
        Set the oversampling rate.
        
        :param value: Oversampling rate value (use BMA400.OVERSAMPLING_X constants)
        :raises ValueError: If the value is invalid
        """
        if value not in [self.OVERSAMPLING_0, self.OVERSAMPLING_1, 
                        self.OVERSAMPLING_2, self.OVERSAMPLING_3]:
            raise ValueError("Value must be a valid oversampling rate setting")
        self._write_register_bits(self._ACC_CONFIG1, 4, 2, value)
        self._oversampling_rate = value
    
    @property
    def acc_range(self) -> str:
        """
        Get the current acceleration range.
        
        :return: String describing the current acceleration range
        """
        range_val = self._read_register_bits(self._ACC_CONFIG1, 6, 2)
        ranges = ["ACC_RANGE_2", "ACC_RANGE_4", "ACC_RANGE_8", "ACC_RANGE_16"]
        return ranges[range_val]
    
    @acc_range.setter
    def acc_range(self, value: int) -> None:
        """
        Set the acceleration range.
        
        :param value: Acceleration range value (use BMA400.ACC_RANGE_X constants)
        :raises ValueError: If the value is invalid
        """
        if value not in [self.ACC_RANGE_2, self.ACC_RANGE_4, 
                        self.ACC_RANGE_8, self.ACC_RANGE_16]:
            raise ValueError("Value must be a valid acceleration range setting")
        self._write_register_bits(self._ACC_CONFIG1, 6, 2, value)
        self._acc_range = value
    
    @property
    def filter_bandwidth(self) -> str:
        """
        Get the current filter bandwidth.
        
        :return: String describing the current filter bandwidth
        """
        bw = self._read_register_bits(self._ACC_CONFIG0, 7, 1)
        bandwidths = ["ACC_FILT_BW0", "ACC_FILT_BW1"]
        return bandwidths[bw]
    
    @filter_bandwidth.setter
    def filter_bandwidth(self, value: int) -> None:
        """
        Set the filter bandwidth.
        
        :param value: Filter bandwidth value (use BMA400.ACC_FILT_BWx constants)
        :raises ValueError: If the value is invalid
        """
        if value not in [self.ACC_FILT_BW0, self.ACC_FILT_BW1]:
            raise ValueError("Value must be a valid filter bandwidth setting")
        self._write_register_bits(self._ACC_CONFIG0, 7, 1, value)
        self._filter_bandwidth = value
    
    @property
    def source_data_registers(self) -> str:
        """
        Get the current source data register setting.
        
        :return: String describing the current source data register
        """
        src = self._read_register_bits(self._ACC_CONFIG2, 2, 2)
        sources = ["ACC_FILT1", "ACC_FILT2", "ACC_FILT_LP"]
        return sources[src] if src < len(sources) else f"UNKNOWN (0x{src:02X})"
    
    @source_data_registers.setter
    def source_data_registers(self, value: int) -> None:
        """
        Set the source data register.
        
        :param value: Source data register value (use BMA400.ACC_FILT_X constants)
        :raises ValueError: If the value is invalid
        """
        if value not in [self.ACC_FILT1, self.ACC_FILT2, self.ACC_FILT_LP]:
            raise ValueError("Value must be a valid source data register setting")
        self._write_register_bits(self._ACC_CONFIG2, 2, 2, value)
        self._source_data_registers = value
    
    @property
    def acceleration(self) -> Tuple[float, float, float]:
        """
        Get the current acceleration measurements in m/s².
        
        :return: Tuple of (x, y, z) acceleration values
        """
        # Read 6 bytes (2 bytes per axis, x, y, z)
        data = self._read_bytes(self._ACCEL_DATA, 6)
        
        # Convert data to 16-bit signed integers
        rawx = (data[1] << 8) | data[0]
        rawy = (data[3] << 8) | data[2]
        rawz = (data[5] << 8) | data[4]
        
        # Apply two's complement if needed
        if rawx > 2047:
            rawx -= 4096
        if rawy > 2047:
            rawy -= 4096
        if rawz > 2047:
            rawz -= 4096
        
        # Convert to m/s² based on range setting
        factor = self._acc_range_factor[self._acc_range] * self._ACC_CONVERSION
        
        return (rawx / factor, rawy / factor, rawz / factor)
    
    @property
    def temperature(self) -> float:
        """
        Get the current temperature in Celsius.
        The temperature sensor is calibrated with a precision of ±5°C.
        
        :return: Temperature in Celsius
        """
        raw_temp = self._read_byte(self._TEMP_DATA)
        
        # Apply two's complement for 8-bit value
        if raw_temp > 127:
            raw_temp -= 256
        
        # Convert to Celsius according to datasheet
        return (raw_temp * 0.5) + 23.0
    
    def soft_reset(self) -> None:
        """
        Perform a soft reset of the sensor.
        This resets all registers to their default values.
        """
        # Command register and soft reset command
        CMD_REG = 0x7E
        SOFT_RESET_CMD = 0xB6
        
        self._write_byte(CMD_REG, SOFT_RESET_CMD)
        time.sleep(0.2)  # Wait for reset to complete