#!/usr/bin/env python3
# -*- coding: utf-8 -*-

# Licensed to the Apache Software Foundation (ASF) under one
# or more contributor license agreements.  See the NOTICE file
# distributed with this work for additional information
# regarding copyright ownership.  The ASF licenses this file
# to you 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.

import argparse
import tools.find_mxnet
import mxnet as mx
import os
import sys
from detect.detector import Detector
from symbol.symbol_factory import get_symbol
from dataset.cv2Iterator import CameraIterator
import logging
import cv2

def get_detector(net, prefix, epoch, data_shape, mean_pixels, ctx, num_class,
                 nms_thresh=0.5, force_nms=True, nms_topk=400):
    """
    wrapper for initialize a detector

    Parameters:
    ----------
    net : str
        test network name
    prefix : str
        load model prefix
    epoch : int
        load model epoch
    data_shape : int
        resize image shape
    mean_pixels : tuple (float, float, float)
        mean pixel values (R, G, B)
    ctx : mx.ctx
        running context, mx.cpu() or mx.gpu(?)
    num_class : int
        number of classes
    nms_thresh : float
        non-maximum suppression threshold
    force_nms : bool
        force suppress different categories
    """
    if net is not None:
        if isinstance(data_shape, tuple):
            data_shape = data_shape[0]
        net = get_symbol(net, data_shape, num_classes=num_class, nms_thresh=nms_thresh,
            force_nms=force_nms, nms_topk=nms_topk)
    detector = Detector(net, prefix, epoch, data_shape, mean_pixels, ctx=ctx)
    return detector

def parse_args():
    parser = argparse.ArgumentParser(description='Single-shot detection network demo')
    parser.add_argument('--network', dest='network', type=str, default='resnet50',
                        help='which network to use')
    parser.add_argument('--images', dest='images', type=str, default='./data/demo/dog.jpg',
                        help='run demo with images, use comma to seperate multiple images')
    parser.add_argument('--dir', dest='dir', nargs='?',
                        help='demo image directory, optional', type=str)
    parser.add_argument('--ext', dest='extension', help='image extension, optional',
                        type=str, nargs='?')
    parser.add_argument('--epoch', dest='epoch', help='epoch of trained model',
                        default=0, type=int)
    parser.add_argument('--batch-size', dest='batch_size', help='batch size',
                        default=1, type=int)
    parser.add_argument('--prefix', dest='prefix', help='trained model prefix',
                        default=os.path.join(os.getcwd(), 'model', 'ssd_'),
                        type=str)
    parser.add_argument('--cpu', dest='cpu', help='(override GPU) use CPU to detect',
                        action='store_true', default=False)
    parser.add_argument('--gpu', dest='gpu_id', type=int, default=0,
                        help='GPU device id to detect with')
    parser.add_argument('--data-shape', dest='data_shape', type=str, default='512',
                        help='set image shape')
    parser.add_argument('--mean-r', dest='mean_r', type=float, default=123,
                        help='red mean value')
    parser.add_argument('--mean-g', dest='mean_g', type=float, default=117,
                        help='green mean value')
    parser.add_argument('--mean-b', dest='mean_b', type=float, default=104,
                        help='blue mean value')
    parser.add_argument('--thresh', dest='thresh', type=float, default=0.5,
                        help='object visualize score threshold, default 0.6')
    parser.add_argument('--nms', dest='nms_thresh', type=float, default=0.5,
                        help='non-maximum suppression threshold, default 0.5')
    parser.add_argument('--no-force', dest='force_nms', action='store_false',
                        help='dont force non-maximum suppression on different class')
    parser.add_argument('--no-timer', dest='show_timer', action='store_false',
                        help='dont show detection time')
    parser.add_argument('--deploy', dest='deploy_net', action='store_true', default=False,
                        help='Load network from json file, rather than from symbol')
    parser.add_argument('--class-names', dest='class_names', type=str,
                        default='aeroplane, bicycle, bird, boat, bottle, bus, \
                        car, cat, chair, cow, diningtable, dog, horse, motorbike, \
                        person, pottedplant, sheep, sofa, train, tvmonitor',
                        help='string of comma separated names, or text filename')
    parser.add_argument('--camera', action='store_true',
                        help="use camera for image capturing")
    parser.add_argument('--frame-resize', type=str, default=None,
                        help="resize camera frame to x,y pixels or a float scaling factor")
    args = parser.parse_args()
    return args

