# 
# Author: klaute -Kai Lauterbach - @kailauterbach - me@klaute.de
# Date: 09/2016
# License: GPLv3
# 

###############################################################################

import argparse
import threading
import time
import serial
import copy
import binascii
import matplotlib.pyplot as plt
import sys
import operator

###############################################################################

parser = argparse.ArgumentParser(description='SWR meter helper tool.')
parser.add_argument("-d", "--device", type=str, help="The control device like /dev/ttyUSB0 or COM3.")
# start_freq
parser.add_argument("-s", "--start_freq", type=int, help="")
# end_freq
parser.add_argument("-e", "--end_freq", type=int, help="")
# intervall
parser.add_argument("-i", "--intervall", type=int, help="")
# step_freq
parser.add_argument("-p", "--step_freq", type=int, help="")
# drive_str
parser.add_argument("-a", "--drive_str", type=int, help="")
# start measurement
parser.add_argument("-m", "--start_meas", default=False, help="", action='store_true')
# output file (CSV)
parser.add_argument("-o", "--output_file", type=str, help="")
# show graphical results
parser.add_argument("-g", "--show_graph", default=False, help="", action='store_true')
# get config
parser.add_argument("-c", "--get_config", default=False, help="", action='store_true')
# enable clk
parser.add_argument("-l", "--enable_clk", type=int, help="")
# disable clk
parser.add_argument("-L", "--disable_clk", type=int, help="")
# save default config
parser.add_argument("-S", "--save_config", default=False, help="", action='store_true')

###############################################################################

MSG_SOD1 = 0x3c
MSG_SOD2 = 0x3e
MSG_EOD1 = 0x0d
MSG_EOD2 = 0x0a

MSG_TYPE_ANSWER_OK      = 0x01
MSG_TYPE_ANSWER_NOK     = 0x02
MSG_TYPE_MEAS_FREQ_INFO = 0x03
MSG_TYPE_MEAS_END_INFO  = 0x04
MSG_TYPE_CONFIG         = 0x05

MSG_TYPE_ANSWER_OK_DATA_TO_RECV      = 0
MSG_TYPE_ANSWER_NOK_DATA_TO_RECV     = 0
MSG_TYPE_MEAS_FREQ_INFO_DATA_TO_RECV = 8
MSG_TYPE_MEAS_END_INFO_DATA_TO_RECV  = 0
MSG_TYPE_CONFIG_DATA_TO_RECV         = 15

CC_CMD_SET_START_FREQ     = 0x01
CC_CMD_SET_END_FREQ       = 0x02
CC_CMD_SET_INTERVALL      = 0x03
CC_CMD_SET_DRIVE_STRENGTH = 0x04
CC_CMD_SET_FREQ_STEP      = 0x05
CC_CMD_START_MEASUREMENT  = 0x06
CC_CMD_GET_CONFIG         = 0x10
CC_CMD_EN_CLK             = 0x20
CC_CMD_DIS_CLK            = 0x21
CC_CMD_SAV_DFLT           = 0x22

###############################################################################

TIMEOUT_CNT_MAX     = 150
MAIN_LOOP_DELAY_S   = 0.05
THREAD_LOOP_DELAY_S = 0.01

###############################################################################

ser    = None
device = "/dev/ttyUSB0"

###############################################################################

CC_STATE_WAIT_SOD1 = 0x01
CC_STATE_WAIT_SOD2 = 0x02
CC_STATE_READ_DATA = 0x03
CC_STATE_WAIT_EOD1 = 0x04
CC_STATE_WAIT_EOD2 = 0x05
CC_STATE_GET_TYPE  = 0x06

cc_state = CC_STATE_WAIT_SOD1

cc_state_list = [ CC_STATE_WAIT_SOD1,
                  CC_STATE_WAIT_SOD2,
                  CC_STATE_READ_DATA,
                  CC_STATE_WAIT_EOD1,
                  CC_STATE_WAIT_EOD2 ]

cc_state_fn = {}

msg_type_list = [ MSG_TYPE_ANSWER_OK,
                  MSG_TYPE_ANSWER_NOK,
                  MSG_TYPE_MEAS_FREQ_INFO,
                  MSG_TYPE_CONFIG,
                  MSG_TYPE_MEAS_END_INFO, ]

