from __future__ import absolute_import, division, print_function

import pdb
import time

import cv2
import numpy as np
import torch
from models.model import create_model, load_model
from progress.bar import Bar
from utils.debugger import Debugger
from utils.image import get_affine_transform


class BaseDetector(object):
  def __init__(self, opt):
    if opt.gpus[0] >= 0:
      opt.device = torch.device('cuda')
    else:
      opt.device = torch.device('cpu')
    
    print('Creating model...')
    self.model = create_model(opt.arch, opt.heads, opt.head_conv, opt.inp_channel)
    self.model = load_model(self.model, opt.load_model)
    self.model = self.model.to(opt.device)
    self.model.eval()

    self.mean = np.array(opt.mean, dtype=np.float32).reshape(1, 1, opt.inp_channel)
    self.std = np.array(opt.std, dtype=np.float32).reshape(1, 1, opt.inp_channel)
    self.max_per_image = 100
    self.num_classes = opt.num_classes
    self.scales = opt.test_scales
    self.opt = opt
    self.pause = True

  def pre_process(self, image, scale, meta=None):
    # original shape: [600, 1116, 3]
    height, width = image.shape[0:2]
    new_height = int(height * scale)
    new_width  = int(width * scale)
    if self.opt.fix_res:
      # here they are the dataset resolution - (512, 512) in the multi_pose case
      inp_height, inp_width = self.opt.input_h, self.opt.input_w
      # c: the center coordinate of the original image
      c = np.array([new_width / 2., new_height / 2.], dtype=np.float32)
      # s = 1116 = max(original_size)
      s = max(height, width) * 1.0
    else:
      inp_height = (new_height | self.opt.pad) + 1
      inp_width = (new_width | self.opt.pad) + 1
      c = np.array([new_width // 2, new_height // 2], dtype=np.float32)
      s = np.array([inp_width, inp_height], dtype=np.float32)

    # crop to resize to the dataset default resolution,
    # finally normalization(range to 0-1) and standardization (0-mean 1-var)
    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)

    # normalize and standardize
    if self.opt.input_mod == "RGD":
      inp_image[:, :, :2]  = inp_image[:,:,:2] / 255.
    elif self.opt.input_mod == "RGBD":
      inp_image[:, :, :3] = inp_image[:, :, :3] / 255.
    elif self.opt.input_mod == "RGB":
      inp_image = inp_image / 255.
    elif self.opt.input_mod == "D":
      inp_image = inp_image[:, :, None]   # since cv2.resize will drop the one channel axes
    inp_image = ((inp_image - self.mean) / self.std).astype(np.float32)

    images = inp_image.transpose(2, 0, 1).reshape(1, self.opt.inp_channel, inp_height, inp_width)
    if self.opt.flip_test:
      images = np.concatenate((images, images[:, :, :, ::-1]), axis=0)
    images = torch.from_numpy(images)
    meta = {'c': c, 's': s, 
            'out_height': inp_height // self.opt.down_ratio, 
            'out_width': inp_width // self.opt.down_ratio}
    return images, meta

  def process(self, images, return_time=False):
    raise NotImplementedError

  def post_process(self, dets, meta, scale=1):
    raise NotImplementedError

  def merge_outputs(self, detections):
    raise NotImplementedError

  def debug(self, debugger, images, dets, output, scale=1):
    raise NotImplementedError

  def show_results(self, debugger, image, results, display=True):
   raise NotImplementedError
  
  def save_results(self, debugger, prefix):
    raise NotImplementedError

  def run(self, image_or_path_or_tensor, meta=None):
    # prepare
    load_time, pre_time, net_time, dec_time, post_time = 0, 0, 0, 0, 0
    merge_time, tot_time = 0, 0
    debugger = Debugger(dataset=self.opt.dataset, ipynb=(self.opt.debug==3),
                        theme=self.opt.debugger_theme, kpt_type=self.opt.kpt_type)
    start_time = time.time()
    # load the image data - Assuming the data from tensor(dataset?) is already preprocessed
    pre_processed = False
    if isinstance(image_or_path_or_tensor, np.ndarray):
      image = image_or_path_or_tensor
    elif type(image_or_path_or_tensor) == type (''): 
      image = cv2.imread(image_or_path_or_tensor)
    else:
      image = image_or_path_or_tensor['image'][0].numpy()
      pre_processed_images = image_or_path_or_tensor
      pre_processed = True
    
    loaded_time = time.time()
    load_time += (loaded_time - start_time)

    # start the multiscale processing 
    detections = []
    for scale in self.scales:
      scale_start_time = time.time()

      # prprocess
      if not pre_processed:
        images, meta = self.pre_process(image, scale, meta)
      else:
        # import pdb; pdb.set_trace()
        images = pre_processed_images['images'][scale][0]
        meta = pre_processed_images['meta'][scale]
        meta = {k: v.numpy()[0] for k, v in meta.items()}
        if "img_idx" in pre_processed_images['meta'].keys():
          meta["img_idx"] = pre_processed_images['meta']["img_idx"].numpy().astype(np.int32)[0]
      images = images.to(self.opt.device)
      torch.cuda.synchronize()
      pre_process_time = time.time()
      pre_time += pre_process_time - scale_start_time

      # debug image
      if self.opt.input_mod == "RGBD":
        image_debug = image[:, :, :3]
        images_debug = images[:, :3, :, :] 
      else:
        image_debug = image
        images_debug = images

      # process      
      output, dets, forward_time = self.process(images, return_time=True)

      torch.cuda.synchronize()
      net_time += forward_time - pre_process_time
      decode_time = time.time()
      dec_time += decode_time - forward_time
      
      if self.opt.debug >= 2:
        self.debug(debugger, images_debug, dets, output, scale)

      # post process 
      dets = self.post_process(dets, meta, scale)
      torch.cuda.synchronize()
      post_process_time = time.time()
      post_time += post_process_time - decode_time

      detections.append(dets)

    # merge multiscale detection to get the final results. 
    results = self.merge_outputs(detections)
    torch.cuda.synchronize()
    end_time = time.time()
    merge_time += end_time - post_process_time
    tot_time += end_time - start_time

    if self.opt.debug >= 1:
      self.show_results(debugger, image_debug, results, display=(self.opt.debug != 5))

    if self.opt.debug >= 4:
      self.save_results(debugger, prefix=meta["img_idx"])
    
    return {'results': results, 'tot': tot_time, 'load': load_time,
            'pre': pre_time, 'net': net_time, 'dec': dec_time,
            'post': post_time, 'merge': merge_time}