#!/usr/bin/env python

# Copyright (c) 2017-present, Facebook, Inc.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
##############################################################################

"""Perform inference on a single image or all images with a certain extension
(e.g., .jpg) in a folder.
"""

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

from collections import defaultdict
import argparse
import cv2  # NOQA (Must import before importing caffe2 due to bug in cv2)
import glob
import logging
import os
import sys
import time

from caffe2.python import workspace

from core.config import assert_and_infer_cfg
from core.config import cfg
from core.config import merge_cfg_from_file
from utils.io import cache_url
from utils.timer import Timer
import core.test_engine as infer_engine
import datasets.dummy_datasets as dummy_datasets
from datasets.json_dataset import JsonDataset
import utils.c2 as c2_utils
import utils.logging
import utils.vis as vis_utils
import os.path as osp
c2_utils.import_detectron_ops()
# OpenCL may be enabled by default in OpenCV3; disable it because it's not
# thread safe and causes unwanted GPU memory allocations.
cv2.ocl.setUseOpenCL(False)


def parse_args():
    parser = argparse.ArgumentParser(description='End-to-end inference')
    parser.add_argument(
        '--cfg',
        dest='cfg',
        help='cfg model file (/path/to/model_config.yaml)',
        default=None,
        type=str
    )
    parser.add_argument(
        '--wts',
        dest='weights',
        help='weights model file (/path/to/model_weights.pkl)',
        default=None,
        type=str
    )
    parser.add_argument(
        '--output-dir',
        dest='output_dir',
        help='directory for visualization pdfs (default: /tmp/infer_simple)',
        default='/tmp/infer_simple',
        type=str
    )
    parser.add_argument(
        '--image-ext',
        dest='image_ext',
        help='image file name extension (default: jpg)',
        default='jpg',
        type=str
    )
    parser.add_argument(
        '--im_or_folder', dest='im_or_folder', help='image or folder of images', default=None
    )
    parser.add_argument(
        '--cls_thrsh_file', dest='cls_thrsh_file', help='image or folder of images', default=None
        '--always-out',
        dest='out_when_no_box',
        help='output image even when no object is found',
        action='store_true'
    )
    parser.add_argument(
        '--output-ext',
        dest='output_ext',
        help='output image file format (default: pdf)',
        default='pdf',
        type=str
    )
    parser.add_argument(
        '--thresh',
        dest='thresh',
        help='Threshold for visualizing detections',
        default=0.7,
        type=float
    )
    parser.add_argument(
        '--kp-thresh',
        dest='kp_thresh',
        help='Threshold for visualizing keypoints',
        default=2.0,
        type=float
    )
    parser.add_argument(
        'im_or_folder', help='image or folder of images', default=None
    )
    if len(sys.argv) == 1:
        parser.print_help()
        sys.exit(1)
    return parser.parse_args()
    