msg_type_data_to_read = { MSG_TYPE_ANSWER_OK      : MSG_TYPE_ANSWER_OK_DATA_TO_RECV,
                          MSG_TYPE_ANSWER_NOK     : MSG_TYPE_ANSWER_NOK_DATA_TO_RECV,
                          MSG_TYPE_MEAS_FREQ_INFO : MSG_TYPE_MEAS_FREQ_INFO_DATA_TO_RECV,
                          MSG_TYPE_CONFIG         : MSG_TYPE_CONFIG_DATA_TO_RECV,
                          MSG_TYPE_MEAS_END_INFO  : MSG_TYPE_MEAS_END_INFO_DATA_TO_RECV, }

msg_type       = 0

cc_data_read   = 0
cc_data_buffer = []

# yes a separate counter to manage the order of the received messages
cc_message_cnt       = 0
cc_received_messages = []

###############################################################################

thread_obj     = None
thread_lock    = None
thread_started = False
thread_stop    = False

###############################################################################

def cc_init():

  global cc_state_fn
  global cc_state

  cc_state = CC_STATE_WAIT_SOD1

  cc_state_fn = { CC_STATE_WAIT_SOD1 : cc_state_fn_wait_for_sod1,
                  CC_STATE_WAIT_SOD2 : cc_state_fn_wait_for_sod2,
                  CC_STATE_WAIT_EOD1 : cc_state_fn_wait_for_eod1,
                  CC_STATE_WAIT_EOD2 : cc_state_fn_wait_for_eod2,
                  CC_STATE_GET_TYPE  : cc_state_fn_get_type,
                  CC_STATE_READ_DATA : cc_state_fn_read_data, }

########## function to call by the thread
def cc_dataReceiverThread():

  global ser
  global cc_state
  global cc_state_fn
  global thread_started
  global thread_stop

  thread_started = True

  while thread_stop == False:

    # 1. read byte from serial port into incoming
    incoming = []
    bytesToRead = ser.inWaiting()
    if bytesToRead > 0:
      incoming = list(ser.read(64))

      # 2. process the received data
      for c in incoming:
        c = int(binascii.hexlify(c), 16)

        # call the cc_state specific function to process the currently received byte
        cc_state_fn[cc_state](c)

      if cc_state not in cc_state_list:
        cc_state = CC_STATE_WAIT_SOD1

    time.sleep(THREAD_LOOP_DELAY_S)

  thread_started = False

##########
def cc_startReceiverThread():

  global thread_obj
  global thread_lock
  global thread_stop

  if thread_started == False:
    thread_lock = threading.Lock()
    thread_obj = threading.Thread(target=cc_dataReceiverThread)
    thread_obj.start()
    thread_stop = False

##########
def cc_stopReceiverThread():

  global thread_obj
  global thread_started
  global thread_stop

  if thread_started == True:
    thread_stop = True
    thread_obj.join() # wait for the thread to finish

##### CC_STATE_WAIT_SOD1
def cc_state_fn_wait_for_sod1(c):

  global cc_data_read
  global msg_type
  global cc_data_buffer
  global cc_state

  cc_data_read   = 0
  msg_type       = 0
  cc_data_buffer = []

  if c == MSG_SOD1:
    cc_state = CC_STATE_WAIT_SOD2
  else:
    cc_state = CC_STATE_WAIT_SOD1

##### CC_STATE_WAIT_SOD2
def cc_state_fn_wait_for_sod2(c):

  global cc_state

  if c == MSG_SOD2:
    cc_state = CC_STATE_GET_TYPE
  else:
    cc_state = CC_STATE_WAIT_SOD1

##### CC_STATE_GET_TYPE
def cc_state_fn_get_type(c):

  global msg_type
  global cc_state

  if c in msg_type_list:
    msg_type = c

    if msg_type_data_to_read[msg_type] > 0:
      cc_state = CC_STATE_READ_DATA
    else:
      cc_state = CC_STATE_WAIT_EOD1
  else:
    cc_state = CC_STATE_WAIT_SOD1

##### CC_STATE_READ_DATA
def cc_state_fn_read_data(c):

  global cc_data_buffer
  global cc_data_read
  global cc_state

  if cc_data_read <= msg_type_data_to_read[msg_type] - 1:
    cc_data_buffer.append(c)
    cc_data_read = cc_data_read + 1

  if cc_data_read == msg_type_data_to_read[msg_type]:
    cc_state = CC_STATE_WAIT_EOD1

