import struct, logging, traceback
from serial import Serial
import numpy as np

logger = logging.getLogger(__name__)

class GX5_25():

    POLL = bytes.fromhex("7565 0C07 0703 0101 05 0000 FEF8")  # without ack, returns 22 bytes

    POLL_WITH_ACK = bytes.fromhex("7565 0C07 0703 0001 05 0000 FDF3")  # with ack, returns 32 bytes; ack is in the last 10 bytes

    def __init__(self, COM='COM6', baud=115200, timeout=0.5):
        self._fd = None         # serial port file descriptor 
        self.open(COM, baud, timeout)

    def __enter__(self):
        return self

    def __exit__(self, type, value, tb):
        if type or value or tb:
            traceback.print_exception(type, value, tb)
        logger.debug('In AHRS __exit__.')
        self.close()
        return False # rethrow
    

    def open(self, COM, baud, timeout):
        try:
            self._fd = Serial(COM, baud, timeout=timeout)      
        except:
            logger.exception('AHRS serial connection failed.')
            raise


    def close(self):
        if self._fd:
            self._fd.close()

    
    def get_data(self):   # returns a 3x1 numpy array [roll,pitch,yaw], in radians
        self._fd.reset_input_buffer()
        self._fd.write(GX5_25.POLL)
        reply = self._fd.read(22)
        data_type = np.dtype(np.float32).newbyteorder('>')   # sensor sends the data as 4-byte little-endian floats
        return np.frombuffer(reply[6:18], dtype = data_type).astype(float)   # data starts at byte 6

    
    def get_data_with_ack(self):   # returns a 3x1 numpy array [roll,pitch,yaw], in radians
        self._fd.reset_input_buffer()
        self._fd.write(GX5_25.POLL_WITH_ACK)
        reply = self._fd.read(32)
        data = struct.unpack_from('>fff', reply, 6)
        error = struct.unpack_from('>b', reply, 29)[0] # the ack/nack is at index 29
        if not error:
            return np.array(data)
        else:
            return None
    
''' OTHER COMMANDS (TESTED BUT NOT NEEDED)
PING = bytes.fromhex("7565 0102 0201 E0C6") # this one and those below return a 10 byte ack
IDLE = bytes.fromhex("7565 0102 0202 E1C7")
RESUME = bytes.fromhex("7565 0102 0206 E5CB")
SET_FORMAT = bytes.fromhex("7565 0C06 060A 0101 057D  808D") # EKF, Euler (05), 4Hz (7D)
SAVE_FORMAT = bytes.fromhex("7565 0C04 040A 0300  FBFB")
STREAM_ON =  bytes.fromhex("7565 0C05 0511 010301 061E")
STREAM_OFF = bytes.fromhex("7565 0C05 0511 010300 051D")
'''