import numpy as np
def convert_from_cls_format(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
    
def output_detbbox_one_image(
        im, im_name, output_dir, boxes, segms=None, keypoints=None, class_thresholds=None, thresh=0.5, 
        FilterPrecScore=True, FilterSize=True,
        kp_thresh=2, dpi=200, box_alpha=0.0, dataset=None, show_class=True,
        ext='tsv'):
    """Visual debugging of detections."""
    scoreTopK = 5
    if not os.path.exists(output_dir):
        os.makedirs(output_dir)

    if isinstance(boxes, list):
        boxes, segms, keypoints, classes = convert_from_cls_format(
            boxes, segms, keypoints)

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

    # Display in largest to smallest order to reduce occlusion
    areas = (boxes[:, 2] - boxes[:, 0]) * (boxes[:, 3] - boxes[:, 1])
    out_scores = boxes[:,-1]

    #sorted_inds = np.argsort(-areas)
    sorted_inds = np.argsort(-out_scores)
    height, width, channels = im.shape
    size_ratio_thrsh = 0.15
    mask_color_id = 0
    urlPrefix = "http://cached.blob.core.windows.net/image-snapshots/"
    imgName = osp.splitext(osp.basename(im_name))[0]
    url = urlPrefix + imgName
    print (url)
    outStrings = ''
    
    for i in sorted_inds[:scoreTopK]:
        bbox = boxes[i, :4]
        score = boxes[i, -1]
        cls_name = dataset.classes[classes[i]]
        cls_thrsh = class_thresholds[cls_name]
        if score < thresh: #cls_thrsh:
            continue


        if score < cls_thrsh:
            if FilterPrecScore:
                continue
            ecolor = 'm'
            tcolor = 'm' # text color    
        
        box_width = bbox[2] - bbox[0]
        box_height = bbox[3] - bbox[1]
        lstyle = '-'
        ecolor = 'r'
        tcolor = 'g'
        if float(box_width)/float(width) < size_ratio_thrsh or float(box_height)/float(height) < size_ratio_thrsh:
            if FilterSize:
                continue
            #print ("bbox width ratio = {}, height ratio = {}".format(float(box_width)/float(width), float(box_height)/float(height)))

        #print("l,t,r,b,score = {}: ({:.3f}, {:.3f}, {:.3f}, {:.3f}, {:.3f})".format(cls_name, bbox[0], bbox[1], bbox[2], bbox[3], score))
        line = "{}\t{:.4f}\t{:.4f}\t{:.4f}\t{:.4f}\t{}\n".format(url, bbox[0]/float(width), bbox[1]/float(height), bbox[2]/float(width), bbox[3]/float(height), cls_name) # need to normalized
        print (line)
        outStrings += line
        # show box (off by default)
        
    return outStrings


def checkMkdir(dir):
    if not osp.isdir(dir):
        os.makedirs(dir)

def main(args):
    datasetName = 'furniture_val'
    logger = logging.getLogger(__name__)
    merge_cfg_from_file(args.cfg)
    cfg.NUM_GPUS = 1
    vis = True #False 
    args.weights = cache_url(args.weights, cfg.DOWNLOAD_CACHE)
    assert_and_infer_cfg(cache_urls=False)
    model = infer_engine.initialize_model_from_cfg(args.weights)
    if args.cls_thrsh_file is not None:
        class_thresholds = {l.split('\t')[0]:float(l.rstrip().split('\t')[1]) for l in open(args.cls_thrsh_file,'r').readlines()}
        print (class_thresholds)
    else:
        class_thresholds = None
    dummy_coco_dataset = dummy_datasets.get_coco_dataset()
    dataset = JsonDataset(datasetName)
    print (args.im_or_folder)
    if os.path.isdir(args.im_or_folder):
        im_list = glob.iglob(args.im_or_folder + '/*.' + args.image_ext)
    else:
        im_list = [args.im_or_folder]
    checkMkdir(args.output_dir)
    outTable = osp.join(args.output_dir, 'HF_CT_Measurement_Detected_Boxes.tsv')
    with open(outTable,'wb') as fout:
        for i, im_name in enumerate(im_list):
            out_name = os.path.join(
                args.output_dir, '{}'.format(os.path.basename(im_name) + '.pdf')
            )
            #logger.info('Processing {} -> {}'.format(im_name, out_name))
            im = cv2.imread(im_name)
            timers = defaultdict(Timer)
            t = time.time()
            with c2_utils.NamedCudaScope(0):
                cls_boxes, cls_segms, cls_keyps = infer_engine.im_detect_all(
                    model, im, None, timers=timers
                )
            #logger.info('Inference time: {:.3f}s'.format(time.time() - t))
            #for k, v in timers.items():
            #    logger.info(' | {}: {:.3f}s'.format(k, v.average_time))
            if i == 0:
                logger.info(
                    ' \ Note: inference on the first image will be slower than the '
                    'rest (caches and auto-tuning need to warm up)'
                )
            
            
            if vis:
                vis_utils.vis_one_image(
                    im[:, :, ::-1],  # BGR -> RGB for visualization
                    im_name,
                    args.output_dir,
                    cls_boxes,
                    cls_segms,
                    cls_keyps,
                    dataset=dummy_coco_dataset,
                    box_alpha=0.3,
                    show_class=True,
                    thresh=0.7,
                    kp_thresh=2, 
                    ext=args.output_ext,
                    out_when_no_box=args.out_when_no_box
                )



if __name__ == '__main__':
    workspace.GlobalInit(['caffe2', '--caffe2_log_level=0'])
    utils.logging.setup_logging(__name__)
    args = parse_args()
    main(args)
