import time
from tbd_dataloader import TBDv2Dataset

import numpy as np

import matplotlib.pyplot as plt
import matplotlib.patches as patches

def point2screen(points):
    K = [607.13232421875, 0.0, 638.6468505859375, 0.0, 607.1067504882812, 367.1607360839844, 0.0, 0.0, 1.0]
    K = np.reshape(np.array(K), [3, 3])
    rot_points = np.array(points) + np.array([0, 0.2, 0])
    rot_points = rot_points
    points_camera = rot_points.reshape(3, 1)

    project_matrix = np.array(K).reshape(3, 3)
    points_prj = project_matrix.dot(points_camera)
    points_prj = points_prj.transpose()
    if not points_prj[:, 2][0] == 0.0:
        points_prj[:, 0] = points_prj[:, 0] / points_prj[:, 2]
        points_prj[:, 1] = points_prj[:, 1] / points_prj[:, 2]
    points_screen = points_prj[:, :2]
    assert points_screen.shape == (1, 2)
    points_screen = points_screen.reshape(-1)
    return points_screen

if __name__ == '__main__':
    data = TBDv2Dataset(number_frames_to_sample=1, resize_img=None)
    index = np.random.randint(0, len(data))
    start = time.time()
    (
        kinect_imgs, # <- len x 720 x 1280 x 3
        tracker_imgs,
        battery_imgs,
        skele1,
        skele2,
        bbox,
        tracker_skeID_sample, # <- This is the tracker skeleton ID 
        tracker2d,
        label,
        experiment_id, # From here for debugging
        timestep,
        obj_id, # <- This is the object ID as a string
    ) = data[index]
    end = time.time()
    print(f"Time for one sample: {end-start}")

    img = kinect_imgs[-1]
    bbox = bbox[-1]
    print(label.shape)

    print(skele1.shape)
    print(skele2.shape)

    skele1 = skele1[-1, :,:]
    skele2 = skele2[-1, :,:]

    print(skele1.shape)

    

    # reshape img from c, h, w to h, w, c
    img = img.permute(1, 2, 0)

    fig, ax = plt.subplots(1)
    ax.imshow(img)
    print(bbox[0], bbox[1], bbox[2], bbox[3]) # t(top left x, top left y, width, height)
    top_left_x, top_left_y, width, height = bbox[0], bbox[1], bbox[2], bbox[3]
    x_min, y_min, x_max, y_max = bbox[0], bbox[1], bbox[2], bbox[3]

    


    for i in range(26):
        print(skele1[i,0], skele1[i,1])
        print(skele1[i,:].shape)
        print(point2screen(skele1[i,:]))
        x, y = point2screen(skele1[i,:])[0], point2screen(skele1[i,:])[1]
        ax.text(x, y, f"{i}", fontsize=5, color='w')

        wedge = patches.Wedge((x,y), 10, 0, 360, width=10, color='b')
        ax.add_patch(wedge)

    for i in range(26):
        x, y = point2screen(skele2[i,:])[0], point2screen(skele2[i,:])[1]
        ax.text(x, y, f"{i}", fontsize=5, color='w')
        wedge = patches.Wedge((point2screen(skele2[i,:])[0], point2screen(skele2[i,:])[1]), 10, 0, 360, width=10, color='r')
        ax.add_patch(wedge)

    # Create a Rectangle patch
    # rect = patches.Rectangle((top_left_x, top_left_y-height), width, height, linewidth=1, edgecolor='r', facecolor='none')
    # ax.add_patch(rect)
    # rect = patches.Rectangle((x_min, y_max), x_max-x_min, y_max-y_min, linewidth=1, edgecolor='g', facecolor='none')
    # ax.add_patch(rect)
    fig.savefig(f"bbox_{obj_id}_{index}_{experiment_id}.png")