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 PASCAL_COCO_tensorboard_added(data.Dataset):
  num_classes = 20
  default_resolution = [384, 384]
  mean = np.array([0.485, 0.456, 0.406],
                   dtype=np.float32).reshape(1, 1, 3)
  std  = np.array([0.229, 0.224, 0.225],
                   dtype=np.float32).reshape(1, 1, 3)

  def __init__(self, opt, split):
    super(PASCAL_COCO_tensorboard_added, self).__init__()
    self.data_dir = os.path.join(opt.data_dir, 'voc_coco')
    self.img_dir = os.path.join(self.data_dir, 'images', 'JPEGImages')
    _ann_name = {'train': 'trainval2012', 'val': 'test2007'}
    self.annot_path = os.path.join(
      self.data_dir, 'annotations', 
      'pascal_{}.json').format(_ann_name[split])
    self.max_objs = 50
    # self.class_name = [
    #   '__background__', 'person', 'bicycle', 'car', 'motorcycle', 'airplane',
    #   'bus', 'train', 'truck', 'boat', 'traffic light', 'fire hydrant',
    #   'stop sign', 'parking meter', 'bench', 'bird', 'cat', 'dog', 'horse',
    #   'sheep', 'cow', 'elephant', 'bear', 'zebra', 'giraffe', 'backpack',
    #   'umbrella', 'handbag', 'tie', 'suitcase', 'frisbee', 'skis',
    #   'snowboard', 'sports ball', 'kite', 'baseball bat', 'baseball glove',
    #   'skateboard', 'surfboard', 'tennis racket', 'bottle', 'wine glass',
    #   'cup', 'fork', 'knife', 'spoon', 'bowl', 'banana', 'apple', 'sandwich',
    #   'orange', 'broccoli', 'carrot', 'hot dog', 'pizza', 'donut', 'cake',
    #   'chair', 'couch', 'potted plant', 'bed', 'dining table', 'toilet', 'tv',
    #   'laptop', 'mouse', 'remote', 'keyboard', 'cell phone', 'microwave',
    #   'oven', 'toaster', 'sink', 'refrigerator', 'book', 'clock', 'vase',
    #   'scissors', 'teddy bear', 'hair drier', 'toothbrush']
    self.class_name = ['__background__', "aeroplane", "bicycle", "bird", "boat",
     "bottle", "bus", "car", "cat", "chair", "cow", "diningtable", "dog", 
     "horse", "motorbike", "person", "pottedplant", "sheep", "sofa", 
     "train", "tvmonitor"]
    self._valid_ids = np.arange(1, 21, dtype=np.int32)
    self.cat_ids = {v: i for i, v in enumerate(self._valid_ids)}
    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.mean = np.array([0.485, 0.456, 0.406], np.float32).reshape(1, 1, 3)
    # self.std = np.array([0.229, 0.224, 0.225], np.float32).reshape(1, 1, 3)

    self.split = split
    self.opt = opt

    # print('==> initializing coco 2017 {} data.'.format(split))
    print('==> initializing PASCAL VOC2012 {} data.'.format(split))
    self.coco = coco.COCO(self.annot_path)
    self.images = sorted(self.coco.getImgIds())
    #self.images = os.listdir('/home/coffeemix/Desktop/Gordon/CenterNet-1/data/coco/VOC2012_train')
    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 = self._valid_ids[cls_ind - 1]
        for bbox in all_bboxes[image_id][cls_ind]:
          bbox[2] -= bbox[0]
          bbox[3] -= bbox[1]
          score = bbox[4]
          bbox_out  = list(map(self._to_float, bbox[0:4]))

          detection = {
              "image_id": int(image_id),
              "category_id": int(category_id),
              "bbox": bbox_out,
              "score": float("{:.2f}".format(score))
          }
          if len(bbox) > 5:
              extreme_points = list(map(self._to_float, bbox[5:13]))
              detection["extreme_points"] = extreme_points
          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(save_dir, "results.json")
    detections  = self.convert_eval_format(results)
    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, "bbox")
    coco_eval.evaluate()
    coco_eval.accumulate()
    coco_eval.summarize()
    return coco_eval.stats[:12]


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 PASCAL_COCO_tensorboard_added(data.Dataset):
  num_classes = 20
  default_resolution = [384, 384]
  mean = np.array([0.485, 0.456, 0.406],
                   dtype=np.float32).reshape(1, 1, 3)
  std  = np.array([0.229, 0.224, 0.225],
                   dtype=np.float32).reshape(1, 1, 3)

  def __init__(self, opt, split):
    super(PASCAL_COCO_tensorboard_added, self).__init__()
    self.data_dir = os.path.join(opt.data_dir, 'voc_coco')
    self.img_dir = os.path.join(self.data_dir, 'images', 'JPEGImages')
    _ann_name = {'train': 'trainval2012', 'val': 'test2007'}
    self.annot_path = os.path.join(
      self.data_dir, 'annotations', 
      'pascal_{}.json').format(_ann_name[split])
    self.max_objs = 50
    # self.class_name = [
    #   '__background__', 'person', 'bicycle', 'car', 'motorcycle', 'airplane',
    #   'bus', 'train', 'truck', 'boat', 'traffic light', 'fire hydrant',
    #   'stop sign', 'parking meter', 'bench', 'bird', 'cat', 'dog', 'horse',
    #   'sheep', 'cow', 'elephant', 'bear', 'zebra', 'giraffe', 'backpack',
    #   'umbrella', 'handbag', 'tie', 'suitcase', 'frisbee', 'skis',
    #   'snowboard', 'sports ball', 'kite', 'baseball bat', 'baseball glove',
    #   'skateboard', 'surfboard', 'tennis racket', 'bottle', 'wine glass',
    #   'cup', 'fork', 'knife', 'spoon', 'bowl', 'banana', 'apple', 'sandwich',
    #   'orange', 'broccoli', 'carrot', 'hot dog', 'pizza', 'donut', 'cake',
    #   'chair', 'couch', 'potted plant', 'bed', 'dining table', 'toilet', 'tv',
    #   'laptop', 'mouse', 'remote', 'keyboard', 'cell phone', 'microwave',
    #   'oven', 'toaster', 'sink', 'refrigerator', 'book', 'clock', 'vase',
    #   'scissors', 'teddy bear', 'hair drier', 'toothbrush']
    self.class_name = ['__background__', "aeroplane", "bicycle", "bird", "boat",
     "bottle", "bus", "car", "cat", "chair", "cow", "diningtable", "dog", 
     "horse", "motorbike", "person", "pottedplant", "sheep", "sofa", 
     "train", "tvmonitor"]
    self._valid_ids = np.arange(1, 21, dtype=np.int32)
    self.cat_ids = {v: i for i, v in enumerate(self._valid_ids)}
    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.mean = np.array([0.485, 0.456, 0.406], np.float32).reshape(1, 1, 3)
    # self.std = np.array([0.229, 0.224, 0.225], np.float32).reshape(1, 1, 3)

    self.split = split
    self.opt = opt

    # print('==> initializing coco 2017 {} data.'.format(split))
    print('==> initializing PASCAL VOC2012 {} data.'.format(split))
    self.coco = coco.COCO(self.annot_path)
    self.images = sorted(self.coco.getImgIds())
    #self.images = os.listdir('/home/coffeemix/Desktop/Gordon/CenterNet-1/data/coco/VOC2012_train')
    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 = self._valid_ids[cls_ind - 1]
        for bbox in all_bboxes[image_id][cls_ind]:
          bbox[2] -= bbox[0]
          bbox[3] -= bbox[1]
          score = bbox[4]
          bbox_out  = list(map(self._to_float, bbox[0:4]))

          detection = {
              "image_id": int(image_id),
              "category_id": int(category_id),
              "bbox": bbox_out,
              "score": float("{:.2f}".format(score))
          }
          if len(bbox) > 5:
              extreme_points = list(map(self._to_float, bbox[5:13]))
              detection["extreme_points"] = extreme_points
          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(save_dir, "results.json")
    detections  = self.convert_eval_format(results)
    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, "bbox")
    coco_eval.evaluate()
    coco_eval.accumulate()
    coco_eval.summarize()
    return coco_eval.stats[:12]


