import sys
import os
import time
from collections import OrderedDict

import numpy as np
import logging
import argparse
import json, re
from tqdm import tqdm

from tf_pose.common import read_imgfile
from tf_pose.estimator import TfPoseEstimator
from tf_pose.networks import model_wh, get_graph_path

from pycocotools.coco import COCO
from pycocotools.cocoeval import COCOeval

logging.basicConfig(level=logging.INFO, format="%(asctime)s %(levelname)s %(message)s")

logger = logging.getLogger("TfPoseEstimator-Video")
logger.setLevel(logging.INFO)
ch = logging.StreamHandler()
ch.setLevel(logging.DEBUG)
formatter = logging.Formatter("[%(asctime)s] [%(name)s] [%(levelname)s] %(message)s")
ch.setFormatter(formatter)
logger.addHandler(ch)

eval_size = -1


def round_int(val):
    return int(round(val))


def write_coco_json(human, image_w, image_h):
    keypoints = []
    coco_ids = [0, 15, 14, 17, 16, 5, 2, 6, 3, 7, 4, 11, 8, 12, 9, 13, 10]
    for coco_id in coco_ids:
        if coco_id not in human.body_parts.keys():
            keypoints.extend([0, 0, 0])
            continue
        body_part = human.body_parts[coco_id]
        keypoints.extend(
            [round_int(body_part.x * image_w), round_int(body_part.y * image_h), 2]
        )
    return keypoints


if __name__ == "__main__":
    parser = argparse.ArgumentParser(description="Tensorflow Openpose Inference")
    parser.add_argument(
        "--resize",
        type=str,
        default="0x0",
        help="if provided, resize images before they are processed. default=0x0, Recommends : 432x368 or 656x368 or 1312x736 ",
    )
    parser.add_argument(
        "--resize-out-ratio",
        type=float,
        default=8.0,
        help="if provided, resize heatmaps before they are post-processed. default=8.0",
    )
    parser.add_argument(
        "--model",
        type=str,
        default="cmu",
        help="cmu / mobilenet_thin / mobilenet_v2_large",
    )
    parser.add_argument("--cocoyear", type=str, default="2014")
    parser.add_argument("--coco-dir", type=str, default="/data/public/rw/coco/")
    parser.add_argument("--data-idx", type=int, default=-1)
    parser.add_argument("--multi-scale", type=bool, default=False)
    args = parser.parse_args()

    cocoyear_list = ["2014", "2017"]
    if args.cocoyear not in cocoyear_list:
        logger.error("cocoyear should be one of %s" % str(cocoyear_list))
        sys.exit(-1)

    # TODO : Scales

    image_dir = args.coco_dir + "val%s/" % args.cocoyear
    coco_json_file = (
        args.coco_dir + "annotations/person_keypoints_val%s.json" % args.cocoyear
    )
    cocoGt = COCO(coco_json_file)
    catIds = cocoGt.getCatIds(catNms=["person"])
    keys = cocoGt.getImgIds(catIds=catIds)
    if args.data_idx < 0:
        if eval_size > 0:
            keys = keys[:eval_size]  # only use the first #eval_size elements.
        pass
    else:
        keys = [keys[args.data_idx]]
    logger.info("validation %s set size=%d" % (coco_json_file, len(keys)))
    write_json = "../etcs/%s_%s_%0.1f.json" % (
        args.model,
        args.resize,
        args.resize_out_ratio,
    )

    logger.debug("initialization %s : %s" % (args.model, get_graph_path(args.model)))
    w, h = model_wh(args.resize)
    if w == 0 or h == 0:
        e = TfPoseEstimator(get_graph_path(args.model), target_size=(432, 368))
    else:
        e = TfPoseEstimator(get_graph_path(args.model), target_size=(w, h))

    print("FLOPs: ", e.get_flops())

    result = []
    tqdm_keys = tqdm(keys)
    for i, k in enumerate(tqdm_keys):
        img_meta = cocoGt.loadImgs(k)[0]
        img_idx = img_meta["id"]

        img_name = os.path.join(image_dir, img_meta["file_name"])
        image = read_imgfile(img_name, None, None)
        if image is None:
            logger.error("image not found, path=%s" % img_name)
            sys.exit(-1)

        # inference the image with the specified network
        t = time.time()
        humans = e.inference(
            image,
            resize_to_default=(w > 0 and h > 0),
            upsample_size=args.resize_out_ratio,
        )
        elapsed = time.time() - t

        scores = 0
        ann_idx = cocoGt.getAnnIds(imgIds=[img_idx], catIds=[1])
        anns = cocoGt.loadAnns(ann_idx)
        for human in humans:
            item = {
                "image_id": img_idx,
                "category_id": 1,
                "keypoints": write_coco_json(
                    human, img_meta["width"], img_meta["height"]
                ),
                "score": human.score,
            }
            result.append(item)
            scores += item["score"]

        avg_score = scores / len(humans) if len(humans) > 0 else 0
        tqdm_keys.set_postfix(
            OrderedDict({"inference time": elapsed, "score": avg_score})
        )
        if args.data_idx >= 0:
            logger.info("score:", k, len(humans), len(anns), avg_score)

            import matplotlib.pyplot as plt

            fig = plt.figure()
            a = fig.add_subplot(2, 3, 1)
            plt.imshow(e.draw_humans(image, humans, True))

            a = fig.add_subplot(2, 3, 2)
            # plt.imshow(cv2.resize(image, (e.heatMat.shape[1], e.heatMat.shape[0])), alpha=0.5)
            tmp = np.amax(e.heatMat[:, :, :-1], axis=2)
            plt.imshow(tmp, cmap=plt.cm.gray, alpha=0.5)
            plt.colorbar()

            tmp2 = e.pafMat.transpose((2, 0, 1))
            tmp2_odd = np.amax(np.absolute(tmp2[::2, :, :]), axis=0)
            tmp2_even = np.amax(np.absolute(tmp2[1::2, :, :]), axis=0)

            a = fig.add_subplot(2, 3, 4)
            a.set_title("Vectormap-x")
            # plt.imshow(CocoPose.get_bgimg(inp, target_size=(vectmap.shape[1], vectmap.shape[0])), alpha=0.5)
            plt.imshow(tmp2_odd, cmap=plt.cm.gray, alpha=0.5)
            plt.colorbar()

            a = fig.add_subplot(2, 3, 5)
            a.set_title("Vectormap-y")
            # plt.imshow(CocoPose.get_bgimg(inp, target_size=(vectmap.shape[1], vectmap.shape[0])), alpha=0.5)
            plt.imshow(tmp2_even, cmap=plt.cm.gray, alpha=0.5)
            plt.colorbar()

            plt.show()

    fp = open(write_json, "w")
    json.dump(result, fp)
    fp.close()

    cocoDt = cocoGt.loadRes(write_json)
    cocoEval = COCOeval(cocoGt, cocoDt, "keypoints")
    cocoEval.params.imgIds = keys
    cocoEval.evaluate()
    cocoEval.accumulate()
    cocoEval.summarize()

    print("".join(["%11.4f |" % x for x in cocoEval.stats]))

    pred = json.load(open(write_json, "r"))
