import serial
import time
import tkinter as tk
import threading
import math

# Constants for easy maintenance
COM_PORT = 'COM9'  # Replace with your Arduino's serial port
STEP_PER_REVOLUTION = 48  # Example value, adjust as needed
JOINTS = [
    {'id': 1, 'name': 'Base', 'step_range': [0, 2838], 'step_per_degree': 8},
    {'id': 2, 'name': 'Shoulder', 'step_range': [-1000, 1000], 'step_per_degree': 8},
    {'id': 6, 'name': 'Elbow', 'step_range': [-450, 450], 'step_per_degree': 5},
    {'id': 3, 'name': 'Gripper 1', 'step_range': [0, 250]},
    {'id': 4, 'name': 'Gripper 2', 'step_range': [-180, 180]},
    {'id': 5, 'name': 'Gripper 3', 'step_range': [-180, 180]},
]
COORDINATES = {
    'x': {'min': -50, 'max': 50},
    'y': {'min': -50, 'max': 50},
    'z': {'min': 0, 'max': 200},
}

test_mode = False

def configure_serial_port():
    global test_mode, ser
    try:
        ser = serial.Serial(COM_PORT, 9600)
        time.sleep(2)
    except serial.SerialException as e:
        print(f"Error opening serial port: {e}")
        print("Starting Test Mode")
        test_mode = True

def send_potentiometer_values(values):
    if len(values) != 6:
        raise ValueError("Exactly 6 values are required")
    
    values_str = ','.join(map(str, values))
    if not test_mode:
        try:
            ser.write((values_str + '\n').encode())
            log_serial(f"-> {values_str}")
        except serial.SerialException as e:
            print(f"Error writing to serial port: {e}")
    else:
        log_serial(f"-> {values_str} (test)")

def log_serial(message, color="black"):
    serial_text.config(state=tk.NORMAL)
    serial_text.insert(tk.END, message + '\n', ("color",))
    serial_text.tag_config("color", foreground=color)
    serial_text.config(state=tk.DISABLED)
    serial_text.see(tk.END)

def read_serial():
    while True:
        if ser.in_waiting > 0:
            try:
                line = ser.readline().decode('utf-8').strip()
                log_serial(f"<- {line}")
            except serial.SerialException as e:
                print(f"Error reading from serial port: {e}")
        time.sleep(0.1)

def cartesian_to_angles(x, y, z):
    try:
        theta1 = math.atan2(y, x)
        acos_arg = (x**2 + y**2 + z**2 - 53 * z - 420.75) / 900
        if acos_arg < -1 or acos_arg > 1:
            raise ValueError(f"acos argument out of range: {acos_arg}")
        theta3 = math.acos(acos_arg)
        theta2 = math.atan2(z + 26.5, math.sqrt(x**2 + y**2)) - math.atan2(15 * math.sin(theta3), 30 + 15 * math.cos(theta3))
        theta2 = theta2 - math.pi / 2

        theta1 = theta1%(2 * math.pi)
        theta2 = theta2%(2 * math.pi)
        theta3 = theta3%(2 * math.pi)
        
        if theta2 > math.pi:
            theta2 -= 2 * math.pi
        alpha1 = theta1* 180 / math.pi
        alpha2 = theta2 * 180 / math.pi
        alpha3 = theta3 * 180 / math.pi
        log_serial(f"Angles: {alpha1:.2f}, {alpha2:.2f}, {alpha3:.2f}", "blue")
        return [alpha1, alpha2, alpha3] # in degree
    except ValueError as e:
        print(f"Error in cartesian_to_angles: {e}")
        raise e
    
def angles_to_steps(angles):
    steps = []
    for i, angle in enumerate(angles):
        step_range = JOINTS[i]['step_range']
        try:
            step_value = int(angle * JOINTS[i]["step_per_degree"])
            if step_value < step_range[0] or step_value > step_range[1]:
                raise ValueError(f"Joint {JOINTS[i]['id']} out of range: {step_value} (Range: {step_range[0]} to {step_range[1]})")
            steps.append(step_value)
        except ValueError as e:
            print(f"Error in angles_to_steps: {e}")
            raise e
    return steps

