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

import torch.utils.data as data
import numpy as np
import torch
import json
import cv2
import os
from utils.image import flip, color_aug
from utils.image import get_affine_transform, affine_transform
from utils.image import gaussian_radius, draw_umich_gaussian, draw_msra_gaussian
from utils.image import draw_dense_reg
import math

def pre_process(image, scale, mean, std, meta=None):
    height, width = image.shape[0:2]
    new_height = int(height * scale)
    new_width  = int(width * scale)
    inp_height = (new_height | 31) + 1
    inp_width = (new_width | 31) + 1
    c = np.array([new_width // 2, new_height // 2], dtype=np.float32)
    s = np.array([inp_width, inp_height], dtype=np.float32)

    trans_input = get_affine_transform(c, s, 0, [inp_width, inp_height])
    resized_image = cv2.resize(image, (new_width, new_height))
    inp_image = cv2.warpAffine(
      resized_image, trans_input, (inp_width, inp_height),
      flags=cv2.INTER_LINEAR)
    inp_image = ((inp_image / 255. - mean) / std).astype(np.float32)
    images = inp_image.transpose(2, 0, 1).reshape(1, 3, inp_height, inp_width)
  
    images = torch.from_numpy(images)
    meta = {'c': c, 's': s, 
            'out_height': inp_height // 4, 
            'out_width': inp_width // 4}
    return images, meta
reverse = lambda x : x[...,2]*100+x[...,1]*10+x[...,0]
x = np.expand_dims(np.array([x for x in range(512)]), 0).repeat(512, 0)
y = np.expand_dims(np.array([x for x in range(512)]), 1).repeat(512, 1)
class CTDetDataset(data.Dataset):
  def _coco_box_to_bbox(self, box):
    bbox = np.array([box[0], box[1], box[0] + box[2], box[1] + box[3]],
                    dtype=np.float32)
    return bbox

  def _get_border(self, border, size):
    i = 1
    while size - border // i <= border // i:
        i *= 2
    return border // i

  def __getitem__(self, index):
    img_id = self.images[index]
    file_name = self.coco.loadImgs(ids=[img_id])[0]['file_name']
    img_path = os.path.join(self.img_dir, file_name)
    gx_path = os.path.join(self.data_dir, 'panoptic', 'converted', 'gx', file_name.replace('jpg', 'png'))
    gy_path = os.path.join(self.data_dir, 'panoptic', 'converted', 'gy', file_name.replace('jpg', 'png'))
    ann_ids = self.coco.getAnnIds(imgIds=[img_id])
    anns = self.coco.loadAnns(ids=ann_ids)
    num_objs = min(len(anns), self.max_objs)

    img = cv2.imread(img_path)
    if not self.split == 'train':
      inp, meta = pre_process(img, 1, self.mean, self.std)
      meta['img_id']=img_id
      return inp, meta
    gx = cv2.imread(gx_path)
    gy = cv2.imread(gy_path)

    height, width = img.shape[0], img.shape[1]
    c = np.array([img.shape[1] / 2., img.shape[0] / 2.], dtype=np.float32)
    s = max(img.shape[0], img.shape[1]) * 1.0
    input_h, input_w = [512, 512]
    
    flipped = False
    if self.split == 'train':
      s = s * np.random.choice(np.arange(0.6, 1.4, 0.1))
      w_border = self._get_border(128, img.shape[1])
      h_border = self._get_border(128, img.shape[0])
      c[0] = np.random.randint(low=w_border, high=img.shape[1] - w_border)
      c[1] = np.random.randint(low=h_border, high=img.shape[0] - h_border)
      
      if np.random.random() < 0.5:
        flipped = True
        img = img[:, ::-1, :]
        gx = gx[:, ::-1]
        gy = gy[:, ::-1]
        c[0] =  width - c[0] - 1
        

    trans_input = get_affine_transform(
      c, s, 0, [input_w, input_h])
    inp = cv2.warpAffine(img, trans_input, 
                         (input_w, input_h),
                         flags=cv2.INTER_LINEAR)
    gx = cv2.warpAffine(gx, trans_input, 
                      (input_w, input_h),
                      flags=cv2.INTER_NEAREST)
    gy = cv2.warpAffine(gy, trans_input, 
                      (input_w, input_h),
                      flags=cv2.INTER_NEAREST)
    gx = reverse(gx.astype(np.float32))
    gy = reverse(gy.astype(np.float32))
    grad = np.stack([gx,gy, np.ones_like(gx)],-1)

    inp = (inp.astype(np.float32) / 255.)
    if self.split == 'train':
      color_aug(self._data_rng, inp, self._eig_val, self._eig_vec)
    inp = (inp - self.mean) / self.std
    inp = inp.transpose(2, 0, 1)

    output_h = input_h // 4
    output_w = input_w // 4
    num_classes = self.num_classes
    trans_output = get_affine_transform(c, s, 0, [output_w, output_h])

    t = get_affine_transform(c, s, 0, [input_w, input_h])
    if flipped:
        grad[...,0] = width - grad[...,0] - 1
    grad = np.einsum("ij,...j->...i", t, grad)
    grad = np.clip(grad, 0, 511)
    grad[...,0] -= x
    grad[...,1] -= y
    grad[np.where((gx+gy)==0)] = [0,0]

    hm = np.zeros((num_classes, output_h, output_w), dtype=np.float32)
    wh = np.zeros((self.max_objs, 2), dtype=np.float32)
    dense_wh = np.zeros((2, output_h, output_w), dtype=np.float32)
    reg = np.zeros((self.max_objs, 2), dtype=np.float32)
    ind = np.zeros((self.max_objs), dtype=np.int64)
    reg_mask = np.zeros((self.max_objs), dtype=np.uint8)
    cat_spec_wh = np.zeros((self.max_objs, num_classes * 2), dtype=np.float32)
    cat_spec_mask = np.zeros((self.max_objs, num_classes * 2), dtype=np.uint8)
    
    draw_gaussian = draw_umich_gaussian

    gt_det = []
    for k in range(num_objs):
      ann = anns[k]
      bbox = self._coco_box_to_bbox(ann['bbox'])
      cls_id = int(self.cat_ids[ann['category_id']])
      if flipped:
        bbox[[0, 2]] = width - bbox[[2, 0]] - 1
      bbox[:2] = affine_transform(bbox[:2], trans_output)
      bbox[2:] = affine_transform(bbox[2:], trans_output)
      bbox[[0, 2]] = np.clip(bbox[[0, 2]], 0, output_w - 1)
      bbox[[1, 3]] = np.clip(bbox[[1, 3]], 0, output_h - 1)
      h, w = bbox[3] - bbox[1], bbox[2] - bbox[0]
      if h > 0 and w > 0:
        radius = gaussian_radius((math.ceil(h), math.ceil(w)))
        radius = max(0, int(radius))
        ct = np.array(
          [(bbox[0] + bbox[2]) / 2, (bbox[1] + bbox[3]) / 2], dtype=np.float32)
        ct_int = ct.astype(np.int32)
        draw_gaussian(hm[cls_id], ct_int, radius)
        wh[k] = 1. * w, 1. * h
        ind[k] = ct_int[1] * output_w + ct_int[0]
        reg[k] = ct - ct_int
        reg_mask[k] = 1
        cat_spec_wh[k, cls_id * 2: cls_id * 2 + 2] = wh[k]
        cat_spec_mask[k, cls_id * 2: cls_id * 2 + 2] = 1

        gt_det.append([ct[0] - w / 2, ct[1] - h / 2, 
                       ct[0] + w / 2, ct[1] + h / 2, 1, cls_id])

    return inp, hm, reg_mask, ind, wh, reg, grad