##### CC_STATE_WAIT_EOD1
def cc_state_fn_wait_for_eod1(c):

  global cc_state

  if c == MSG_EOD1:
    cc_state = CC_STATE_WAIT_EOD2
  else:
    cc_state = CC_STATE_WAIT_SOD1

##### CC_STATE_WAIT_EOD2
def cc_state_fn_wait_for_eod2(c):

  global thread_lock
  global cc_message_cnt
  global cc_state

  if c == MSG_EOD2:
    is_message_read = False
    thread_lock.acquire()
    cc_received_messages.append([ cc_message_cnt,
                                  msg_type,
                                  is_message_read,
                                  copy.deepcopy(cc_data_buffer) ])
    thread_lock.release()
    cc_message_cnt = cc_message_cnt + 1

    cc_state = CC_STATE_WAIT_SOD1

###############################################################################

#####
def openSerialDevice(d):

  global ser

  try:
    if "com" in d.lower():
      d = "\\\\.\\"+d

    ser = serial.Serial(d)
  except:
    print "ERROR (1): Can't open the serial device " + d
    exit(1)

  # Toggle DTR to reset Arduino
  ser.setDTR(False)
  time.sleep(1)
  # toss any data already received, see
  # http://pyserial.sourceforge.net/pyserial_api.html#serial.Serial.flushInput
  ser.flushInput()
  ser.setDTR(True)

  try:
    if "com" in d.lower():
      ser.close()
    ser = serial.Serial(
          port=d,\
          baudrate=115200,\
          parity=serial.PARITY_NONE,\
          stopbits=serial.STOPBITS_ONE,\
          bytesize=serial.EIGHTBITS,\
          rtscts=0,\
          timeout=0)
  except:
    print "ERROR (2): Can't open the serial device " + d
    exit(2)

#####
def closeSerialDevice():

  global ser

  ser.flush()
  ser.read(1000)

  ser.close()

#####
def sendSerialData(data):

  global ser

  ser.write(bytearray([ MSG_SOD1, MSG_SOD2 ]))
  ser.write(bytearray(data))
  ser.write(bytearray([ MSG_EOD1, MSG_EOD2 ]))

  ser.flush()

###############################################################################

def user_friendly_freq(f):

  if f >= 1000000:
    return str(f / 1000000.0) + " MHz"

  elif f >= 1000:
    return str(f / 1000.0) + " kHz"

  return str(f) + " Hz"

###############################################################################

def gen_progress_bar(n, sf, ef, fs):
  steps = 1 + ((end_freq - start_freq) / step_freq)
  m = 40.0 / steps
  ret = "[ "

  for i in range(0, int(m * (steps - n))):
    ret += "#"

  for i in range(0, int(m * n)):
    ret += " "

  ret += " ] %0.2f %% " % ( (100.0 * (steps - n)/ steps) )

  return ret

###############################################################################

