import matplotlib.pyplot as plt
import numpy as np

from skelda import utils_view

# ==================================================================================================


def show_poses2d(bodies, images, joint_names, title=""):
    num_imgs = len(images)
    rowbreak = int(num_imgs / 2.0 + 0.5)
    fig, axs = plt.subplots(2, rowbreak, figsize=(30, 20))
    fig.suptitle(title, fontsize=20)

    if isinstance(bodies, np.ndarray):
        bodies = bodies.tolist()

    # Draw skeletons into images
    for i, image in enumerate(images):
        colors = plt.cm.hsv(np.linspace(0, 1, len(bodies[i]), endpoint=False)).tolist()
        colors = [[int(c[0] * 255), int(c[1] * 255), int(c[2] * 255)] for c in colors]

        for j, body in enumerate(bodies[i]):
            image = utils_view.draw_body_in_image(image, body, joint_names, colors[j])

    # Rescale image range for plotting
    images = [img / 255.0 for img in images]

    if rowbreak == 1:
        axs[0].imshow(images[0])
        if len(images) == 2:
            axs[1].imshow(images[1])
        else:
            # Optionally delete last empty plot
            fig.delaxes(axs[1])

    else:
        for i in range(rowbreak):
            axs[0][i].imshow(images[i])
            if i + rowbreak < num_imgs:
                axs[1][i].imshow(images[i + rowbreak])
            else:
                # Optionally delete last empty plot
                fig.delaxes(axs[1][i])

    return fig