from __future__ import absolute_import
from __future__ import division
from __future__ import print_function

import pycocotools.coco as coco
from pycocotools.cocoeval import COCOeval
import numpy as np
import json
import os

import torch.utils.data as data

class HRSCHP(data.Dataset):
  num_classes =33#33 #1#24 change
  num_joints = 1
  default_resolution =  [1024, 1024]#[800, 1333]#
  mean = np.array([0.40789654, 0.44719302, 0.47026115],
                   dtype=np.float32).reshape(1, 1, 3)
  std  = np.array([0.28863828, 0.27408164, 0.27809835],
                   dtype=np.float32).reshape(1, 1, 3)
  flip_idx =[[1, 1]] #[[1, 2], [3, 4], [5, 6], [7, 8], [9, 10], 
             # [11, 12], [13, 14], [15, 16]]
  def __init__(self, opt, split):
    super(HRSCHP, self).__init__()
    self.edges = [[0, 1], [0, 2], [1, 3], [2, 4], 
                  [4, 6], [3, 5], [5, 6], 
                  [5, 7], [7, 9], [6, 8], [8, 10], 
                  [6, 12], [5, 11], [11, 12], 
                  [12, 14], [14, 16], [11, 13], [13, 15]]
    
    self.acc_idxs =[1]# [1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16]
    self.data_dir = os.path.join(opt.data_dir, 'coco')
    self.img_dir = os.path.join(self.data_dir, '{}2017'.format(split))
    NAME_LABEL_MAP_HRSC_en = {
        'ship': 1,
        'aircraft carrier': 2,
        'warcraft': 3,
        'merchant ship': 4,
        'Nimitz': 5,
        'Enterprise': 6,
        'Arleigh Burke': 7,
        'WhidbeyIsland': 8,
        'Perry': 9,
        'Sanantonio': 10,
        'Ticonderoga': 11,
        'Kitty Hawk': 12,
        'Kuznetsov': 13,
        'Abukuma': 14,
        'Austen': 15,
        'Tarawa': 16,
        'Blue Ridge': 17,
        'Container ship': 18,
        'OXo|--)': 19,
        'Car carrier([]==[])': 20,
        'jian':21,
        'Hovercraft': 22,
        'hang':23,
        'yacht': 24,
        'CntShip(_|.--.--|_]=':25,
        'Cruise':26,
        'submarine':27,
        'lute':28,
        'Medical':29,
        'Car carrier(======|':30,
        'Ford-class':31,
        'Midway-class':32,
        'Invincible-class':33
        }
    self.class_name=[]
    for name, label in NAME_LABEL_MAP_HRSC_en.items():
        self.class_name.append(name)
    if split == 'test':
      self.annot_path = os.path.join(
          self.data_dir, 'annotations', 
          'image_info_test-dev2017.json').format(split)
    else:
      self.annot_path = os.path.join(
        self.data_dir, 'annotations', 
        'person_keypoints_{}2017.json').format(split)
    self.max_objs = 32
    self._data_rng = np.random.RandomState(123)
    self._eig_val = np.array([0.2141788, 0.01817699, 0.00341571],
                             dtype=np.float32)
    self._eig_vec = np.array([
        [-0.58752847, -0.69563484, 0.41340352],
        [-0.5832747, 0.00994535, -0.81221408],
        [-0.56089297, 0.71832671, 0.41158938]
    ], dtype=np.float32)
    self.split = split
    self.opt = opt

    print('==> initializing coco 2017 {} data.'.format(split))
    self.coco = coco.COCO(self.annot_path)
    image_ids = self.coco.getImgIds()

    if split == 'train':
      self.images = []
      for img_id in image_ids:
        idxs = self.coco.getAnnIds(imgIds=[img_id])
        if len(idxs) > 0:
          self.images.append(img_id)
    else:
      self.images = image_ids
    self.num_samples = len(self.images)
    print('Loaded {} {} samples'.format(split, self.num_samples))

  def _to_float(self, x):
    return float("{:.2f}".format(x))

  def convert_eval_format(self, all_bboxes):
    # import pdb; pdb.set_trace()
    detections = []
    for image_id in all_bboxes:
      for cls_ind in all_bboxes[image_id]:
        category_id = cls_ind#change 1
        for dets in all_bboxes[image_id][cls_ind]:
          bbox = dets[:4]
          bbox[2] -= bbox[0]
          bbox[3] -= bbox[1]
          score = dets[4]
          bbox_out  = list(map(self._to_float, bbox))
          keypoints = np.concatenate([
            #change
            np.array(dets[5:7], dtype=np.float32).reshape(-1, 2), #change39
            np.ones((1, 1), dtype=np.float32)], axis=1).reshape(3).tolist()#change 17 51
          keypoints  = list(map(self._to_float, keypoints))

          detection = {
              "image_id": int(image_id),
              "category_id": int(category_id),
              "bbox": bbox_out,
              "score": float("{:.2f}".format(score)),
              "keypoints": keypoints
          }
          detections.append(detection)
    return detections

  def __len__(self):
    return self.num_samples

  def save_results(self, results, save_dir):
    json.dump(self.convert_eval_format(results), 
              open('{}/results.json'.format(save_dir), 'w'))


  def run_eval(self, results, save_dir):
    # result_json = os.path.join(opt.save_dir, "results.json")
    # detections  = convert_eval_format(all_boxes)
    # json.dump(detections, open(result_json, "w"))
    self.save_results(results, save_dir)
    coco_dets = self.coco.loadRes('{}/results.json'.format(save_dir))
    coco_eval = COCOeval(self.coco, coco_dets, "keypoints")
    coco_eval.evaluate()
    coco_eval.accumulate()
    coco_eval.summarize()
    coco_eval = COCOeval(self.coco, coco_dets, "bbox")
    coco_eval.evaluate()
    coco_eval.accumulate()
    coco_eval.summarize()