#!/usr/bin/env python3

# spi_explore.py
#
# Python3 script to exercise the spidev module and explore the SPI protocol
#
# This example code is in the public domain.
#
# Michel Deslierres <https://sigmdel.ca/michel>
#
# May 5, 2020
#
# July 10, 2020
#   Added -r --repeat option to more easily investigate the effect of
#   the /sys/module/spidev/parameters/bufsiz parameter

import spidev
import time
import argparse

DEFAULT_SPI_BUS = 0          # spidev0
DEFAULT_SPI_CS  = 0          # spidev0.0
DEFAULT_SPI_CLOCK = 1000000  # 1 Mhz
DEFAULT_SPI_MODE = 0
DEFAULT_SPI_BITS = 8

DEFAULT_BUFFER_SIZE = 4      # 32 bytes or less
DEFAULT_FIRST_BYTE  = 254
DEFAULT_TEST = "writebytes"

# write Bytes as hex bytes
def BytesToHex(Bytes):
  return ''.join(["%02X " % x for x in Bytes]).strip()


def dumpAttributes(msg):
  print()
  print(msg)
  print("  bits_per_word: {}".format(spi.bits_per_word))
  print("  cshigh: {} ".format(spi.cshigh))
  print("  loop: {}".format(spi.loop))
  print("  lsbfirst: {}".format(spi.lsbfirst))
  print("  max_speed_hz: {}".format(spi.max_speed_hz))
  print("  mode: {}".format(spi.mode))
  print("  nc_cs: {}".format(spi.no_cs))
  print("  threewire: {}".format(spi.threewire))


parser = argparse.ArgumentParser()
parser.add_argument("-b", "--bus", type=int, default=DEFAULT_SPI_BUS, metavar='', help="SPI bus (0 for spidev0, 1 for spidev1 ...)")
parser.add_argument("-c", "--cs", type=int, default=DEFAULT_SPI_CS, choices=range(0, 3), metavar='', help="Chip select (0, 1 or 2)")
parser.add_argument("-i", "--cshigh",default=False, action='store_true', help="Chip select active high")
parser.add_argument("-n", "--nocs", default=False, action='store_true', help="No chip select signal")
parser.add_argument("-s", "--speed", type=int, default=DEFAULT_SPI_CLOCK, metavar='', help="Maximum speed (Hz)")
parser.add_argument("-m", "--mode", type=int, default=DEFAULT_SPI_MODE,  choices=range(0, 4), metavar='', help="Mode (0,1,2 or 3)")
parser.add_argument("-w", "--word", type=int, default=DEFAULT_SPI_BITS,  choices=range(8, 17), metavar='', help="Bits per word (8 to 16) (read-only)")
parser.add_argument("-l", "--lsb", default=False, action='store_true', help="Least significant bits sent first, (read-only)")
parser.add_argument("-o", "--loop", default=False, action='store_true', help="Loop test, (read-only)")
parser.add_argument("-3", "--threewire", default=False, action='store_true', help="Three wire mode")
parser.add_argument("-t", "--test", type=str, default=DEFAULT_TEST, choices=["writebytes", "writebytes2", "xfer", "xfer2", "xfer3"],
                                    metavar='', help="Test: ('writebytes', 'writebytes2', 'xfer', 'xfer2' or 'xfer3')")
parser.add_argument("-r", "--repeat", type=int, default=0, metavar='', help="number of repeat transmissions, 0 (default) for endless repetitions")
parser.add_argument("-C", "--xclock", type=int, default=0, metavar='', help="xfer speed (Hz)")
parser.add_argument("-D", "--xdelay", type=int, default=0, metavar='', help="xfer delay before disactivating slave select (microseconds)")
parser.add_argument("-W", "--xword", type=int, default=8, metavar='', help="xfer bits per word, (8 only)")
parser.add_argument("-L", "--length", type=int, default=DEFAULT_BUFFER_SIZE, metavar='', help="Buffer length")
parser.add_argument("-F", "--first", type=int, default=DEFAULT_FIRST_BYTE, choices=range(0, 256), metavar='', help="First byte in buffer (0 to 256)")
parser.add_argument("-v", "--verbose", default=False, action='store_true', help="Display object attributes")
args = parser.parse_args()


# setup SPI
if args.verbose:
  spi = spidev.SpiDev()
  dumpAttributes("Attributes after spidev.SpiDev()")

  spi.open(args.bus, args.cs)
  dumpAttributes("Attributes after open({}, {})".format(args.bus, args.cs))
else:
  spi = spidev.SpiDev(args.bus, args.cs)

# set attributes
spi.bits_per_word = args.word
spi.cshigh = args.cshigh
spi.lsbfirst = args.lsb
spi.loop = args.loop
spi.max_speed_hz = args.speed
spi.mode = args.mode
spi.no_cs = args.nocs
spi.threewire = args.threewire

dumpAttributes("Attributes for test")
print()
print("Test parameters")
print("  Buffer size: {} bytes".format(args.length))
print("  First byte in buffer: {}".format(args.first))
print("  Testing method: {}".format(args.test))
if args.repeat == 0:
  print("  Repeat transmission endlessly")
else:
  print("  Repeat transmission {} times".format(args.repeat))

# set up buffer
bufferSize = args.length
bufferFirst = args.first

send = [*range(bufferSize)]
for x in range(bufferSize):
  send[x] = bufferFirst
  bufferFirst += 1
  if bufferFirst > 255:
    bufferFirst = 0
sentdata = BytesToHex(send)


# common xfer execution
def perform(trsfer):
  if args.xword:
    resp = trsfer(list(send), args.xclock, args.xdelay, args.xword)
  elif args.xdelay:
    resp = trsfer(list(send), args.xclock, args.xdelay)
  elif args.xclock:
    resp = trsfer(list(send), args.xclock)
  else:
    resp = trsfer(list(send))  # transfer a copy of send, otherwise xfer(send) will overwrite send


loopcount = 0
maxcount = args.repeat

try:
  while (maxcount == 0) or (loopcount < maxcount):
    print("")
    print("TX:", sentdata)

    # select how to write list of bytes
    if args.test == "writebytes":
      spi.writebytes(send)
    elif args.test == "writebytes2":
      spi.writebytes2(send)
    elif args.test == "xfer":
      perform(spi.xfer)
    elif args.test == "xfer2":
      perform(spi.xfer2)
    elif args.test == "xfer3":
      perform(spi.xfer3)
    else:
      print("how did this happen?")

    if (loopcount == 0) and args.verbose:
      dumpAttributes("Attributes after first transmission")
    firstLoop = False
    time.sleep(0.2)
    loopcount += 1

finally:
  spi.close()