if __name__ == "__main__":

  start_freq = 0
  end_freq   = 0
  step_freq  = 0

  config_read = False

  meas_data = []

  cc_init()

  # parse the commandline arguments
  args = parser.parse_args()

  dataSend = 0
  timeout  = 0

  # 1. open serial device or abort
  if args.device != None:
    device = args.device

  print "SWR meter measurement software v0.1 by Kai Lauterbach (me@klaute.de)\n---\n"

  openSerialDevice(device)

  # 2. start thread to poll cc_dataReceiverThread()
  cc_startReceiverThread()

  time.sleep(1.5)

  # 3. get and process the commandline arguments/parameter
  if args.start_freq != None:
    print "Set start frequency to: " + user_friendly_freq(args.start_freq)
    sendSerialData([CC_CMD_SET_START_FREQ,
                    (args.start_freq & 0xff000000) >> 24,
                    (args.start_freq & 0x00ff0000) >> 16,
                    (args.start_freq & 0x0000ff00) >>  8,
                    (args.start_freq & 0x000000ff)])
    dataSend = dataSend + 1

  if args.end_freq != None:
    print "Set the end frequency to: " + user_friendly_freq(args.end_freq)
    sendSerialData([CC_CMD_SET_END_FREQ,
                    (args.end_freq & 0xff000000) >> 24,
                    (args.end_freq & 0x00ff0000) >> 16,
                    (args.end_freq & 0x0000ff00) >>  8,
                    (args.end_freq & 0x000000ff)])
    dataSend = dataSend + 1

  if args.step_freq != None:
    print "Set the frequency step size to: " + user_friendly_freq(args.step_freq)
    sendSerialData([CC_CMD_SET_FREQ_STEP,
                    (args.step_freq & 0xff000000) >> 24,
                    (args.step_freq & 0x00ff0000) >> 16,
                    (args.step_freq & 0x0000ff00) >>  8,
                    (args.step_freq & 0x000000ff)])
    dataSend = dataSend + 1

  if args.intervall != None:
    print "Set the time intervall to %d milliseconds" % (args.intervall)
    sendSerialData([CC_CMD_SET_INTERVALL,
                    (args.intervall & 0x0000ff00) >>  8,
                    (args.intervall & 0x000000ff)])
    dataSend = dataSend + 1

  if args.drive_str != None:
    print "Set the output drive strength to %d mA" % ((args.drive_str + 1) * 2)
    sendSerialData([CC_CMD_SET_DRIVE_STRENGTH,
                    args.drive_str])
    dataSend = dataSend + 1

  if args.start_meas == True:
    print "\nStarting the measurement process..."
    sendSerialData([CC_CMD_START_MEASUREMENT])
    dataSend = dataSend + 1

  if args.get_config == True and args.start_meas == False:
    print "Read configuration values..."
    sendSerialData([CC_CMD_GET_CONFIG])
    dataSend = dataSend + 1

  if args.enable_clk != None:
    if args.enable_clk < 0 or args.enable_clk > 2:
      args.enable_clk = 0
    print "Enabling clock output channel: %d" % (args.enable_clk)
    sendSerialData([CC_CMD_EN_CLK, args.enable_clk])
    dataSend = dataSend + 1

  if args.disable_clk != None and args.enable_clk == None:
    if args.disable_clk < 0 or args.disable_clk > 2:
      args.disable_clk = 0
    print "Disabling clock output channel: %d" % (args.disable_clk)
    sendSerialData([CC_CMD_DIS_CLK, args.disable_clk])
    dataSend = dataSend + 1

  if args.save_config == True:
    print "Save default configuration values..."
    sendSerialData([CC_CMD_SAV_DFLT])
    dataSend = dataSend + 1

  # 4. start main loop
  while dataSend > 0 and timeout < TIMEOUT_CNT_MAX:

    thread_lock.acquire()
    tmp_messages = copy.deepcopy(cc_received_messages)
    thread_lock.release()

    # 4.1 test for the response(s)
    for e in tmp_messages:
      if e[2] == False: # test for unread message

        # process it and set the data to read
        if e[1] == MSG_TYPE_ANSWER_OK:
          #print "recv: OK"
          pass

        elif e[1] == MSG_TYPE_ANSWER_NOK:
          print "recv: NOT OK on %d" % (dataSend)

        elif e[1] == MSG_TYPE_MEAS_FREQ_INFO:
          #print "recv: FREQ INFO"
          freq  = e[3][0] << 24
          freq += e[3][1] << 16
          freq += e[3][2] <<  8
          freq += e[3][3]
          a0    = e[3][4] << 8
          a0   += e[3][5]
          a1    = e[3][6] << 8
          a1   += e[3][7]
          #print "freq: " + user_friendly_freq(freq)
          #print "a0: " + str(a0)
          #print "a1: " + str(a1)
          sys.stdout.write("\r" + gen_progress_bar(dataSend, start_freq, end_freq, step_freq))
          meas_data.append([ freq, a0, a1 ])

        elif e[1] == MSG_TYPE_CONFIG:
          #print "recv: CONFIG"
          if args.start_meas == True:
            print "\nConfiguration used for measurement:"
          else:
            print "\nConfiguration:"
          start_freq  = e[3][0] << 24
          start_freq += e[3][1] << 16
          start_freq += e[3][2] <<  8
          start_freq += e[3][3]
          end_freq  = e[3][4] << 24
          end_freq += e[3][5] << 16
          end_freq += e[3][6] <<  8
          end_freq += e[3][7]
          step_freq  = e[3][8]  << 24
          step_freq += e[3][9]  << 16
          step_freq += e[3][10] <<  8
          step_freq += e[3][11]
          intervall  = e[3][12] << 8
          intervall += e[3][13]
          drive_str = e[3][14]

          print "start_freq = " + user_friendly_freq(start_freq)
          print "end_freq   = " + user_friendly_freq(end_freq)
          print "step_freq  = " + user_friendly_freq(step_freq)
          print "intervall  = " + str(intervall) + " ms"
          print "drive_str  = " + str((drive_str + 1) * 2) + " mA"
          print ""

          if args.start_meas == True and config_read == False:
            dataSend = dataSend + 1 + ((end_freq - start_freq) / step_freq)
            config_read = True

        elif e[1] == MSG_TYPE_MEAS_END_INFO:
          #print "recv: END INFO"
          sys.stdout.write("\r100.00 % done                                       \n")
          print ""

          meas_freq = []
          meas_ratio = []
          meas_r = []
          meas_a0 = []
          meas_a1 = []
          meas_ratio_f = {}

          min_vswr = [ 10, 0 ] # the default VSWR is 10 and the default freq is 0]

          i = ((drive_str + 1) * 2.0) / 1000.0

          ##### calculate the results
          if args.output_file != None or args.show_graph == True:
            for m in meas_data:

              meas_freq.append(m[0])
              vswr = 0

              meas_a0.append(m[1])
              meas_a1.append(m[2])

              if m[1] > 0 and m[2] > 0:
                if m[1] > m[2]:
                  vswr = (1.0 * m[1] / m[2])
                  meas_ratio.append(vswr)
                elif m[1] < m[2]:
                  vswr = (1.0 * m[2] / m[1])
                  meas_ratio.append(vswr)
                else:
                  vswr = 1
                  meas_ratio.append(1)
              else:
                vswr = 1
                meas_ratio.append(1)

              meas_ratio_f[m[0]] = vswr

              if vswr < min_vswr[0]:
                min_vswr[0] = vswr
                min_vswr[1] = m[0] # the frequency

              # impedance
              r = 50.0 * vswr
              meas_r.append(r)

          ##### generate the output CSV file
          if args.output_file != None:
            FILE = open(args.output_file, "w")
            FILE.write("freqency;ratio;impedance;watt;drive;a0;a1\n")
            j = 0
            for m in meas_ratio:
              FILE.write("%f;%f;%f;%f;%d;%d\n" % (meas_freq[j], m, meas_r[j], i, meas_data[j][1], meas_data[j][2]))
              j = j + 1
            FILE.close()
            print "Output file " + args.output_file + " written."

          print "First minimum VSWR %0.6f found at freqency %s" % (min_vswr[0], user_friendly_freq(min_vswr[1]))

          ##### show the graph
          if args.show_graph == True:

            f, axarr = plt.subplots(3, sharex=True)

            f.canvas.set_window_title("SWR meter measurement results")


            vswr_marker = []
            i = 0
            old = 1000
            for r in sorted(meas_ratio_f.items(), key=operator.itemgetter(1)):
            #for r in sorted(meas_ratio_f):
              vswr_marker.append(meas_freq.index(r[0]))
              if r > old:
                i += 1
              old = r
              if i == 5:
                break

            lv, = axarr[0].plot(meas_freq, meas_ratio, label='VSWR', markevery=vswr_marker, markersize=4, marker="o", markerfacecolor="r")
            lr, = axarr[1].plot(meas_freq, meas_r, label='impedance', markevery=[vswr_marker[0]], markersize=4, marker="o", markerfacecolor="r")
            la0, = axarr[2].plot(meas_freq, meas_a0, label='a0')
            la1, = axarr[2].plot(meas_freq, meas_a1, label='a1')

            axarr[0].legend(handles=[lv])
            #axarr[0].scatter(meas_freq, meas_ratio, 1)
            axarr[0].set_ylim([0, 5])
            axarr[1].legend(handles=[lr])
            #axarr[1].scatter(meas_freq, meas_r, 1)
            axarr[1].set_ylim([20, 150])
            axarr[2].legend(handles=[la0, la1])
            axarr[2].set_ylim([0, 1024])

            print "Please close the mathplot window to exit..."
            plt.show()

        else:
          print "err: unknown type 0x%02x" % (e[1])
          break

        thread_lock.acquire()
        cc_received_messages[e[0]][2] = True
        thread_lock.release()

        timeout = 0 # reset the timeout

        # reduce the number of messages to receive
        dataSend = dataSend - 1

    # manage the timeout behaviour
    time.sleep(MAIN_LOOP_DELAY_S)
    timeout = timeout + 1

  if timeout >= TIMEOUT_CNT_MAX:
    print "Timeout happened"

  # 5. stop data processing thread
  cc_stopReceiverThread()

  # 6. close serial device
  closeSerialDevice()

  exit(0)