#include <iostream>
#include <fstream>
#include <sstream>

#include <assert.h>
#include <sys/time.h>
#include <unistd.h>

#include <ros/ros.h>
#include <image_transport/image_transport.h>
#include <cv_bridge/cv_bridge.h>

#include "dnn_detect/DetectedObject.h"
#include "dnn_detect/DetectedObjectArray.h"
#include "dnn_detect/Detect.h"

#include <opencv2/highgui.hpp>
#include <opencv2/dnn.hpp>
#include <opencv2/calib3d.hpp>

#include <list>
#include <string>
#include <boost/algorithm/string.hpp>
#include <boost/format.hpp>

#include <thread>
#include <mutex>
#include <condition_variable>

#include "objectDetection.hpp"

using namespace std;
using namespace cv;

string yoloBasePath = "/home/phil/Development/RosDev/opencv_ws/src/opencv_processing/dat/"; // relative paths not working
string yoloClassesFile = yoloBasePath + "coco.names";
string yoloModelConfiguration = yoloBasePath + "yolov3-tiny.cfg";
string yoloModelWeights = yoloBasePath + "yolov3-tiny.weights";
float confThreshold = 0.01f;
float nmsThreshold = 0.9f;

std::condition_variable cond;
std::mutex mutx;

class DnnNode {
  private:
    ros::Publisher results_pub;

    image_transport::ImageTransport it;
    image_transport::Subscriber img_sub;

    // if set, we publish the images that contain objects
    bool publish_images;

    int frame_num;
    float min_confidence;
    int im_size;
    int rotate_flag;
    float scale_factor;
    float mean_val;
    std::vector<std::string> class_names;

    image_transport::Publisher image_pub;

    cv::dnn::Net net;
    cv::Mat resized_image;
    cv::Mat rotated_image;

    bool single_shot;
    volatile bool triggered;
    volatile bool processed;

    dnn_detect::DetectedObjectArray results;

    ros::ServiceServer detect_srv;

    bool trigger_callback(dnn_detect::Detect::Request &req,
                          dnn_detect::Detect::Response &res);

    void image_callback(const sensor_msgs::ImageConstPtr &msg);

  public:
    DnnNode(ros::NodeHandle &nh);
};

bool DnnNode::  trigger_callback(dnn_detect::Detect::Request &req,
                                 dnn_detect::Detect::Response &res)
{
    ROS_INFO("Got service request");
    triggered = true;

    std::unique_lock<std::mutex> lock(mutx);

    while (!processed) {
      cond.wait(lock);
    }
    res.result = results;
    processed = false;
    return true;
}


