import serial
import time
import numpy as np
import pandas as pd
import matplotlib.pyplot as plt
from sklearn.neighbors import KNeighborsClassifier
from sklearn.tree import DecisionTreeClassifier
import os
import pyaudio
import wave

def add_imu_data(data, accx, accy, accz):
    try:        
        if len(data) == 3 and data[0] and data[1] and data[2]:
            accx.append(float(data[0]))
            accy.append(float(data[1]))
            accz.append(float(data[2]))
    except:
        pass

    return accx, accy, accz

def save_gesture(num_samples, accx, accy, accz, training_data):
    row = accx
    row.extend(accy)
    row.extend(accz)
    row.append("")
    columns = [str(i) for i in range(3 * num_samples)]
    columns.append('gesture')
    row = pd.DataFrame([row], columns = columns)

    if training_data is None:
        row.to_csv("gestures.csv", index = False)
    else:
        training_data = training_data.append(row)
        training_data.to_csv("gestures.csv", index = False)
    
    return training_data

def playsound(filename):
    wf = wave.open(filename, 'rb')
    # Create an interface to PortAudio
    p = pyaudio.PyAudio()
    # Open a .Stream object to write the WAV file to
    # 'output = True' indicates that the sound will be played rather than recorded
    stream = p.open(format = p.get_format_from_width(wf.getsampwidth()),
                    channels = wf.getnchannels(),
                    rate = wf.getframerate(),
                    output = True)
    # Read data in chunks
    data = wf.readframes(1024)
    # Play the sound by writing the audio data to the stream
    while data != b'':
        stream.write(data)
        data = wf.readframes(1024)
    # Close and terminate the stream
    stream.close()
    p.terminate()


arduino = serial.Serial(port = 'COM10', baudrate = 115200, timeout = 0.1)
accx = []
accy = []
accz = []
num_samples = 300
mode = "test" 
x = 0

training_data = None
if os.path.exists("gestures.csv"):
    training_data = pd.read_csv("gestures.csv")

    if mode == "test":
        X = training_data.iloc[:,:-1]
        y = training_data.iloc[:,-1]
        knn = KNeighborsClassifier(n_neighbors = 3)
        knn.fit(X, y)

        clf = DecisionTreeClassifier()
        clf.fit(X, y)

while True:
    try:
        data = arduino.readline().decode('utf-8').rstrip().split(" ")
    except:
        pass

    if data[0] == "DONE":
        if mode == "train":
            arduino.write(bytes(str(-1), 'utf-8'))
            training_data = save_gesture(num_samples, accx, accy, accz, training_data)
        else:
            row = accx
            row.extend(accy)
            row.extend(accz)
            gesturek = knn.predict([row])[0]
            gesture = clf.predict([row])[0]
            print(f"prediction clf {x}: ", gesture, "prediction knn: ", gesturek)
            x += 1
            # send gesture to arduino
            if gesture == gesturek:
                arduino.write(bytes(str(gesture), 'utf-8'))
                if gesture == 1:
                    playsound("dracula_short.wav")
                elif gesture == 3:
                    playsound("disco.wav")
            else:
                arduino.write(bytes(str(-1), 'utf-8'))

        accx = []
        accy = []
        accz = []
    else:
        accx, accy, accz = add_imu_data(data, accx, accy, accz)

# training_data = pd.read_csv("gestures.csv")
# for i in range(5, len(training_data)):
#     plt.plot(training_data.iloc[i])
#     plt.show()