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

from collections import defaultdict
import cv2
import os
import sys
import numpy as np
from caffe2.python import workspace
from detectron.core.config import assert_and_infer_cfg
from detectron.core.config import cfg
from detectron.core.config import merge_cfg_from_file
from detectron.utils.io import cache_url
from detectron.utils.logging import setup_logging
from detectron.utils.timer import Timer
import detectron.core.test_engine as infer_engine
import detectron.datasets.dummy_datasets as dummy_datasets
import detectron.utils.c2 as c2_utils
import detectron.utils.vis as vis_utils

class TextureFactory(object):

    def __init__(self):
        self.model = self.setup()

    def setup(self):
        c2_utils.import_detectron_ops()
        cv2.ocl.setUseOpenCL(False)
        merge_cfg_from_file('/DensePose/configs/DensePose_ResNet101_FPN_s1x-e2e.yaml')
        cfg.NUM_GPUS = 1
        weights = cache_url('https://s3.amazonaws.com/densepose/DensePose_ResNet101_FPN_s1x-e2e.pkl', cfg.DOWNLOAD_CACHE)
        assert_and_infer_cfg(cache_urls=False)
        model = infer_engine.initialize_model_from_cfg(weights,1)
        dummy_coco_dataset = dummy_datasets.get_coco_dataset()
        return model

    def get_iuv(self, img):
        timers = defaultdict(Timer)
        with c2_utils.NamedCudaScope(0):
            cls_boxes, cls_segms, cls_keyps, cls_bodys = infer_engine.im_detect_all(self.model, img, None, timers=timers)
        iuv = self.make_iuv(img, cls_boxes, cls_segms, cls_keyps, cls_bodys)
        if iuv is None:
            raise Exception("Nothing found")
        return iuv

    def make_iuv(self, im, boxes, segms, keypoints, body_uv):
        if isinstance(boxes, list):
            boxes, segms, keypoints, classes = self.convert_from_cls_format(
                boxes, segms, keypoints)

        if boxes is None or boxes.shape[0] == 0 or max(boxes[:, 4]) < 0.9:
            return

        IUV_fields = body_uv[1]
    	All_Coords = np.zeros(im.shape)
    	inds = np.argsort(boxes[:,4])

    	for i, ind in enumerate(inds):
            entry = boxes[ind,:]

            if entry[4] > 0.65:
                entry=entry[0:4].astype(int)
                output = IUV_fields[ind]
		All_Coords_Old = All_Coords[ entry[1] : entry[1]+output.shape[1],entry[0]:entry[0]+output.shape[2],:]
	        All_Coords_Old[All_Coords_Old==0]=output.transpose([1,2,0])[All_Coords_Old==0]
		All_Coords[ entry[1] : entry[1]+output.shape[1],entry[0]:entry[0]+output.shape[2],:]= All_Coords_Old

	All_Coords[:,:,1:3] = 255. * All_Coords[:,:,1:3]
        All_Coords[All_Coords>255] = 255.
	All_Coords = All_Coords.astype(np.uint8)
        return All_Coords

    def convert_from_cls_format(self, cls_boxes, cls_segms, cls_keyps):
        """Convert from the class boxes/segms/keyps format generated by the testing
        code.
        """
        box_list = [b for b in cls_boxes if len(b) > 0]
        if len(box_list) > 0:
            boxes = np.concatenate(box_list)
        else:
            boxes = None
        if cls_segms is not None:
            segms = [s for slist in cls_segms for s in slist]
        else:
            segms = None
        if cls_keyps is not None:
            keyps = [k for klist in cls_keyps for k in klist]
        else:
            keyps = None
        classes = []
        for j in range(len(cls_boxes)):
            classes += [j] * len(cls_boxes[j])
        return boxes, segms, keyps, classes