void DnnNode::image_callback(const sensor_msgs::ImageConstPtr & msg)
{
    if (single_shot && !triggered) {
        return;
    }
    triggered = false;

    ROS_INFO("Got image %d", msg->header.seq);
    frame_num++;

    cv_bridge::CvImagePtr cv_ptr;

    try {
        cv_ptr = cv_bridge::toCvCopy(msg, sensor_msgs::image_encodings::BGR8);

        int w = cv_ptr->image.cols;
        int h = cv_ptr->image.rows;

        if (rotate_flag >= 0) {
          cv::rotate(cv_ptr->image, rotated_image, rotate_flag);
          rotated_image.copyTo(cv_ptr->image);
        }

        // load class names from file
        vector<string> classes;
        ifstream ifs(yoloClassesFile.c_str());
        string line;
        while (getline(ifs, line)) classes.push_back(line);

        cv::resize(cv_ptr->image, resized_image, cvSize(im_size, im_size));
        //cv::Mat img = cv_ptr->image;
        cv::Mat blob;
        vector<cv::Mat> netOutput;
        double scalefactor = 1/255.0;
        cv::Size size = cv::Size(320, 320); // 416/416, 320/320, 608/608
        cv::Scalar mean = cv::Scalar(0,0,0);
        bool swapRB = false;
        bool crop = false;
        cv::dnn::blobFromImage(resized_image, blob, scalefactor, size, mean, swapRB, crop);

        // Get names of output layers
        vector<cv::String> names;
        vector<int> outLayers = net.getUnconnectedOutLayers(); // get  indices of  output layers, i.e.  layers with unconnected outputs
        vector<cv::String> layersNames = net.getLayerNames(); // get  names of all layers in the network

        names.resize(outLayers.size());
        for (size_t i = 0; i < outLayers.size(); ++i) // Get the names of the output layers in names
            names[i] = layersNames[outLayers[i] - 1];

        net.setInput(blob);
        net.forward(netOutput, names);

        // Scan through all bounding boxes and keep only the ones with high confidence
        vector<int> classIds; vector<float> confidences; vector<cv::Rect> boxes;
        for (size_t i = 0; i < netOutput.size(); ++i)
        {
            float* data = (float*)netOutput[i].data;
            for (int j = 0; j < netOutput[i].rows; ++j, data += netOutput[i].cols)
            {
                cv::Mat scores = netOutput[i].row(j).colRange(5, netOutput[i].cols);
                cv::Point classId;
                double confidence;

                // Get the value and location of the maximum score
                cv::minMaxLoc(scores, 0, &confidence, 0, &classId);
                if (confidence > confThreshold)
                {
                    cv::Rect box; int cx, cy;
                    cx = (int)(data[0] * img.cols);
                    cy = (int)(data[1] * img.rows);
                    box.width = (int)(data[2] * img.cols);
                    box.height = (int)(data[3] * img.rows);
                    box.x = cx - box.width/2; // left
                    box.y = cy - box.height/2; // top

                    boxes.push_back(box);
                    classIds.push_back(classId.x);
                    confidences.push_back((float)confidence);
                }
            }
        }


        std::unique_lock<std::mutex> lock(mutx);
        results.header.frame_id = msg->header.frame_id;
        results.objects.clear();

        // perform non-maxima suppression
        vector<int> indices;
        cv::dnn::NMSBoxes(boxes, confidences, confThreshold, nmsThreshold, indices);

        for(auto it=indices.begin(); it!=indices.end(); ++it) {

            std::string label = str(boost::format{"%1% %2%"} % classIds[*it] % confidences[*it]);
            ROS_INFO("%s", label.c_str());

            dnn_detect::DetectedObject obj;
            int x_min = boxes[*it].x;
            int x_max = boxes[*it].x + boxes[*it].width;
            int y_min = boxes[*it].y;
            int y_max = boxes[*it].y + boxes[*it].height;
            obj.x_min = x_min;
            obj.x_max = x_max;
            obj.y_min = y_min;
            obj.y_max = y_max;
            obj.class_name = classIds[*it];
            obj.confidence = confidences[*it];
            results.objects.push_back(obj);

            Rect object(x_min, y_min, x_max-x_min, y_max-y_min);

            rectangle(cv_ptr->image, object, Scalar(0, 255, 0));
            int baseline=0;
            cv::Size text_size = cv::getTextSize(label,
                                 FONT_HERSHEY_SIMPLEX, 0.75, 2, &baseline);
            putText(cv_ptr->image, label, Point(x_min, y_min-text_size.height),
                    FONT_HERSHEY_SIMPLEX, 0.75, Scalar(0, 255, 0));
        }


    results_pub.publish(results);

    image_pub.publish(cv_ptr->toImageMsg());

    }
    catch(cv_bridge::Exception & e) {
        ROS_ERROR("cv_bridge exception: %s", e.what());
    }
    catch(cv::Exception & e) {
        ROS_ERROR("cv exception: %s", e.what());
    }
    ROS_DEBUG("Notifying condition variable");
    processed = true;
    cond.notify_all();
}

DnnNode::DnnNode(ros::NodeHandle & nh) : it(nh)
{
    frame_num = 0;

    std::string dir;
    std::string proto_net_file;
    std::string caffe_model_file;
    std::string classes("background,"
       "aeroplane,bicycle,bird,boat,bottle,bus,car,cat,chair,"
       "cow,diningtable,dog,horse,motorbike,person,pottedplant,"
       "sheep,sofa,train,tvmonitor");



    nh.param<bool>("single_shot", single_shot, false);
    nh.param<bool>("publish_images", publish_images, false);
    nh.param<string>("data_dir", dir, "");
    nh.param<string>("protonet_file", proto_net_file, "MobileNetSSD_deploy.prototxt.txt");
    nh.param<string>("caffe_model_file", caffe_model_file, "MobileNetSSD_deploy.caffemodel");
    nh.param<float>("min_confidence", min_confidence, 0.01);
    nh.param<int>("im_size", im_size, 320);
    nh.param<int>("rotate_flag", rotate_flag, -1);
    nh.param<float>("scale_factor", scale_factor, 1/255.0);
    nh.param<float>("mean_val", mean_val, 127.5f);
    nh.param<std::string>("class_names", classes, classes);

    boost::split(class_names, classes, boost::is_any_of(","));
    ROS_INFO("Read %d class names", (int)class_names.size());

    try {
        net = cv::dnn::readNetFromDarknet(yoloModelConfiguration, yoloModelWeights);
        net.setPreferableBackend(cv::dnn::DNN_BACKEND_CUDA);
        net.setPreferableTarget(cv::dnn::DNN_TARGET_CUDA_FP16);
    }
    catch(cv::Exception & e) {
        ROS_ERROR("cv exception: %s", e.what());
        exit(1);
    }

    triggered = false;

    detect_srv = nh.advertiseService("detect", &DnnNode::trigger_callback, this);

    results_pub =
        nh.advertise<dnn_detect::DetectedObjectArray>("/dnn_objects", 20);

    image_pub = it.advertise("/dnn_images", 1);

    img_sub = it.subscribe("/camera", 1,
                           &DnnNode::image_callback, this);

    ROS_INFO("DNN detection ready");
}

int main(int argc, char ** argv) {
    ros::init(argc, argv, "dnn_detect");
    ros::NodeHandle nh("~");

    DnnNode node = DnnNode(nh);
    ros::MultiThreadedSpinner spinner(2);
    spinner.spin();

    return 0;
}