"""

laser.py

Opis: analizowanie danych wejściowych audio przy użyciu FFT i wysyłanie informacji informacji dotyczących prędkości-kierunku 
              port szeregowy. Jest to wykorzystywane do tworzenia krzywych podobnych do 
              figur Lissajous-like przy użyciu dwóch silniczników z lusterkami, 
              wskaźnika laserowego i mikrokontrolera Arduino.

Autor: Mahesh Venkitachalam
Strona WWW: electronut.in

Przykładowe uruchomienie: 

$python laser.py --port /dev/tty.usbserial-A7006Yqb
opening  /dev/tty.usbserial-A7006Yqb
otwieranie strumienia...
^Czatrzymywanie...
czyszczenie
"""

import sys, serial, struct
import pyaudio
import numpy
import math
from time import sleep
import argparse

# ręczny test dla przesyłania prędkości silniczków
def manualTest(ser):
    print('rozpoczynanie ręcznego testu...')
    try:
        while True:
            print('wpisz dane dotyczące sterowania silniczkami, np. < 100 1 120 0 >')
            strIn = raw_input()
            vals = [int(val) for val in strIn.split()[:4]]
            vals.insert(0, ord('H'))
            data = struct.pack('BBBBB', *vals)
            ser.write(data)
    except:
        print('zamykanie...')
        # wyłączenie silniczków
        vals = [ord('H'), 0, 1, 0, 1]
        data = struct.pack('BBBBB', *vals)
        ser.write(data)
        ser.close()

# automatyczny test przesyłania prędkości silniczków
def autoTest(ser):
    print('rozpoczynanie automatycznego testu...')
    try:
        while True:
            # dla każdej kombinacji kierunków
            for dr in [(0, 0), (1, 0), (0, 1), (1, 1)]:
                # dla zakresu prędkości
                for j in range(25, 180, 10):
                    for i in range(25, 180, 10):
                        vals = [ord('H'), i, dr[0], j, dr[1]]
                        print(vals[1:])
                        data = struct.pack('BBBBB', *vals)
                        ser.write(data)
                        sleep(0.1)
    except KeyboardInterrupt:
        print('zamykanie...')
        # wyłączenie silniczków
        vals = [ord('H'), 0, 1, 0, 1]
        data = struct.pack('BBBBB', *vals)
        ser.write(data)
        ser.close()


# uzyskanie urządzenia wejściowego pyaudio
def getInputDevice(p):
    index = None
    nDevices = p.get_device_count()
    print('Znaleziono urządzeń: %d. Wybierz urządzenie wejściowe:' % nDevices)
    # wyświetlenie wszystkich znalezionych urządzeń
    for i in range(nDevices):
        deviceInfo = p.get_device_info_by_index(i)
        devName = deviceInfo['name']
        print("%d: %s" % (i, devName))
    # uzyskanie wyboru użytkownika
    try:
        # konwersja na liczbę całkowitą
        index = int(input())
    except:
        pass

    # wyświetlenie nazwy wybranego urządzenia
    if index is not None:
        devName = p.get_device_info_by_index(index)["name"]
        print("Wybrane urządzenie wejściowe: %s" % devName)
    return index

# fft strumienia audio na żywo
def fftLive(ser):
  # inicjowanie pyaudio
  p = pyaudio.PyAudio()

  # uzyskanie indeksu urządzenia wejściowego pyAudio
  inputIndex = getInputDevice(p)

  # ustawienie długości próbki FFT
  fftLen = 2**11
  # ustawienie częstotliwości próbkowania
  sampleRate = 44100

  print('otwieranie strumienia...')
  stream = p.open(format = pyaudio.paInt16,
                  channels = 1,
                  rate = sampleRate,
                  input = True,
                  frames_per_buffer = fftLen,
                  input_device_index = inputIndex)
  try:
      while True:
          # odczyt fragmentu danych
          data  = stream.read(fftLen)
          # konwersja danych na tablicę numpy
          dataArray = numpy.frombuffer(data, dtype=numpy.int16)

          # uzyskanie FFT danych
          fftVals = numpy.fft.rfft(dataArray)*2.0/fftLen
          # uzyskanie bezwzględnych wartości liczb zespolonych
          fftVals = numpy.abs(fftVals)
          # uzyskanie średniej trzech zakresów częstotliwości
          # 0-100 Hz, 100-1000 Hz oraz 1000-2500 Hz
          levels = [numpy.sum(fftVals[0:100])/100,
                    numpy.sum(fftVals[100:1000])/900,
                    numpy.sum(fftVals[1000:2500])/1500]


          # wysyłane dane mają następującą postać:
          # 'H' (nagłówek), s1 (prędkość 1), d1 (kierunek 1), s2 (prędkość 2), d2 (kierunek 2)
          vals = [ord('H'), 100, 1, 100, 1]

          # The code below sets speed/direction based on information 
          # in the frequency bin. This is totally arbitary - depends 
          # on the effect you are trying to create. The goal is to 
          # generate 4 values (s1 [0-255], d1[0/1], s2[0-255], d2[0/1])
          # from the constantly changing frequency content, so the motors 
          # are in sync with the music.

          # prędkość 1
          vals[1] = int(5*levels[0]) % 255
          # prędkość 2
          vals[3] = int(100 + levels[1]) % 255

          # kierunek
          d1 = 0
          if levels[2] > 0.1:
              d1 = 1
          vals[2] = d1
          vals[4] = 0

          # pakowanie danych
          data = struct.pack('BBBBB', *vals)
          # wypisywanie danych do portu szeregowego
          ser.write(data)
          # krótka pauza
          sleep(0.001)
  except KeyboardInterrupt:
      print('zatrzymywanie...')
  finally:
      print('czyszczenie')
      stream.close()
      p.terminate()
      # wyłączanie silniczków
      vals = [ord('H'), 0, 1, 0, 1]
      data = struct.pack('BBBBB', *vals)
      ser.write(data)
      # zamykanie portu szeregowego
      ser.flush()
      ser.close()

# metoda main
def main():
    # parsowanie argumentów
    parser = argparse.ArgumentParser(description='Analizowanie wejścia audio i wysyłanie informacji dotyczących sterowania silniczkami przez port szeregowy')
    # dodawanie argumentów
    parser.add_argument('--port', dest='serial_port_name', required=True)
    parser.add_argument('--mtest', action='store_true', default=False)
    parser.add_argument('--atest', action='store_true', default=False)
    args = parser.parse_args()

    # otwieranie portu szeregowego
    strPort = args.serial_port_name
    print('otwieranie ', strPort)
    ser = serial.Serial(strPort, 9600)
    if args.mtest:
        manualTest(ser)
    elif args.atest:
        autoTest(ser)
    else:
        fftLive(ser)
        
# wywołanie funkcji main
if __name__ == '__main__':
    main()
