#!/usr/bin/env python
# coding: utf-8
#
# Author:   Kazuto Nakashima
# URL:      https://github.com/kazuto1011
# Created:  2016-07-05

import os
import os.path as osp
import sys

if 'FRCN_ROOT' not in os.environ:
    print "Could not find $FRCN_ROOT."
    sys.exit(1)

FRCN_ROOT = os.environ['FRCN_ROOT']
print "FRCNN=" + FRCN_ROOT
sys.path.insert(0, osp.join(FRCN_ROOT, 'caffe-fast-rcnn/python'))
sys.path.insert(0, osp.join(FRCN_ROOT, 'lib'))

from fast_rcnn.config import cfg
from fast_rcnn.test import im_detect
from fast_rcnn.nms_wrapper import nms
from utils.timer import Timer
import numpy as np
import caffe
import cv2
import argparse
import rospy as rp
import tms_ss_rcnn.srv as srv
import tms_ss_rcnn.msg as msg
import pymongo

db_client = pymongo.MongoClient('localhost', 27017)
db = db_client.faster_rcnn_experiment

CLASSES = ('__background__',
           'Water bottle',
           'Wine bottle',
           'Book',
           'Teapot',
           'Watering pot',
           'Kettle',
           'Coffee can',
           'Glass',
           'Cup',
           'Dish',
           'Controller'
           )

NETS = {'vgg16': ('VGG16', 'vgg16_housewares.caffemodel'),
        'zf': ('ZF', 'zf_housewares.caffemodel')}

def parse_args():
    parser = argparse.ArgumentParser(description='Faster R-CNN demo')
    parser.add_argument('--gpu', dest='gpu_id',
                        help='GPU device id to use [0]',
                        default=0, type=int)
    parser.add_argument('--cpu', dest='cpu_mode',
                        help='Use CPU mode (overrides --gpu)',
                        action='store_true')
    parser.add_argument('--net', dest='demo_net',
                        help='Network to use [vgg16]',
                        choices=NETS.keys(), default='zf')
    parser.add_argument('--conf', dest='conf_thresh',
                        default=0.8, type=float)
    parser.add_argument('--nms', dest='nms_thresh',
                        default=0.3, type=float)
    return parser.parse_args()


def load_net(args):
    cfg.TEST.HAS_RPN = True  # Use RPN for proposals

    data_dir = osp.join(osp.abspath(osp.dirname(__file__)), '../models')
    models_dir = osp.join(data_dir, 'houseware')
    prototxt = osp.join(models_dir, NETS[args.demo_net][
                        0], 'faster_rcnn_end2end', 'test.prototxt')
    caffemodel = osp.join(data_dir, NETS[args.demo_net][1])

    if args.cpu_mode:
        caffe.set_mode_cpu()
    else:
        caffe.set_mode_gpu()
        caffe.set_device(args.gpu_id)
        cfg.GPU_ID = args.gpu_id

    return caffe.Net(prototxt, caffemodel, caffe.TEST)


class FasterRCNN:

    def __init__(self, name, args, conf_thresh, nms_thresh):
        self._args = args
        self._name = name
        self._init = False
        self.conf_thresh = conf_thresh
        self.nms_thresh = nms_thresh

        rp.loginfo("Ready to start")
        self._server = rp.Service(
            self._name, srv.obj_detection, self._callback)

    def _callback(self, req):
        rp.loginfo("Received an image")

        if self._init is False:
            self._net = load_net(self._args)
            self._init = True

        if self._args.cpu_mode: