#include "DnDetect.h"

DnDetect::DnDetect(const int gpuIndex):m_net(nullptr),m_thresh(0.5),m_hier_thresh(0.5),m_nms(0.45),names(nullptr),alphabet(nullptr)
{
    cuda_set_device(gpuIndex);

    char *datacfg = "cfg/coco.data";
    char *cfgfile = "cfg/yolov3.cfg";
    char *weightfile = "./yolov3.weights";

    double loadtime = what_time_is_it_now();
    list *options = read_data_cfg(datacfg);
    char *name_list = option_find_str(options, "names", "data/names.list");
    names = get_labels(name_list);

    alphabet = load_alphabet();
    m_net = load_network(cfgfile, weightfile, 0);
    set_batch_network(m_net, 1);
    printf("load mod use %f seconds.\n", what_time_is_it_now()-loadtime);

    srand(2222222);
    double attime;
    float nms=.45;
}

DnDetect::~DnDetect()
{

}

std::vector<cv::Rect2f> DnDetect::detect(cv::Mat & img)
{

    double bttime=what_time_is_it_now();
    image im = matToImg(img);
    image sized = letterbox_image(im, m_net->w, m_net->h);
    layer l = m_net->layers[m_net->n-1];

    float *X = sized.data;
    //attime=what_time_is_it_now();
    network_predict(m_net, X);
    //printf("Predicted in %f seconds.\n", what_time_is_it_now()-attime);
    int nboxes = 0;
    detection *dets = get_network_boxes(m_net, im.w, im.h, 0.5f, 0.1f, 0, 1, &nboxes);
    if (m_nms) do_nms_sort(dets, nboxes, l.classes, m_nms);
  //  draw_detections(im, dets, nboxes, m_thresh, names, alphabet, l.classes);
    std::vector<cv::Rect2f> rects;
    for(int i=0;i<nboxes;i++)
    {
        std::vector<float> vec(80);
        memcpy(&vec[0],dets[i].prob,sizeof(float)*80);
        int type = -1;
        for(int j = 0; j < l.classes; ++j){
            if(j != 0){
                continue;
            }
            if (dets[i].prob[j] > m_thresh){
                if (type < 0) {
                    type = j;
                } else {
                }
            } else{
            }
        }
        if(type >= 0){
            if(type != 0){
                continue;
            }
            float left  = (dets[i].bbox.x-dets[i].bbox.w/2.);
            float top   = (dets[i].bbox.y-dets[i].bbox.h/2.);
            cv::Rect2f rect(left,top,dets[i].bbox.w,dets[i].bbox.h);
            std::cout<<"rect.x : "<<left<<" rect.y : "<<top<<" rect.w : "<<rect.width<<" rect.h : "<<rect.height<<std::endl;
            rects.push_back(rect);
        }


    }
    free_detections(dets, nboxes);
//    show_image(im, "Video");
//    cv::waitKey(10);
    free_image(im);
    free_image(sized);
    printf("all time use %f seconds.\n", what_time_is_it_now()-bttime);
    return rects;
}

image DnDetect::matToImg(cv::Mat& RefImg) {
    CV_Assert(RefImg.depth() == CV_8U);

    int h = RefImg.rows;
    int w = RefImg.cols;
    int channels = RefImg.channels();
    image im = make_image(w, h, 3);
    int count = 0;
    switch(channels){
        case 1:{
            cv::MatIterator_<unsigned char> it, end;
            for (it = RefImg.begin<unsigned char>(), end = RefImg.end<unsigned char>(); it != end; ++it){
                im.data[count] = im.data[w*h + count] = im.data[w*h*2 + count] = (float)(*it)/255.0;

                ++count;
            }
            break;
        }
        case 3:{
            cv::MatIterator_<cv::Vec3b> it, end;
            for (it = RefImg.begin<cv::Vec3b>(), end = RefImg.end<cv::Vec3b>(); it != end; ++it){
                im.data[count] = (float)(*it)[2]/255.0;
                im.data[w*h + count] = (float)(*it)[1]/255.0;
                im.data[w*h*2 + count] = (float)(*it)[0]/255.0;
                ++count;
            }
            break;
        }

        default:
            printf("Channel number not supported.\n");
            break;
    }
    return im;
}