import cv2
import numpy as np
import torch
import pyzed.sl as sl

from models.with_mobilenet import PoseEstimationWithMobileNet
from modules.keypoints import extract_keypoints, group_keypoints
from modules.load_state import load_state
from modules.pose import Pose, track_poses
from val import normalize, pad_width


def infer_fast(net, img, net_input_height_size, stride, upsample_ratio, cpu,
               pad_value=(0, 0, 0), img_mean=(128, 128, 128), img_scale=1/256):
    height, width, _ = img.shape
    scale = net_input_height_size / height

    scaled_img = cv2.resize(img, (0, 0), fx=scale, fy=scale, interpolation=cv2.INTER_CUBIC)
    scaled_img = normalize(scaled_img, img_mean, img_scale)
    min_dims = [net_input_height_size, max(scaled_img.shape[1], net_input_height_size)]
    padded_img, pad = pad_width(scaled_img, stride, pad_value, min_dims)

    tensor_img = torch.from_numpy(padded_img).permute(2, 0, 1).unsqueeze(0).float()
    if not cpu:
        tensor_img = tensor_img.cuda()

    stages_output = net(tensor_img)

    stage2_heatmaps = stages_output[-2]
    heatmaps = np.transpose(stage2_heatmaps.squeeze().cpu().data.numpy(), (1, 2, 0))
    heatmaps = cv2.resize(heatmaps, (0, 0), fx=upsample_ratio, fy=upsample_ratio, interpolation=cv2.INTER_CUBIC)

    stage2_pafs = stages_output[-1]
    pafs = np.transpose(stage2_pafs.squeeze().cpu().data.numpy(), (1, 2, 0))
    pafs = cv2.resize(pafs, (0, 0), fx=upsample_ratio, fy=upsample_ratio, interpolation=cv2.INTER_CUBIC)

    return heatmaps, pafs, scale, pad


def run_demo(net, img, height_size, cpu, track, smooth):
    if not cpu:
        net = net.cuda()

    stride = 8
    upsample_ratio = 4
    num_keypoints = Pose.num_kpts
    previous_poses = []
    delay = 33

    orig_img = img.copy()
    heatmaps, pafs, scale, pad = infer_fast(net, img, height_size, stride, upsample_ratio, cpu)

    total_keypoints_num = 0
    all_keypoints_by_type = []
    for kpt_idx in range(num_keypoints):  # 19th for bg
        total_keypoints_num += extract_keypoints(heatmaps[:, :, kpt_idx], all_keypoints_by_type, total_keypoints_num)

    pose_entries, all_keypoints = group_keypoints(all_keypoints_by_type, pafs, demo=True)
    for kpt_id in range(all_keypoints.shape[0]):
        all_keypoints[kpt_id, 0] = (all_keypoints[kpt_id, 0] * stride / upsample_ratio - pad[1]) / scale
        all_keypoints[kpt_id, 1] = (all_keypoints[kpt_id, 1] * stride / upsample_ratio - pad[0]) / scale
    current_poses = []
    for n in range(len(pose_entries)):
        if len(pose_entries[n]) == 0:
            continue
        pose_keypoints = np.ones((num_keypoints, 2), dtype=np.int32) * -1
        for kpt_id in range(num_keypoints):
            if pose_entries[n][kpt_id] != -1.0:  # keypoint was found
                pose_keypoints[kpt_id, 0] = int(all_keypoints[int(pose_entries[n][kpt_id]), 0])
                pose_keypoints[kpt_id, 1] = int(all_keypoints[int(pose_entries[n][kpt_id]), 1])

        pose = Pose(pose_keypoints, pose_entries[n][18])
        current_poses.append(pose)

        print(n, pose_keypoints.shape)

    # if track:
    #     track_poses(previous_poses, current_poses, smooth=smooth)
    #     previous_poses = current_poses
    # for pose in current_poses:
    #     pose.draw(img)
    #
    # img = cv2.addWeighted(orig_img, 0.6, img, 0.4, 0)
    # for pose in current_poses:
    #     cv2.rectangle(img, (pose.bbox[0], pose.bbox[1]),
    #                   (pose.bbox[0] + pose.bbox[2], pose.bbox[1] + pose.bbox[3]), (0, 255, 0))
    #     if track:
    #         cv2.putText(img, 'id: {}'.format(pose.id), (pose.bbox[0], pose.bbox[1] - 16),
    #                     cv2.FONT_HERSHEY_COMPLEX, 0.5, (0, 0, 255))
    # cv2.imshow('Lightweight Human Pose Estimation Python Demo', img)
    # key = cv2.waitKey(delay)
    # if key == 27:  # esc
    #     return
    # elif key == 112:  # 'p'
    #     if delay == 33:
    #         delay = 0
    #     else:
    #         delay = 33


if __name__ == '__main__':
    # ----------- init ZED camera ----------- #
    init_parameters = sl.InitParameters()
    # ----------- Read RGB & Depth map from ZED SVO file ---------- #
    init_parameters.set_from_svo_file("/home/gsx/videos/HD720_SN2684380_15-29-39.svo")

    zed = sl.Camera()
    err = zed.open(init_parameters)

    svo_image = sl.Mat()
    svo_depth = sl.Mat()

    # ----------- Init Model ---------- #
    net = PoseEstimationWithMobileNet()
    checkpoint = torch.load("./ckpt/checkpoint_iter_370000.pth")
    load_state(net, checkpoint)
    net = net.eval()

    while True:
        if zed.grab() == sl.ERROR_CODE.SUCCESS:
            zed.retrieve_image(svo_image, sl.VIEW.SIDE_BY_SIDE)
            zed.retrieve_measure(svo_depth, sl.MEASURE.DEPTH)
            svo_image_np = np.asarray(svo_image.get_data())[:, :, :3]
            svo_depth_np = np.asarray(svo_depth.get_data())

            run_demo(net, svo_image_np, 256, cpu=True, track=1, smooth=1)

