import glob
import os
from PIL import Image
import numpy as np
from matplotlib import pyplot as plt
from io import BytesIO

def visualize_hand(all_joints, img, side=["right", "left"], n_avail_joints=21):
    # Define the connections between joints for drawing lines and their corresponding colors
    connections = [
        ((0, 1), "red"),
        ((1, 2), "green"),
        ((2, 3), "blue"),
        ((3, 4), "purple"),
        ((0, 5), "orange"),
        ((5, 6), "pink"),
        ((6, 7), "brown"),
        ((7, 8), "cyan"),
        ((0, 9), "yellow"),
        ((9, 10), "magenta"),
        ((10, 11), "lime"),
        ((11, 12), "indigo"),
        ((0, 13), "olive"),
        ((13, 14), "teal"),
        ((14, 15), "navy"),
        ((15, 16), "gray"),
        ((0, 17), "lavender"),
        ((17, 18), "silver"),
        ((18, 19), "maroon"),
        ((19, 20), "fuchsia"),
    ]
    H, W, C = img.shape

    # Create a figure and axis
    plt.figure()
    ax = plt.gca()
    # Plot joints as points
    ax.imshow(img)
    start_is = []
    if "right" in side:
        start_is.append(0)
    if "left" in side:
        start_is.append(21)
    for start_i in start_is:
        joints = all_joints[start_i : start_i + n_avail_joints]
        if len(joints) == 1:
            ax.scatter(joints[0][0], joints[0][1], color="red", s=10)
        else:
            for connection, color in connections[: len(joints) - 1]:
                joint1 = joints[connection[0]]
                joint2 = joints[connection[1]]
                ax.plot([joint1[0], joint2[0]], [joint1[1], joint2[1]], color=color, linewidth=4)

    ax.set_xlim([0, W])
    ax.set_ylim([0, H])
    ax.grid(False)
    ax.set_axis_off()
    ax.invert_yaxis()
    # plt.subplots_adjust(wspace=0.01)
    # plt.show()
    buf = BytesIO()
    plt.savefig(buf, format="png", bbox_inches="tight", pad_inches=0)
    plt.close()

    # Convert BytesIO object to numpy array
    buf.seek(0)
    img_pil = Image.open(buf)
    img_pil = img_pil.resize((W, H))
    numpy_img = np.array(img_pil)

    return numpy_img

'''put brush example at alpha channel'''
# img_dir = "bad_hands"
# masked_paths = sorted(glob.glob(os.path.join(img_dir, "*_mask.jpg")))
# for masked_pth in masked_paths:
    # img_path = masked_pth.replace("_mask.jpg", ".jpg")
    # assert os.path.exists(img_path), f"Image path {img_path} does not exist."
    # masked = np.array(Image.open(masked_pth))
    # mask = (np.all(masked > 245, axis=-1)).astype(np.uint8)*128 + 64
    # img = np.array(Image.open(img_path))
    # composite = np.concatenate((img, mask[..., None]), axis=-1)
    # composite = Image.fromarray(composite)
    # composite.save(masked_pth.replace("_mask.jpg", "_composite.png"))
    # print(f"Saved composite image {masked_pth.replace('_mask.jpg', '_composite.png')}")

'''visualize keypoint example'''
data_dir = "bad_hands"
kpts_paths = sorted(glob.glob(os.path.join(data_dir, "*.npy")))
for kpts_pth in kpts_paths:
    img_pth = kpts_pth.replace(".npy", ".jpg")
    kpts = np.load(kpts_pth)
    img = np.array(Image.open(img_pth))
    h, w = img.shape[:2]
    kpts = kpts / np.array([256,256]) * np.array([w, h])
    kpts_vis = visualize_hand(kpts, img)
    save_path = kpts_pth.replace(".npy", "_kpts.png")
    kpts_vis = Image.fromarray(kpts_vis).save(save_path)
    print(f"Saved {save_path}")