import csv
import threading
from time import perf_counter, sleep
import serial
serial = serial.Serial("COM5",baudrate=115200)
sleep(5)
serial.readline()
serial.readline()
serial.readline()

def set_pos(pos: float):
    serial.write(("1"+str(pos)+";").encode("ascii"))
def set_velo(velo: float):
    serial.write(("2"+str(velo)+";").encode("ascii"))
def set_Kp(Kp: float):
    serial.write(("3"+str(Kp)+";").encode("ascii"))
def set_Kd(Kd: float):
    serial.write(("4"+str(Kd)+";").encode("ascii"))
def reset_pos():
    serial.write("5;".encode("ascii"))

global theta,theta_dot,time_stamp,current,data,start_time
theta = 0
theta_dot = 0
time_stamp = 0
start_time = 0
current = 0
data = []
def decode_packet(packet: bytes):
    global theta,theta_dot,time_stamp,current,data,start_time
    dp = packet.decode("ascii")
    dspt = dp.split(',')
    try:
        time_stamp = float(dspt[0])
        theta = float(dspt[2])
        theta_dot = float(dspt[3])
        current = float(dspt[4])
        data.append([time_stamp-start_time,theta,theta_dot,current])
    except:
        pass

Kp_set = [0.3,0.4,0.5,1.0,1.5,2.0,3.0,4.0,5.0]
Kd_set = [0.0,0.05,0.1,0.15,0.20,0.25,0.30,0.35]
x_step = 2.0
time_step = 5

def read_serial():
    global theta,theta_dot,time_stamp,current,data,start_time
    while serial.in_waiting:
        decode_packet(serial.readline())
threading.Thread(target=read_serial).start()

for Kd in Kd_set:
    for Kp in Kp_set:
        dset = []
        set_Kd(Kd)
        set_Kp(Kp)
        reset_pos()
        set_pos(x_step)
        start_time = time_stamp;
        perf_start = perf_counter()
        print("Performing Experiment for Kp = "+str(Kp)+", Kd = "+str(Kd))
        while perf_counter()-perf_start < time_step:
            sleep(0.01)
        f = open("vib_result_kp"+str(Kp)+"_kd"+str(Kd)+".csv", 'w', newline='')
        writer = csv.writer(f)
        writer.writerows(data)
        f.close()
        data = []
print("Program Completed!, Data Saved to CSV")