from math import degrees

import xarm
import time

arm = xarm.Controller('USB')
print('Battery voltage of arm in volts:', arm.getBatteryVoltage())

s1 = xarm.Servo(1, 45.0) #-45 to 45?
s2 = xarm.Servo(2, 0.0) # 125 to -125
s3 = xarm.Servo(3, 0.0) # -125 to 125
s4 = xarm.Servo(4, 0.0) # -125 to 125
s5 = xarm.Servo(5, -98.50) #-50 to -125
s6 = xarm.Servo(6, 0.0) # + is clockwise

servo_list = [s1, s2, s3, s4, s5, s6]
print("servo | position")
for servo in servo_list:
    print(servo.servo_id, servo.position)

def claw_close():
    arm.setPosition(1, 45.0, wait=True)

def claw_open():
    arm.setPosition(1, -15.0, wait=True)

def set_grid_dict(x):
    # [x, y] : [s1, s2, s3, s4, s5, s6] positions
    grid_dict = {}
    for i in range(x):
        for j in range(int(x/2)):
            print("next grid point")
            input(f"Set to [{i}, {j}]")
            p1 = arm.getPosition(1, degrees=True)
            p2 = arm.getPosition(2, degrees=True)
            p3 = arm.getPosition(3, degrees=True)
            p4 = arm.getPosition(4, degrees=True)
            p5 = arm.getPosition(5, degrees=True)
            p6 = arm.getPosition(6, degrees=True)
            grid_dict[(i,j)] = [p1, p2, p3, p4, p5, p6]
    return grid_dict

def interpolate_grid(x, y, grid_dict):
    # sorted x an y from grid pts
    x_vals = sorted(set(key[0] for key in grid_dict))
    y_vals = sorted(set(key[1] for key in grid_dict))

    #find square around current location
    x1 = max([xi for xi in x_vals if xi <= x], default=None)
    x2 = min([xi for xi in x_vals if xi >= x], default=None)
    y1 = max([yi for yi in y_vals if yi <= y], default=None)
    y2 = min([yi for yi in y_vals if yi >= y], default=None)

    # Get values from the four corners
    Q11 = grid_dict[(x1, y1)]
    Q21 = grid_dict[(x2, y1)]
    Q12 = grid_dict[(x1, y2)]
    Q22 = grid_dict[(x2, y2)]

     # interpolate
    interpolated = []
    for i in range(6):
        f = (Q11[i] * (x2-x) * (y2-y) +
            Q21[i] * (x-x1) * (y2-y) +
            Q12[i] * (x2-x) * (y-y1) +
            Q22[i] * (x-x1) * (y-y1)
        ) / ((x2-x1) * (y2-y1))
        interpolated.append(f)

    return interpolated

def black_to_arm(bx, by):
    return [3/354*bx, -6/505*by+6]
def white_to_arm(wx, wy):
    return [3/508*wx+3, -3/487*wy+3]