import matplotlib.pyplot as plt
from edge_finder import edge_finder, match_clusters
from webcam import get_image, image_and_crop
import cv2
from arm_control import arm, set_grid_dict, interpolate_grid, claw_open, claw_close, white_to_arm, black_to_arm


def move_to(x, y):
    print("move to")
    print(x, y)
    pos_arr = interpolate_grid(x, y, grid_dict)
    print(pos_arr)
    #return to neutral
    arm.setPosition(3, 0, wait=True)
    arm.setPosition(6, 0, wait=True)
    arm.setPosition(5, 0, wait=True)
    arm.setPosition(4, 0, wait=True)
    arm.setPosition(2, 0, wait=True)
    # arm.setPosition(1, 0, wait=True)
    #go to new
    # arm.set_position(1, pos_arr[1], wait=True)
    arm.setPosition(2, pos_arr[2], wait=True)
    arm.setPosition(3, pos_arr[3], wait=True)
    arm.setPosition(4, pos_arr[4], wait=True)
    arm.setPosition(5, pos_arr[5], wait=True)
    arm.setPosition(6, pos_arr[6], wait=True)

# Get image from camera, crop as needed
image = get_image()
black, white = image_and_crop(image, show_image=False)

grid_dict = set_grid_dict(6)

# #test images
# black = cv2.imread('black.png')
# white = cv2.imread('test_pieces.png')

# Get edges and clusters
b_edges, b_clusters, b_all_points, b_num_clusters = edge_finder(black)
w_edges, w_clusters, w_all_points, w_num_clusters = edge_finder(white)

# Match clusters and put pieces in place
for i in range(b_num_clusters):
    shape_num, shape_angle = match_clusters(i, w_all_points, w_clusters, b_all_points, b_clusters)

    #get piece coordinates
    piece_coords = w_clusters.cluster_centers_[i]
    #get hole coordinates
    hole_coords = b_clusters.cluster_centers_[shape_num]

    #get arm coords
    arm_piece_coords = white_to_arm(piece_coords[0], piece_coords[1])
    arm_hole_coords = black_to_arm(hole_coords[0], hole_coords[1])

    #move arm
    claw_open()
    move_to(arm_piece_coords[0], arm_piece_coords[1])
    claw_close()
    move_to(arm_hole_coords[0],arm_hole_coords[1])
    current_angle = arm.getPosition(2)
    arm.setPosition(2, current_angle + shape_angle)
    claw_open()

print("done")