import matplotlib.pyplot as plt
import numpy as np
import socket
import time

# Replace with the IP address of the WeMos D1 Mini

s = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
s.connect(('localhost', 5555))



# Arm configuration
n_links = 3
gripper_angles = [0, -45, -90, 45, 90]
link_lengths = [0, 165, 105, 40]
link_colors = ['#00D000', '#0000FF', '#FF0000']
origin = (0, 0)
curr_gripper = 0

# Arm joint positions
w = [0] * n_links
z = [0] * n_links
a = [0] * n_links  # angles of links

# Target coordinates
tw, tz = 100, 100
tw0, tz0 = tw, tz

dragging = False  # Track if the mouse is dragging

# Functions for inverse kinematics calculations
def calc_p2():
    global w, z, l12
    w[2] = tw - np.cos(np.radians(gripper_angles[curr_gripper])) * link_lengths[3]
    z[2] = tz - np.sin(np.radians(gripper_angles[curr_gripper])) * link_lengths[3]
    l12 = np.sqrt(w[2]**2 + z[2]**2)

def calc_p1():
    global w, z, a
    a12 = np.arctan2(z[2], w[2])
    a[1] = np.arccos((link_lengths[1]**2 + l12**2 - link_lengths[2]**2) / (2 * link_lengths[1] * l12)) + a12
    w[1] = np.cos(a[1]) * link_lengths[1]
    z[1] = np.sin(a[1]) * link_lengths[1]
    
    # Calculate angle for link 2
    a[2] = np.arctan2(z[2] - z[1], w[2] - w[1]) - a[1]
    
    # Print the calculated angles in degrees
    print(f"Link 1 angle: {np.degrees(a[1]):.2f}°")
    print(f"Link 2 angle: {180 - abs(np.degrees(a[2])):.2f}°")
    print(f"Gripper angle:",  (abs(np.degrees(a[1])) + abs(np.degrees(a[2]))))
    #sc(0, int(np.degrees(a[1])) )
    #sc(12, int(180-(90 - abs(np.degrees(a[2])))))
   # sc(7, int (90 - abs(np.degrees(a[1])) + abs(np.degrees(a[2]))))

def draw_links(ax):
    for i in range(2):
        ax.plot([origin[0] + w[i], origin[0] + w[i+1]], 
                [origin[1] - z[i], origin[1] - z[i+1]], 
                color=link_colors[i], linewidth=2)
    ax.plot([origin[0] + w[2], origin[0] + tw], 
            [origin[1] - z[2], origin[1] - tz], 
            color=link_colors[2], linewidth=2)
    ax.plot(tw + origin[0], origin[1] - tz, 'ko')  # Draw the target point

# Mouse event functions


  # Update plot

# Main plotting function
def plot_arm():
    global tw, tz, tw0, tz0  # Enable interactive mode
    fig, ax = plt.subplots()
    
    ax.set_aspect('equal')
    ax.set_xlim(0, 1000)
    ax.set_ylim(0, 1000)
    plt.gca().invert_yaxis()
    
    
    while True:
        tw = s.recv(1024).decode('utf-8')
        tz = int(input("y:"))
        s.send("given".encode("utf-8"))
        ax.clear()
        ax.plot([origin[0], origin[0] + 259], [origin[1], origin[1]], 'k-')  # x-axis
        ax.plot([origin[0], origin[0]], [origin[1], origin[1] - 259], 'k-')  # y-axis
        
        calc_p2()
        if l12 > link_lengths[1] + link_lengths[2]:  # Out of reach check
            ax.text(240, 70, "No Solution", ha='center', color='red')
            ax.text(240, 84, "Move target closer to origin", ha='center', color='red')
            tw, tz = tw0, tz0
            calc_p2()
        
        calc_p1()
        draw_links(ax)
        
        tw0, tz0 = tw, tz
        plt.pause(0.1)  # Pause for update

plot_arm()