#----------------------------------------------------------------------------
# Copyright (c) Kauai Labs 2015. All Rights Reserved.
#
# Created in support of Team 2465 (Kauaibots).  Go Purple Wave!
#
# Open Source Software - may be modified and shared by FRC teams. Any
# modifications to this code must be accompanied by the \License.txt file
# in the root directory of the project
#----------------------------------------------------------------------------

from ._impl import AHRSProtocol, IMUProtocol, IMURegisters

from wpilib.timer import Timer

import logging
logger = logging.getLogger('navx')

__all__ = ['RegisterIO']

IO_TIMEOUT_SECONDS = 1.0



class _BoardId:
    type = 0
    hw_rev = 0
    fw_ver_major = 0
    fw_ver_minor = 0
    fw_revision = 0
    unique_id = [0]*12

class _BoardState:
    op_status = 0
    sensor_status = 0
    cal_status = 0
    selftest_status = 0
    capability_flags = 0
    update_rate_hz = 0
    accel_fsr_g = 0
    gyro_fsr_dps = 0

class RegisterIO:
    
    def __init__(self, io_provider, update_rate_hz, notify_sink, board_capabilities):
        """
            :param board_capabilities: must have the following callable attributes:
                _isOmniMountSupported, _isBoardYawResetSupported,
                _isDisplacementSupported
        
            :param notify_sink: must have the following callable attributes: 
                _setYawPitchRoll, _setAHRSData, _setAHRSPosData,
                _setRawData, _setBoardID, _setBoardState
        """
        
        self.io_provider = io_provider
        self.update_rate_hz = update_rate_hz
        self.board_capabilities = board_capabilities
        self.notify_sink = notify_sink
        
        self.raw_data_update = IMUProtocol.GyroUpdate()
        self.ahrspos_update = AHRSProtocol.AHRSPosUpdate()
        self.board_state = _BoardState()
        self.board_id = _BoardId()
        
        self.last_update_time = 0
        self.byte_count = 0
        self.update_count = 0
        self._stop = False
        
    def stop(self):
        self._stop = True
        
    def shutdown(self):
        self.io_provider.shutdown()
        
    def run(self):
        logger.info("NavX io thread starting")
        
        try:
            self.io_provider.init()
            
            # initial device configuration
            self.setUpdateRateHz(self.update_rate_hz)
            if not self.getConfiguration():
                logger.warn("-- Did not get configuration data")
            else:
                logger.info("-- Board is %s (rev %s)",
                            IMURegisters.model_type(self.board_id.type),
                            self.board_id.hw_rev)
                logger.info("-- Firmware %s.%s", self.board_id.fw_ver_major,
                                                 self.board_id.fw_ver_minor)
            
            log_error = True
            
            # IO Loop
            while not self._stop:
                if self.board_state.update_rate_hz != self.update_rate_hz:
                    self.setUpdateRateHz(self.update_rate_hz)
                
                try:
                    self.getCurrentData()
                except IOError:
                    if log_error:
                        logger.exception("Error getting data")
                        log_error = False
                else:
                    log_error = True
                    
                    
                Timer.delay(1.0/self.update_rate_hz)
        except Exception:
            logger.exception("Unhandled exception in NavX thread")
        finally:
            logger.info("NavX i/o thread exiting")
            
    def getConfiguration(self):
        success = False
        retry_count = 0
        
        while retry_count < 5 and not success:
            try:
                config = self.io_provider.read(IMURegisters.NAVX_REG_WHOAMI,
                                               IMURegisters.NAVX_REG_SENSOR_STATUS_H+1)
            except IOError as e:
                logger.warn("Error reading configuration data, retrying (%s)", e)
                success = False
                Timer.delay(0.5)
            else:
                board_id = self.board_id
                board_id.hw_rev                 = config[IMURegisters.NAVX_REG_HW_REV]
                board_id.fw_ver_major           = config[IMURegisters.NAVX_REG_FW_VER_MAJOR]
                board_id.fw_ver_minor           = config[IMURegisters.NAVX_REG_FW_VER_MINOR]
                board_id.type                   = config[IMURegisters.NAVX_REG_WHOAMI]
                self.notify_sink._setBoardID(board_id)
                
                board_state = self.board_state
                board_state.cal_status          = config[IMURegisters.NAVX_REG_CAL_STATUS]
                board_state.op_status           = config[IMURegisters.NAVX_REG_OP_STATUS]
                board_state.selftest_status     = config[IMURegisters.NAVX_REG_SELFTEST_STATUS]
                board_state.sensor_status       = AHRSProtocol.decodeBinaryUint16(config,IMURegisters.NAVX_REG_SENSOR_STATUS_L)
                board_state.gyro_fsr_dps        = AHRSProtocol.decodeBinaryUint16(config,IMURegisters.NAVX_REG_GYRO_FSR_DPS_L)
                board_state.accel_fsr_g         = config[IMURegisters.NAVX_REG_ACCEL_FSR_G]
                board_state.update_rate_hz      = config[IMURegisters.NAVX_REG_UPDATE_RATE_HZ]
                board_state.capability_flags    = AHRSProtocol.decodeBinaryUint16(config,IMURegisters.NAVX_REG_CAPABILITY_FLAGS_L)
                self.notify_sink._setBoardState(board_state)
                success = True
        
            retry_count += 1
        
        return success
    
    def getCurrentData(self):
        first_address = IMURegisters.NAVX_REG_UPDATE_RATE_HZ
        displacement_registers = self.board_capabilities._isDisplacementSupported()
        
        # If firmware supports displacement data, acquire it - otherwise implement
        # similar (but potentially less accurate) calculations on this processor.
        if displacement_registers:
            read_count = IMURegisters.NAVX_REG_LAST + 1 - first_address
        else:
            read_count = IMURegisters.NAVX_REG_QUAT_OFFSET_Z_H + 1 - first_address
        
        curr_data = self.io_provider.read(first_address, read_count)
    
        #timestamp_low = AHRSProtocol.decodeBinaryUint16(curr_data, IMURegisters.NAVX_REG_TIMESTAMP_L_L-first_address)
        #timestamp_high = AHRSProtocol.decodeBinaryUint16(curr_data, IMURegisters.NAVX_REG_TIMESTAMP_H_L-first_address)
        #sensor_timestamp            = (timestamp_high << 16) + timestamp_low
        ahrspos_update = self.ahrspos_update
        ahrspos_update.op_status       = curr_data[IMURegisters.NAVX_REG_OP_STATUS - first_address]
        ahrspos_update.selftest_status = curr_data[IMURegisters.NAVX_REG_SELFTEST_STATUS - first_address]
        ahrspos_update.cal_status      = curr_data[IMURegisters.NAVX_REG_CAL_STATUS]
        ahrspos_update.sensor_status   = curr_data[IMURegisters.NAVX_REG_SENSOR_STATUS_L - first_address]
        ahrspos_update.yaw             = AHRSProtocol.decodeProtocolSignedHundredthsFloat(curr_data, IMURegisters.NAVX_REG_YAW_L-first_address)
        ahrspos_update.pitch           = AHRSProtocol.decodeProtocolSignedHundredthsFloat(curr_data, IMURegisters.NAVX_REG_PITCH_L-first_address)
        ahrspos_update.roll            = AHRSProtocol.decodeProtocolSignedHundredthsFloat(curr_data, IMURegisters.NAVX_REG_ROLL_L-first_address)
        ahrspos_update.compass_heading = AHRSProtocol.decodeProtocolUnsignedHundredthsFloat(curr_data, IMURegisters.NAVX_REG_HEADING_L-first_address)
        ahrspos_update.mpu_temp_c      = AHRSProtocol.decodeProtocolSignedHundredthsFloat(curr_data, IMURegisters.NAVX_REG_MPU_TEMP_C_L - first_address)
        ahrspos_update.world_linear_accel_x  = AHRSProtocol.decodeProtocolSignedThousandthsFloat(curr_data, IMURegisters.NAVX_REG_LINEAR_ACC_X_L-first_address)
        ahrspos_update.world_linear_accel_y  = AHRSProtocol.decodeProtocolSignedThousandthsFloat(curr_data, IMURegisters.NAVX_REG_LINEAR_ACC_Y_L-first_address)
        ahrspos_update.world_linear_accel_z  = AHRSProtocol.decodeProtocolSignedThousandthsFloat(curr_data, IMURegisters.NAVX_REG_LINEAR_ACC_Z_L-first_address)
        ahrspos_update.altitude        = AHRSProtocol.decodeProtocol1616Float(curr_data, IMURegisters.NAVX_REG_ALTITUDE_D_L - first_address)
        ahrspos_update.baro_pressure   = AHRSProtocol.decodeProtocol1616Float(curr_data, IMURegisters.NAVX_REG_PRESSURE_DL - first_address)
        ahrspos_update.fused_heading   = AHRSProtocol.decodeProtocolUnsignedHundredthsFloat(curr_data, IMURegisters.NAVX_REG_FUSED_HEADING_L-first_address)
        ahrspos_update.quaternionW     = AHRSProtocol.decodeBinaryInt16(curr_data, IMURegisters.NAVX_REG_QUAT_W_L-first_address)
        ahrspos_update.quaternionX     = AHRSProtocol.decodeBinaryInt16(curr_data, IMURegisters.NAVX_REG_QUAT_X_L-first_address)
        ahrspos_update.quaternionY     = AHRSProtocol.decodeBinaryInt16(curr_data, IMURegisters.NAVX_REG_QUAT_Y_L-first_address)
        ahrspos_update.quaternionZ     = AHRSProtocol.decodeBinaryInt16(curr_data, IMURegisters.NAVX_REG_QUAT_Z_L-first_address)
        if displacement_registers:
            ahrspos_update.vel_x       = AHRSProtocol.decodeProtocol1616Float(curr_data, IMURegisters.NAVX_REG_VEL_X_I_L-first_address)
            ahrspos_update.vel_y       = AHRSProtocol.decodeProtocol1616Float(curr_data, IMURegisters.NAVX_REG_VEL_Y_I_L-first_address)
            ahrspos_update.vel_z       = AHRSProtocol.decodeProtocol1616Float(curr_data, IMURegisters.NAVX_REG_VEL_Z_I_L-first_address)
            ahrspos_update.disp_x      = AHRSProtocol.decodeProtocol1616Float(curr_data, IMURegisters.NAVX_REG_DISP_X_I_L-first_address)
            ahrspos_update.disp_y      = AHRSProtocol.decodeProtocol1616Float(curr_data, IMURegisters.NAVX_REG_DISP_Y_I_L-first_address)
            ahrspos_update.disp_z      = AHRSProtocol.decodeProtocol1616Float(curr_data, IMURegisters.NAVX_REG_DISP_Z_I_L-first_address)            
             
            self.notify_sink._setAHRSPosData(ahrspos_update)
        else:
            self.notify_sink._setAHRSData(ahrspos_update)
        
        board_state = self.board_state
        board_state.cal_status      = curr_data[IMURegisters.NAVX_REG_CAL_STATUS-first_address]
        board_state.op_status       = curr_data[IMURegisters.NAVX_REG_OP_STATUS-first_address]
        board_state.selftest_status = curr_data[IMURegisters.NAVX_REG_SELFTEST_STATUS-first_address]
        board_state.sensor_status   = AHRSProtocol.decodeBinaryUint16(curr_data,IMURegisters.NAVX_REG_SENSOR_STATUS_L-first_address)
        board_state.update_rate_hz  = curr_data[IMURegisters.NAVX_REG_UPDATE_RATE_HZ-first_address]
        board_state.gyro_fsr_dps    = AHRSProtocol.decodeBinaryUint16(curr_data,IMURegisters.NAVX_REG_GYRO_FSR_DPS_L)
        board_state.accel_fsr_g     = curr_data[IMURegisters.NAVX_REG_ACCEL_FSR_G]
        board_state.capability_flags= AHRSProtocol.decodeBinaryUint16(curr_data,IMURegisters.NAVX_REG_CAPABILITY_FLAGS_L-first_address)
        self.notify_sink._setBoardState(board_state)
        
        raw_data_update = self.raw_data_update
        raw_data_update.raw_gyro_x      = AHRSProtocol.decodeBinaryInt16(curr_data,  IMURegisters.NAVX_REG_GYRO_X_L-first_address)
        raw_data_update.raw_gyro_y      = AHRSProtocol.decodeBinaryInt16(curr_data,  IMURegisters.NAVX_REG_GYRO_Y_L-first_address)
        raw_data_update.raw_gyro_z      = AHRSProtocol.decodeBinaryInt16(curr_data,  IMURegisters.NAVX_REG_GYRO_Z_L-first_address)
        raw_data_update.raw_accel_x     = AHRSProtocol.decodeBinaryInt16(curr_data,  IMURegisters.NAVX_REG_ACC_X_L-first_address)
        raw_data_update.raw_accel_y     = AHRSProtocol.decodeBinaryInt16(curr_data,  IMURegisters.NAVX_REG_ACC_Y_L-first_address)
        raw_data_update.raw_accel_z     = AHRSProtocol.decodeBinaryInt16(curr_data,  IMURegisters.NAVX_REG_ACC_Z_L-first_address)
        raw_data_update.cal_mag_x       = AHRSProtocol.decodeBinaryInt16(curr_data,  IMURegisters.NAVX_REG_MAG_X_L-first_address)
        raw_data_update.cal_mag_y       = AHRSProtocol.decodeBinaryInt16(curr_data,  IMURegisters.NAVX_REG_MAG_Y_L-first_address)
        raw_data_update.cal_mag_z       = AHRSProtocol.decodeBinaryInt16(curr_data,  IMURegisters.NAVX_REG_MAG_Z_L-first_address)
        raw_data_update.mpu_temp_c      = ahrspos_update.mpu_temp
        self.notify_sink._setRawData(raw_data_update)
        
        self.last_update_time = Timer.getFPGATimestamp()
        self.byte_count += len(curr_data)
        self.update_count += 1
        
    def isConnected(self):
        time_since_last_update = Timer.getFPGATimestamp() - self.last_update_time
        return time_since_last_update <= IO_TIMEOUT_SECONDS
    
    def getByteCount(self):
        return self.byte_count
    
    def getUpdateCount(self):
        return self.update_count
    
    def setUpdateRateHz(self, update_rate):
        self.io_provider.write(IMURegisters.NAVX_REG_UPDATE_RATE_HZ, update_rate)
    
    def zeroYaw(self):
        self.io_provider.write( IMURegisters.NAVX_REG_INTEGRATION_CTL, 
                                   AHRSProtocol.NAVX_INTEGRATION_CTL_RESET_YAW )

    def zeroDisplacement(self):
        self.io_provider.write( IMURegisters.NAVX_REG_INTEGRATION_CTL, 
                                (AHRSProtocol.NAVX_INTEGRATION_CTL_RESET_DISP_X |
                                 AHRSProtocol.NAVX_INTEGRATION_CTL_RESET_DISP_Y |
                                 AHRSProtocol.NAVX_INTEGRATION_CTL_RESET_DISP_Z ) )       