def parse_class_names(class_names):
    """ parse # classes and class_names if applicable """
    if len(class_names) > 0:
        if os.path.isfile(class_names):
            # try to open it to read class names
            with open(class_names, 'r') as f:
                class_names = [l.strip() for l in f.readlines()]
        else:
            class_names = [c.strip() for c in class_names.split(',')]
        for name in class_names:
            assert len(name) > 0
    else:
        raise RuntimeError("No valid class_name provided...")
    return class_names

def parse_frame_resize(x):
    if not x:
        return x
    x = list(map(float, x.strip().split(',')))
    assert len(x) >= 1 and len(x) <= 2, "frame_resize should be a float scaling factor or a tuple of w,h pixels"
    if len(x) == 1:
        x = x[0]
    return x

def parse_data_shape(data_shape_str):
    """Parse string to tuple or int"""
    ds = data_shape_str.strip().split(',')
    if len(ds) == 1:
        data_shape = (int(ds[0]), int(ds[0]))
    elif len(ds) == 2:
        data_shape = (int(ds[0]), int(ds[1]))
    else:
        raise ValueError("Unexpected data_shape: %s", data_shape_str)
    return data_shape

def draw_detection(frame, det, class_names):
    (klass, score, x0, y0, x1, y1) = det
    klass_name = class_names[int(klass)]
    h = frame.shape[0]
    w = frame.shape[1]
    # denormalize detections from [0,1] to the frame size
    p0 = tuple(map(int, (x0*w,y0*h)))
    p1 = tuple(map(int, (x1*w,y1*h)))
    logging.info("detection: %s %s", klass_name, score)
    cv2.rectangle(frame, p0, p1, (0,0,255), 2)
    # Where to draw the text, a few pixels above the top y coordinate
    tp0 = (p0[0], p0[1]-5)
    draw_text = "{} {}".format(klass_name, score)
    cv2.putText(frame, draw_text, tp0, cv2.FONT_HERSHEY_COMPLEX_SMALL, 0.5, (0,0,255))


def network_path(prefix, network, data_shape):
    return "{}{}_{}".format(prefix, network, data_shape)

def run_camera(args,ctx):
    assert args.batch_size == 1, "only batch size of 1 is supported"
    logging.info("Detection threshold is {}".format(args.thresh))
    iter = CameraIterator(frame_resize=parse_frame_resize(args.frame_resize))
    class_names = parse_class_names(args.class_names)
    mean_pixels = (args.mean_r, args.mean_g, args.mean_b)
    data_shape = int(args.data_shape)
    batch_size = int(args.batch_size)
    detector = Detector(
        get_symbol(args.network, data_shape, num_classes=len(class_names)),
        network_path(args.prefix, args.network, data_shape),
        args.epoch,
        data_shape,
        mean_pixels,
        batch_size,
        ctx
    )
    for frame in iter:
        logging.info("Frame info: shape %s type %s", frame.shape, frame.dtype)
        logging.info("Generating batch")
        data_batch = detector.create_batch(frame)
        logging.info("Detecting objects")
        detections_batch = detector.detect_batch(data_batch)
        #detections = [mx.nd.array((1,1,0.2,0.2,0.4,0.4))]
        detections = detections_batch[0]
        logging.info("%d detections", len(detections))
        for det in detections:
            obj = det.asnumpy()
            (klass, score, x0, y0, x1, y1) = obj
            if score > args.thresh:
                draw_detection(frame, obj, class_names)
        cv2.imshow('frame', frame)

def run_images(args,ctx):
    # parse image list
    image_list = [i.strip() for i in args.images.split(',')]
    assert len(image_list) > 0, "No valid image specified to detect"

    network = None if args.deploy_net else args.network
    class_names = parse_class_names(args.class_names)
    data_shape = parse_data_shape(args.data_shape)
    if args.prefix.endswith('_'):
        prefix = args.prefix + args.network + '_' + str(data_shape[0])
    else:
        prefix = args.prefix
    detector = get_detector(network, prefix, args.epoch,
                            data_shape,
                            (args.mean_r, args.mean_g, args.mean_b),
                            ctx, len(class_names), args.nms_thresh, args.force_nms)
    # run detection
    detector.detect_and_visualize(image_list, args.dir, args.extension,
                                  class_names, args.thresh, args.show_timer)

def main():
    logging.getLogger().setLevel(logging.INFO)
    logging.basicConfig(format='%(asctime)-15s %(message)s')
    args = parse_args()
    if args.cpu:
        ctx = mx.cpu()
    else:
        ctx = mx.gpu(args.gpu_id)

    if args.camera:
        run_camera(args, ctx)
    else:
        run_images(args, ctx)
    return 0

if __name__ == '__main__':
    sys.exit(main())