#include<iostream>
#include "tkdnn.h"
#include <stdlib.h>     /* srand, rand */
#include <unistd.h>

#include <opencv2/core/core.hpp>
#include <opencv2/highgui/highgui.hpp>
#include <opencv2/imgproc/imgproc.hpp>

const char *reg_bias = "../tests/yolo/layers/g31.bin";

int prob_sort(const void *pa, const void *pb) {
    tk::dnn::box a = *(tk::dnn::box *)pa;
    tk::dnn::box b = *(tk::dnn::box *)pb;
    float diff = a.prob - b.prob;
    if(diff < 0) return 1;
    else if(diff > 0) return -1;
    return 0;
}

cv::Mat GetSquareImage(const cv::Mat& img, int target_width) {
    int width = img.cols, height = img.rows;

    cv::Mat square = cv::Mat::zeros( target_width, target_width, img.type() );

    int max_dim = ( width >= height ) ? width : height;
    float scale = ( ( float ) target_width ) / max_dim;
    cv::Rect roi;
    if ( width >= height )
    {
        roi.width = target_width;
        roi.x = 0;
        roi.height = height * scale;
        roi.y = ( target_width - roi.height ) / 2;
    }
    else
    {
        roi.y = 0;
        roi.height = target_width;
        roi.width = width * scale;
        roi.x = ( target_width - roi.width ) / 2;
    }

    cv::resize( img, square( roi ), roi.size() );

    return square;
}

//return inference time
double compute_image( cv::Mat imageORIG, 
                    tk::dnn::NetworkRT *netRT, tk::dnn::RegionInterpret *rI,
                    dnnType *input, dnnType *output) {

    //Resize with padding and convert to float
    cv::Mat image = GetSquareImage(imageORIG, netRT->input_dim.w);
    cv::Mat imageF;
    image.convertTo(imageF, CV_32FC3, 1/255.0); 

    //split channels
    cv::Mat bgr[3];   //destination array
    cv::split(imageF,bgr);//split source
    
    //write channels
    int idx = 0;
    memcpy((void*)&input[idx], (void*)bgr[2].data, imageF.rows*imageF.cols*sizeof(dnnType));
    idx = imageF.rows*imageF.cols;
    memcpy((void*)&input[idx], (void*)bgr[1].data, imageF.rows*imageF.cols*sizeof(dnnType));
    idx *= 2;    
    memcpy((void*)&input[idx], (void*)bgr[0].data, imageF.rows*imageF.cols*sizeof(dnnType));

    //DO INFERENCE
    printCenteredTitle(" TENSORRT inference ", '=', 30); 
    TIMER_START
    checkCuda(  cudaMemcpyAsync(netRT->buffersRT[netRT->buf_input_idx], input, 
                                netRT->input_dim.tot()*sizeof(float), 
                                cudaMemcpyHostToDevice, netRT->stream));
    netRT->enqueue();
    checkCuda(  cudaMemcpyAsync(output, netRT->buffersRT[netRT->buf_output_idx], 
                                netRT->output_dim.tot()*sizeof(float), 
                                cudaMemcpyDeviceToHost, netRT->stream));
    cudaStreamSynchronize(netRT->stream);
    TIMER_STOP
    

    rI->interpretData(output, imageORIG.cols, imageORIG.rows);

    return t_ns;
}

int print_usage() {
    std::cout<<"usage: ./detection net.rt validation_list.txt"
             <<" [-t <thresh>] [-s] [-i <iterations>]\n"
             <<" -t: set thresh value\n -s: show images as compute\n"
             <<" -i: images to compute\n\n" 
             <<"> validation_list.txt format: \n"
             <<"    path/to/image.jpg path/to/label.txt\n"
             <<"> label.txt format: \n"  
             <<"    <object-class> <x> <y> <width> <height>\n"
             <<"    x and y are the box center, "   
             <<"all values are relative to the image size\n\n";
    return 1;
}


