#!/usr/bin/env python
#-*- coding:utf-8 -*-

"""
  Copyright: 2016-2020 ROS小课堂 www.corvin.cn
  Description: A Python driver for the Arduino microcontroller.
  Author: corvin, jally
  History:
    20200706:创建代码.
"""
from serial.serialutil import SerialException
import thread, smbus, rospy, time, os
from serial import Serial
import sys, traceback
import numpy as np
import roslib, rospy


class Arduino:
    ''' Configuration Arduino DUE Board Parameters
    '''
    N_ANALOG_PORTS  = 10
    N_DIGITAL_PORTS = 54

    def __init__(self, is_use_serial,serial_port="/dev/ttyACM0",baudrate=57600,i2c_smbus_num=1,i2c_slave_addr=8,timeout=0.5):
        self.PID_RATE = 40   #Don't change this ! It is a fixed property of the Arduino PID controller.
        self.PID_INTERVAL = 25

        self.is_use_serial = is_use_serial #与下位机arduino的通信方式是否使用串口还是IIC
        self.serial_port = serial_port
        self.baudrate    = baudrate

        self.i2c_smbus_num  = i2c_smbus_num
        self.i2c_slave_addr = i2c_slave_addr

        self.timeout = timeout

        self.encoder_count    = 0
        self.writeTimeout     = timeout
        self.interCharTimeout = timeout/30.

        # Keep things thread safe
        self.mutex = thread.allocate_lock()

        # An array to cache analog sensor readings
        self.analog_sensor_cache = [None] * self.N_ANALOG_PORTS

        # An array to cache digital sensor readings
        self.digital_sensor_cache = [None] * self.N_DIGITAL_PORTS

    def connect(self):
        if self.is_use_serial:
            self.serial_connect()
        else:
            self.i2c_bus = smbus.SMBus(self.i2c_smbus_num)
            self.i2c_connect()

    def serial_connect(self):
        try:
            rospy.loginfo("Connecting to Arduino on port: " + self.serial_port)
            self.serial_port = Serial(port=self.serial_port, baudrate=self.baudrate, timeout=self.timeout, writeTimeout=self.writeTimeout)
            # The next line is necessary to give the firmware time to wake up.
            time.sleep(1)
            test = self.get_baud()
            if test != self.baudrate:
                time.sleep(1)
                test = self.get_baud()
                if test != self.baudrate:
                    raise SerialException

            self.beep_ring(1)
            rospy.loginfo("Arduino connected baudrate: "+str(self.baudrate))
        except SerialException:
            rospy.logerr("Cannot connect to Arduino !")
            rospy.logerr(sys.exc_info())
            rospy.logerr(traceback.print_exc(file=sys.stdout))
            os._exit(1)

    def i2c_connect(self):
        try:
            rospy.loginfo("Connecting to Arduino on IIC addr:"+str(self.i2c_slave_addr))
            test = self.i2c_get_baud()
            #test2 = self.getCurrentValue()
            #test3 = self.getBatPercent()
            #test4 = self.detect_distance()
            #test5 = self.analog_read(1)
            #test6 = self.digital_read(10)
            #test7 = self.i2c_get_pidin()
            #test8 = self.i2c_get_encoder_counts()
            #print test8
            if test != self.baudrate:
                time.sleep(1)
                test = self.i2c_get_baud()
                if test != self.baudrate:
                    raise Exception

            self.beep_ring(2)
            rospy.sleep(0.1)
            self.beep_ring(3)
            rospy.sleep(0.05)
            self.beep_ring(2)
            rospy.sleep(0.1)
            self.beep_ring(3)
            rospy.loginfo("Arduino is connected by IIC !")
        except Exception:
            rospy.logerr(sys.exc_info())
            rospy.logerr(traceback.print_exc(file=sys.stdout))
            rospy.logerr("Cannot connect to Arduino IIC slave addr !")
            os._exit(1)

    def open(self):
        ''' Open the serial port.
        '''
        self.serial_port.open()

    def close(self):
        ''' Close the serial port or i2c bus connect.
        '''
        self.beep_ring(2)
        rospy.sleep(0.5)
        self.beep_ring(3)
        if self.is_use_serial:
            self.serial_port.close()
        else:
            self.i2c_bus.close()

    def send(self, cmd):
        ''' This command should not be used on its own: it is called by the execute commands
            below in a thread safe manner.
        '''
        self.serial_port.write(cmd + '\r')

    def recv(self, timeout=0.5):
        timeout = min(timeout, self.timeout)
        ''' This command should not be used on its own: it is called by the execute commands
            below in a thread safe manner
        '''
        c = ''
        value = ''
        attempts = 0

        while c != '\r':
            c = self.serial_port.read(1)
            value += c
            attempts += 1
            if attempts * self.interCharTimeout > timeout:
                return None

        value = value.strip('\r')
        return value

    def recv_ack(self):
        ''' This command should not be used on its own: it is called by the execute commands
            below in a thread safe manner.
        '''
        ack = self.recv(self.timeout)
        return ack == 'OK'

    def recv_int(self):
        ''' This command should not be used on its own: it is called by the execute commands
            below in a thread safe manner.
        '''
        value = self.recv(self.timeout)
        try:
            return int(value)
        except:
            return None

    def recv_array(self):
        ''' This command should not be used on its own: it is called by the execute commands
            below in a thread safe manner.
        '''
        try:
            values = self.recv(self.timeout * self.N_ANALOG_PORTS).split()
            return map(float, values)
        except:
            return []

    def execute(self, cmd):
        ''' Thread safe execution of "cmd" on the Arduino returning a single integer value.
        '''
        self.mutex.acquire()

        try:
            self.serial_port.flushInput()
        except:
            pass

        ntries   = 1
        attempts = 0
        try:
            self.serial_port.write(cmd + '\r')
            value = self.recv(self.timeout)
            while attempts < ntries and (value == '' or value == 'Invalid Command' or value == None):
                try:
                    self.serial_port.flushInput()
                    self.serial_port.write(cmd + '\r')
                    value = self.recv(self.timeout)
                except:
                    rospy.logerr("Exception executing command: " + cmd)
                attempts += 1
        except:
            self.mutex.release()
            rospy.logerr("Exception executing command: " + cmd)
            value = None

        self.mutex.release()
        return value

    def execute_array(self, cmd):
        ''' Thread safe execution of "cmd" on the Arduino returning an float array.
        '''
        self.mutex.acquire()
        try:
            self.serial_port.flushInput()
        except:
            pass

        ntries   = 1
        attempts = 0
        try:
            self.serial_port.write(cmd + '\r')
            values = self.recv_array()
            while attempts < ntries and (values == '' or values == 'Invalid Command' or values == [] or values == None):
                try:
                    self.serial_port.flushInput()
                    self.serial_port.write(cmd + '\r')
                    values = self.recv_array()
                except:
                    rospy.logerr("Exception executing command: " + cmd)
                attempts += 1
        except:
            self.mutex.release()
            rospy.logerr("Exception executing command: " + cmd)
            raise SerialException
            return []

        try:
            values = map(float, values)
        except:
            values = []

        self.mutex.release()
        return values

    def execute_ack(self, cmd):
        ''' Thread safe execution of "cmd" on the Arduino returning True if response is ACK.
        '''
        self.mutex.acquire()

        try:
            self.serial_port.flushInput()
        except:
            pass

        ntries   = 1
        attempts = 0
        try:
            self.serial_port.write(cmd + '\r')
            ack = self.recv(self.timeout)
            while attempts < ntries and (ack == '' or ack == 'Invalid Command' or ack == None):
                try:
                    self.serial_port.flushInput()
                    self.serial_port.write(cmd + '\r')
                    ack = self.recv(self.timeout)
                except:
                    rospy.logerr("Exception executing command: " + cmd)
                attempts += 1
        except:
            self.mutex.release()
            rospy.logerr("execute_ack exception when executing" + cmd)
            rospy.logerr(sys.exc_info())
            return 0

        self.mutex.release()
        return ack == 'OK'

    def update_pid(self, AWheel_Kp, AWheel_Kd, AWheel_Ki, AWheel_Ko,
                         BWheel_Kp, BWheel_Kd, BWheel_Ki, BWheel_Ko ):
        ''' Set the PID parameters on the Arduino by serial
        '''
        rospy.loginfo("Updating Motors PID parameters ...")
        cmd = 'u '+str(AWheel_Kp)+':'+str(AWheel_Kd)+':'+str(AWheel_Ki)+':'+str(AWheel_Ko)+':'+str(BWheel_Kp)+':'+str(BWheel_Kd)+':'+str(BWheel_Ki)+':'+str(BWheel_Ko)
        self.execute_ack(cmd)
        #rospy.loginfo("pid params now:"+cmd)

    def i2c_update_pid(self, AWheel_Kp, AWheel_Kd, AWheel_Ki, AWheel_Ko,
                         BWheel_Kp, BWheel_Kd, BWheel_Ki, BWheel_Ko):
        rospy.loginfo("Updating PID parameters by IIC")

    def get_baud(self):
        ''' Get the current baud rate on the serial port.
        '''
        try:
            return int(self.execute('b'));
        except:
            return None

    def i2c_get_baud(self):
        ''' Get the current baud rate on the arduino IIC addr.
        '''
        try:
            #rospy.loginfo("i2c_slave_addr:"+str(self.i2c_slave_addr))
            self.i2c_bus.write_byte(self.i2c_slave_addr, ord('b'))
            #if the program ends here, maybe I/O error, please check $ sudo i2cdetect -y -a 1
            #if 0x08 does not exist all the time or the refresh is so slow, please reset arduino DUE
            self.i2c_bus.write_byte(self.i2c_slave_addr, ord('\r'))
            time.sleep(0.1)
            values = self.i2c_bus.read_i2c_block_data(self.i2c_slave_addr, 0, 5)
            values = map(chr, values)
            ret = ''.join(values)
            #print "ret:"+ret
            return int(ret)
        except:
            rospy.loginfo("get baud failed !")
            return None

    def get_encoder_counts(self):
        values = self.execute_array('e')
        if len(values) != 2:
            rospy.logerr("Encoder count was not 2")
            raise SerialException
            return None
        else:
            return values

    def i2c_get_encoder_counts(self):
        #print "IIC Get Encoder count !"

        try:
            self.i2c_bus.write_byte(self.i2c_slave_addr, int(ord('e')))
            self.i2c_bus.write_byte(self.i2c_slave_addr, ord('\r'))

            result_array = self.i2c_bus.read_i2c_block_data(self.i2c_slave_addr, 0x08)
            #There are result_flag valid bytes.
            result_flag = result_array.index(120)
            #rospy.loginfo("Valid bytes: "+str(result_flag))
            #extract result_flag bytes as strings
            result_string=''.join([chr(ec) for ec in result_array[0:(result_flag-1)]])
            #extract encoder_counts by space
            values=[int(s) for s in result_string.split(" ")]
            if len(values) != 2:
                rospy.logwarn("Encoder count was not 2")
                return None
            else:
                #rospy.loginfo("Encoder counts now: "+ str(values))
                return values
        except:
            rospy.logerr(sys.exe_info())
            traceback.print_exc(file=sys.stdout)
            return None



    def reset_encoders(self):
        ''' Reset the encoder counts to 0 by serial port
        '''
        return self.execute_ack('r')

    def i2c_reset_encoders(self):
        ''' Reset the encoder counts to 0 by IIC
        '''
        try:
            self.i2c_bus.write_byte(self.i2c_slave_addr, ord('r'))
            self.i2c_bus.write_byte(self.i2c_slave_addr, ord('\r'))
            rospy.loginfo("Reset encoders sucessfully")
        except:
            return None

    def drive(self, AWheel, BWheel):
        ''' Speeds are given in encoder ticks per PID interval
        '''
        #rospy.loginfo("drive() A:" + str(AWheel) + ",B:" + str(BWheel))
        return self.execute_ack('m %d %d' %(AWheel, BWheel))

    def i2c_drive(self, AWheel, BWheel):
        ''' Speeds are given in encoder ticks per PID interval
        '''
        #rospy.loginfo("drive() A:" + str(AWheel) + ",B:" + str(BWheel))
        cmd = (' %d %d\r' %(AWheel, BWheel))
        try:
            self.i2c_bus.write_i2c_block_data(self.i2c_slave_addr, ord('m'), [ord(c) for c in cmd])
        except:
            return None

    def stop(self):
        ''' Stop all three motors.
        '''
        self.drive(0, 0)

    def i2c_stop(self):
        cmd = (' 0 0\r')
        try:
            self.i2c_bus.write_i2c_block_data(self.i2c_slave_addr, ord('m'), [ord(c) for c in cmd])
        except:
            return None

    def analog_read(self, pin):
        if self.is_use_serial:
            return self.execute('a %d' %pin)
        else:
            cmd = (' %d\r' %(pin))
            try:
                self.i2c_bus.write_i2c_block_data(self.i2c_slave_addr, ord('a'), [ord(c) for c in cmd])
                pin = self.i2c_bus.read_i2c_block_data(self.i2c_slave_addr, 0, 3)
                pin = map(chr, pin)
                ret = ''.join(pin)
                #print "ret:"+ret
                return int(ret)
            except:
                return None

    def analog_write(self, pin, value):
        if self.is_use_serial:
            return self.execute_ack('x %d %d' %(pin, value))
        else:
            cmd = (' %d %d\r' %(pin, value))
            try:
                self.i2c_bus.write_i2c_block_data(self.i2c_slave_addr, ord('x'), [ord(c) for c in cmd])
            except:
                return None

    def digital_read(self, pin):
        if self.is_use_serial:
            return self.execute('d %d' %pin)
        else:
            cmd = (' %d\r' %(pin))
            try:
                self.i2c_bus.write_i2c_block_data(self.i2c_slave_addr, ord('d'), [ord(c) for c in cmd])
                pin = self.i2c_bus.read_i2c_block_data(self.i2c_slave_addr, 0, 2)
                pin = map(chr, pin)
                #print pin
                ret = ''.join(pin)
                #print "ret:"+ret
                return int(ret)
            except:
                return None

    def digital_write(self, pin, value):
        if self.is_use_serial:
            return self.execute_ack('w %d %d' %(pin, value))
        else:
            cmd = (' %d %d\r' %(pin, value))
            try:
                self.i2c_bus.write_i2c_block_data(self.i2c_slave_addr, ord('w'), [ord(c) for c in cmd])
            except:
                return None

    def pin_mode(self, pin, mode):
        if self.is_use_serial:
            return self.execute_ack('c %d %d' %(pin, mode))
        else:
            cmd = (' %d %d\r' %(pin, mode))
            try:
                self.i2c_bus.write_i2c_block_data(self.i2c_slave_addr, ord('c'), [ord(c) for c in cmd])
            except:
                return None

    def beep_ring(self, value):
        if self.is_use_serial:
            return self.execute_ack('p %d' %value)
        else:
            cmd = (' %d\r' %(value))
            try:
                self.i2c_bus.write_i2c_block_data(self.i2c_slave_addr, ord('p'), [ord(c) for c in cmd])
            except:
                return None


    def detect_distance(self): #检测三个红外测距传感器的返回值,一共返回三个值,分别为前,左,右三个测距值
        if self.is_use_serial:
            return self.execute_array('h')
        else:
            try:
                self.i2c_bus.write_byte(self.i2c_slave_addr, int(ord('j')))
                self.i2c_bus.write_byte(self.i2c_slave_addr, ord('\r'))

                result_array2 = self.i2c_bus.read_i2c_block_data(self.i2c_slave_addr, 0x08)
                #There are result_flag valid bytes.
                result_flag2 = result_array2.index(120)
                #rospy.loginfo("Valid bytes: "+str(result_flag2))
                #extract result_flag bytes as strings
                result_string2=''.join([chr(ch) for ch in result_array2[0:(result_flag2-1)]])
                #extract infrared_counts by space
                values=[float(s) for s in result_string2.split(" ")]
                #print values
                if len(values) != 3:
                    rospy.logwarn("InfraredEncoder count was not 3")
                    return None
                else:
                    #rospy.loginfo("Infrared counts now: "+ str(values))
                    return values
            except:
                return None


    def getBatPercent(self):
        '''get the remaining power percentage
        '''
        if self.is_use_serial:
            percent = self.execute('g')
            #rospy.loginfo("Bat percent:" + str(percent))
            return percent
        else:
            try:
                self.i2c_bus.write_byte(self.i2c_slave_addr, ord('g'))
                self.i2c_bus.write_byte(self.i2c_slave_addr, ord('\r'))
                percent = self.i2c_bus.read_i2c_block_data(self.i2c_slave_addr, 0, 3)
                percent = map(chr, percent)
                ret = ''.join(percent)
                #print "ret:"+ret
                return int(ret)
            except:
                return None

    def get_pidin(self):
        values = self.execute_array('i')
        if len(values) != 2:
            rospy.logerr("pidin was not 2")
            raise SerialException
            return None
        else:
            return values

    def i2c_get_pidin(self):
        rospy.loginfo("IIC Get Pidin !")
        ch = ''
        values = ''
        cnt = 0
        cmd = "\r"

        try:
            self.i2c_bus.write_byte(self.i2c_slave_addr, int(ord('i')))
            self.i2c_bus.write_byte(self.i2c_slave_addr, ord('\r'))

            result_array3 = self.i2c_bus.read_i2c_block_data(self.i2c_slave_addr, 0x08)
            print result_array
            #There are result_flag valid bytes.
            result_flag3 = result_array.index(120)
            #print result_flag3
            #rospy.loginfo("Valid bytes: "+str(result_flag))
            #extract result_flag bytes as strings
            result_string3=''.join([chr(ch) for ch in result_array3[0:(result_flag3-1)]])
            #print result_string3
            #extract Pidin_counts by space
            values=[int(s) for s in result_string3.split(" ")]
            if len(values) != 2:
                rospy.logwarn("Pidin was not 2")
                return None
            else:
                #rospy.loginfo("Pidin now: "+ str(values))
                return values
        except:
            rospy.logerr(sys.exe_info())
            traceback.print_exc(file=sys.stdout)
            return None

    def get_pidout(self):
        values = self.execute_array('o')
        if len(values) != 2:
            rospy.logerr("pidout was not 2")
            raise SerialException
            return None
        else:
            return values

    def i2c_get_pidout(self):
        rospy.loginfo("IIC Get Pidout !")
        ch = ''
        values = ''
        cnt = 0
        cmd = "\r"

        try:
            self.i2c_bus.write_byte(self.i2c_slave_addr, int(ord('o')))
            self.i2c_bus.write_byte(self.i2c_slave_addr, ord('\r'))

            result_array4 = self.i2c_bus.read_i2c_block_data(self.i2c_slave_addr, 0x08)
            print result_array
            #There are result_flag valid bytes.
            result_flag4 = result_array.index(120)
            #print result_flag4
            #rospy.loginfo("Valid bytes: "+str(result_flag))
            #extract result_flag bytes as strings
            result_string4=''.join([chr(ch) for ch in result_array4[0:(result_flag4-1)]])
            #print result_string4
            #extract Pidout_counts by space
            values=[int(s) for s in result_string4.split(" ")]
            if len(values) != 2:
                rospy.logwarn("Pidout was not 2")
                return None
            else:
                #rospy.loginfo("Pidout now: "+ str(values))
                return values
        except:
            rospy.logerr(sys.exe_info())
            traceback.print_exc(file=sys.stdout)
            return None

""" Basic test for connectivity by serial port or IIC bus"""
if __name__ == "__main__":
    is_use_serial = True
    portName = "/dev/ttyACM0"
    baudRate = 57600

    iic_smbus = 1
    iic_num = 8

    myArduino = Arduino(is_use_serial, serial_port=portName, baudrate=baudRate, i2c_smbus_num=iic_smbus, i2c_slave_addr=iic_num, timeout=0.5)
    myArduino.connect()

    rospy.loginfo("Sleeping for 1 second...")
    time.sleep(1)

    rospy.loginfo("Test Beep ring 3 times...")
    for i in range(3):
        myArduino.beep_ring(1)
        time.sleep(2.0)

    rospy.loginfo("Did you heard the beep ring ?")
    rospy.loginfo("Now print infrared sensors value:")
    values = myArduino.detect_distance()
    distances = np.array([values[0], values[1], values[2]])
    rospy.loginf(distances/100.0)

    myArduino.stop()
    myArduino.close()
    rospy.logwarn("Shutting down Arduino !")