def create_main_window():
    root = tk.Tk()
    root.title("Cartesian Potentiometer Sliders" + (f" ({COM_PORT} Not Found)" if test_mode else f"({COM_PORT})"))
    return root

def create_slider_frame(root):
    slider_frame = tk.Frame(root)
    slider_frame.pack(side=tk.LEFT, padx=10, pady=10)
    return slider_frame

def create_serial_frame(root):
    serial_frame = tk.Frame(root)
    serial_frame.pack(side=tk.RIGHT, padx=10, pady=10)
    return serial_frame

def create_sliders(slider_frame):
    global sliders
    slider_values = [tk.IntVar() for _ in range(6)]
    sliders = []
    coordinates = ['x', 'y', 'z']
    for i, coord in enumerate(coordinates):
        slider = tk.Scale(
            slider_frame, 
            from_=COORDINATES[coord]['min'], 
            to=COORDINATES[coord]['max'], 
            orient=tk.HORIZONTAL, 
            variable=slider_values[i], 
            label=f"{coord.upper()} Coordinate", 
            length=400
        )
        slider.pack()
        sliders.append(slider)
    for i in range(3, 6):
        slider = tk.Scale(
            slider_frame, 
            from_=JOINTS[i]['step_range'][0], 
            to=JOINTS[i]['step_range'][1], 
            orient=tk.HORIZONTAL, 
            variable=slider_values[i], 
            label=f"Joint {JOINTS[i]['id']}: {JOINTS[i]['name']}", 
            length=400
        )
        slider.pack()
        sliders.append(slider)
    return slider_values

def home_position():
    for slider in sliders:
        slider.set(0)
    send_potentiometer_values([0, 0, 0, 0, 0, 0])

def send_stop_command():
    try:
        if not test_mode:
            ser.write("STOP\n".encode())
            log_serial("-> STOP")
        else:
            log_serial("-> STOP (test)")
    except serial.SerialException as e:
        print(f"Error writing to serial port: {e}")

def create_serial_text(serial_frame):
    serial_text = tk.Text(serial_frame, state=tk.DISABLED, width=50, height=20)
    serial_text.pack()
    return serial_text

def create_buttons(serial_frame, update_values, toggle_loop):
    button_frame = tk.Frame(serial_frame)
    button_frame.pack(side=tk.BOTTOM, pady=10)
    send_button = tk.Button(button_frame, text="Send", command=update_values)
    send_button.pack(side=tk.LEFT, padx=5)
    loop_button = tk.Button(button_frame, text="Loop: OFF", command=toggle_loop)
    loop_button.pack(side=tk.LEFT, padx=5)
    home_button = tk.Button(button_frame, text="Home", command=home_position)
    home_button.pack(side=tk.LEFT, padx=5)
    stop_button = tk.Button(button_frame, text="Stop", command=send_stop_command)
    stop_button.pack(side=tk.LEFT, padx=5)
    return loop_button

def update_values():
    try: 
        x = slider_values[0].get()
        y = slider_values[1].get()
        z = slider_values[2].get()
        angles = cartesian_to_angles(x, y, z)
        steps = angles_to_steps(angles)
        values = [steps[0], steps[1], slider_values[3].get(), slider_values[4].get(), slider_values[5].get(), steps[2]]
        send_potentiometer_values(values)
    except Exception as e:
        print(f"Error updating values: {e}")
        log_serial(f"Error updating values: {e}", "red")

def toggle_loop():
    global looping
    looping = not looping
    if looping:
        loop_button.config(text="Loop: ON")
        start_loop()
    else:
        loop_button.config(text="Loop: OFF")

def start_loop():
    if looping:
        update_values()
        root.after(1000, start_loop)

# Main execution
configure_serial_port()
root = create_main_window()
slider_frame = create_slider_frame(root)
serial_frame = create_serial_frame(root)
slider_values = create_sliders(slider_frame)
serial_text = create_serial_text(serial_frame)
loop_button = create_buttons(serial_frame, update_values, toggle_loop)

if not test_mode:
    threading.Thread(target=read_serial, daemon=True).start()

try:
    root.mainloop()
except Exception as e:
    print(f"Error in Tkinter event loop: {e}")

try:
    ser.close()
except serial.SerialException as e:
    print(f"Error closing serial port: {e}")