# self.class_name = [ #   '__background__', 'person', 'bicycle', 'car', 'motorcycle', 'airplane', #   'bus', 'train', 'truck', 'boat', 'traffic light', 'fire hydrant', #   'stop sign', 'parking meter', 'bench', 'bird', 'cat', 'dog', 'horse', #   'sheep', 'cow', 'elephant', 'bear', 'zebra', 'giraffe', 'backpack', #   'umbrella', 'handbag', 'tie', 'suitcase', 'frisbee', 'skis', #   'snowboard', 'sports ball', 'kite', 'baseball bat', 'baseball glove', #   'skateboard', 'surfboard', 'tennis racket', 'bottle', 'wine glass', #   'cup', 'fork', 'knife', 'spoon', 'bowl', 'banana', 'apple', 'sandwich', #   'orange', 'broccoli', 'carrot', 'hot dog', 'pizza', 'donut', 'cake', #   'chair', 'couch', 'potted plant', 'bed', 'dining table', 'toilet', 'tv', #   'laptop', 'mouse', 'remote', 'keyboard', 'cell phone', 'microwave', #   'oven', 'toaster', 'sink', 'refrigerator', 'book', 'clock', 'vase', #   'scissors', 'teddy bear', 'hair drier', 'toothbrush'] # self.mean = np.array([0.485, 0.456, 0.406], np.float32).reshape(1, 1, 3) # self.std = np.array([0.229, 0.224, 0.225], np.float32).reshape(1, 1, 3) # print('==> initializing coco 2017 {} data.'.format(split)) #self.images = os.listdir('/home/coffeemix/Desktop/Gordon/CenterNet-1/data/coco/VOC2012_train') # import pdb; pdb.set_trace()