int main(int argc, char *argv[]) {

    //params
    char *tensor_path = NULL;
    char *imageset_path = NULL; 
    float thresh = 0.3f;
    bool  show = false;
    int iterations = INT_MAX;

    //parse params
    int c;
    while ((c = getopt (argc, argv, "t:si:")) != -1) {
        switch(c) {
        case 't':   thresh = atof(optarg);       break;
        case 's':   show = true;                 break;
        case 'i':   iterations = atoi(optarg);   break;
        case '?':   
            return print_usage();
        default:    return print_usage();
        }
    }

    if(argc - optind == 2) {
        tensor_path   = argv[optind];
        imageset_path = argv[optind+1];
    } else {
        std::cout<<"not enough arguments.\n";
        return print_usage();
    }
    //end parsing

    if(!fileExist(tensor_path))
        FatalError("unable to read serialRT file");
    //convert network to tensorRT
    tk::dnn::NetworkRT netRT(NULL, tensor_path);
    tk::dnn::RegionInterpret rI(netRT.input_dim, netRT.output_dim, 80, 4, 5, thresh, reg_bias);

    dnnType *input = new float[netRT.input_dim.tot()];
    dnnType *output = new float[netRT.output_dim.tot()];

    std::string line;
    std::ifstream imageset(imageset_path);
    if(!imageset.is_open())
        FatalError("could not read imageset");
    
    double mTime = 0;
    float mAP = 0;
    int processed_images;

    for(processed_images=1; 
        processed_images-1 < iterations && getline(imageset, line); 
        processed_images++) {
        
        std::string image_path = line.substr(0, line.find(" "));
        std::string label_path = line.substr(line.find(" ")+1, line.size());
        std::cout<<image_path<<"\n"<<label_path<<"\n";

        //LOAD IMAGE
        cv::Mat img = cv::imread(image_path.c_str(), CV_LOAD_IMAGE_COLOR);
        if(!img.data)                              
            FatalError("Could not open image");
        std::cout<<"Image size: ("<<img.cols<<"x"<<img.rows<<")\n";

        mTime += compute_image(img, &netRT, &rI, input, output); 

        std::ifstream labels(label_path.c_str());
        if(!labels.is_open())
            FatalError("could not read labels");


        qsort(rI.res_boxes, rI.res_boxes_n, sizeof(tk::dnn::box), prob_sort);
        for(int i=0; i<rI.res_boxes_n; i++) {
            tk::dnn::box bx = rI.res_boxes[i];
            std::cout<<" ("<<int(bx.prob*100)<<"%) "<<bx.cl
                     <<": "<<bx.x<<" "<<bx.y<<" "<<bx.w<<" "<<bx.h<<"\n";

            cv::rectangle(img, cv::Point(bx.x - bx.w/2, bx.y - bx.h/2), 
                               cv::Point(bx.x + bx.w/2, bx.y + bx.h/2),
            cv::Scalar( 0, 0, 255), 2);
        }

        std::cout<<"GROUND TRUTH\n";
        tk::dnn::box gt[256];
        int gt_n = 0; 
        int cl;
        float x, y, w, h;
        while(labels>>cl) {
            labels>>x>>y>>w>>h;
            w *= img.cols; x *= img.cols; 
            h *= img.rows; y *= img.rows; 
            std::cout<<cl<<": "<<x<<" "<<y<<" "<<w<<" "<<h<<"\n";
            gt[gt_n].x = x;
            gt[gt_n].y = y;
            gt[gt_n].w = w;
            gt[gt_n].h = h;
            gt[gt_n].cl = cl;
            gt_n++;            

            cv::rectangle(img, cv::Point(x -w/2, y -h/2), 
                               cv::Point(x +w/2, y +h/2),
            cv::Scalar( 255, 0, 0), 2);
        }

        //AP calculation
        float AP = 0;
        for(int i=rI.res_boxes_n; i>=1; i--) {  //for each detected evaluate sub group
            
            int prec = 0;
            for(int j=0; j<i; j++) {            //for each detected in sub group
                for(int z=0; z<gt_n; z++) {     //control each ground truth
                    float iou = tk::dnn::RegionInterpret::box_iou(rI.res_boxes[j], gt[z]);
                    if(iou > 0.6f && rI.res_boxes[j].cl == gt[z].cl) {
                        prec++;
                        break;
                    }
                }
            }

            AP += float(prec)/i;
        }
        AP = AP/gt_n;
        std::cout<<"AP: "<<AP<<"\n";        
        
        mAP += AP;
        std::cout<<"#### processed: "<<processed_images
                 <<", mAP: "<<mAP/processed_images<<"\n";        

        //show results
        if(show) {
            cv::namedWindow("result");
            cv::imshow("result", img);
            cv::waitKey(10);
        }        
    }

    //print results to file
    processed_images -= 1;
    std::ofstream res("results.txt", std::ios::app);
    res<<"#### "<<tensor_path<<"\n";
    res<<"processed images: "<<processed_images<<"\n";
    res<<"mean inference time: "<<mTime/processed_images<<"\n";
    res<<"mean AP: "<<mAP/processed_images<<"\n";
    res<<"thesh used: "<<thresh<<"\n\n";
    return 0;
}