#include <opencv2/opencv.hpp>
//txs
#include "opencv2/video/tracking.hpp"
#include "opencv2/highgui/highgui.hpp"
#include "cameraobjectarray.pb.h"
//txs
#include "TrtNet.h"
#include "argsParser.h"
#include "configs.h"
#include <chrono>
#include "YoloLayer.h"
#include "dataReader.h"
#include "eval.h"

#include "imageBuffer.h"

#include "modulecomm.h"
#include "xmlparam.h"
#include "rawpic.pb.h"

#include <vector>
#include <QFile>
#include <QTextStream>
#include "localization.h"
#include <QDateTime>

#include "Tracking.h"

using namespace std;
using namespace argsParser;
using namespace Tn;
using namespace Yolo;

//txs
using namespace cv;
//iv::vision::cameraobjectarray gobj;
int gntemp = 0;
void * gpa;
void * gpcamout;
void * gdetec;
qint64 pictime;
//txs
vector<string> cocoNames;

ConsumerProducerQueue<cv::Mat> *imageBuffer = new ConsumerProducerQueue<cv::Mat>(1, true);

std::string gstreamer_pipeline (int capture_width, int capture_height, int display_width, int display_height, int framerate, int flip_method) {
    return "nvarguscamerasrc ! video/x-raw(memory:NVMM), width=(int)" + std::to_string(capture_width) + ", height=(int)" +
           std::to_string(capture_height) + ", format=(string)NV12, framerate=(fraction)" + std::to_string(framerate) +
           "/1 ! nvvidconv flip-method=" + std::to_string(flip_method) + " ! video/x-raw, width=(int)" + std::to_string(display_width) + ", height=(int)" +
           std::to_string(display_height) + ", format=(string)BGRx ! videoconvert ! video/x-raw, format=(string)BGR ! appsink";
}

static vector<string> GetNamesFromFile(string filePathAndName)
{
    vector<string> retVec;
    QFile qf(QString::fromStdString(filePathAndName));
    if(qf.exists())
    {
        if(qf.open(QIODevice::ReadOnly))
        {
            QTextStream in(&qf);
            while(!in.atEnd())
            {
                retVec.push_back(in.readLine(1024).toStdString());
            }
            qf.close();
        }
    }

    return retVec;
}

vector<float> prepareImage(cv::Mat& img)
{
    using namespace cv;

    int c = parser::getIntValue("C");
    int h = parser::getIntValue("H");   //net h
    int w = parser::getIntValue("W");   //net w

    float scale = min(float(w)/img.cols,float(h)/img.rows);
    auto scaleSize = cv::Size(img.cols * scale,img.rows * scale);

    cv::Mat rgb ;
    cv::cvtColor(img, rgb, CV_BGR2RGB);
    cv::Mat resized;
    cv::resize(rgb, resized,scaleSize,0,0,INTER_CUBIC);

    cv::Mat cropped(h, w,CV_8UC3, 127);
    Rect rect((w- scaleSize.width)/2, (h-scaleSize.height)/2, scaleSize.width,scaleSize.height);
    resized.copyTo(cropped(rect));

    cv::Mat img_float;
    if (c == 3)
        cropped.convertTo(img_float, CV_32FC3, 1/255.0);
    else
        cropped.convertTo(img_float, CV_32FC1 ,1/255.0);

    //HWC TO CHW
    vector<Mat> input_channels(c);
    cv::split(img_float, input_channels);

    vector<float> result(h*w*c);
    auto data = result.data();
    int channelLength = h * w;
    for (int i = 0; i < c; ++i) {
        memcpy(data,input_channels[i].data,channelLength*sizeof(float));
        data += channelLength;
    }

    return result;
}

void DoNms(vector<Detection>& detections,int classes ,float nmsThresh)
{
    auto t_start = chrono::high_resolution_clock::now();

    vector<vector<Detection>> resClass;
    resClass.resize(classes);

    for (const auto& item : detections)
        resClass[item.classId].push_back(item);

    auto iouCompute = [](float * lbox, float* rbox)
    {
        float interBox[] = {
            max(lbox[0] - lbox[2]/2.f , rbox[0] - rbox[2]/2.f), //left
            min(lbox[0] + lbox[2]/2.f , rbox[0] + rbox[2]/2.f), //right
            max(lbox[1] - lbox[3]/2.f , rbox[1] - rbox[3]/2.f), //top
            min(lbox[1] + lbox[3]/2.f , rbox[1] + rbox[3]/2.f), //bottom
        };
        
        if(interBox[2] > interBox[3] || interBox[0] > interBox[1])
            return 0.0f;

        float interBoxS =(interBox[1]-interBox[0])*(interBox[3]-interBox[2]);
        return interBoxS/(lbox[2]*lbox[3] + rbox[2]*rbox[3] -interBoxS);
    };

    vector<Detection> result;
    for (int i = 0;i<classes;++i)
    {
        auto& dets =resClass[i];
        if(dets.size() == 0)
            continue;

        sort(dets.begin(),dets.end(),[=](const Detection& left,const Detection& right){
            return left.prob > right.prob;
        });

        for (unsigned int m = 0;m < dets.size() ; ++m)
        {
            auto& item = dets[m];
            result.push_back(item);
            for(unsigned int n = m + 1;n < dets.size() ; ++n)
            {
                if (iouCompute(item.bbox,dets[n].bbox) > nmsThresh)
                {
                    dets.erase(dets.begin()+n);
                    --n;
                }
            }
        }
    }

    //swap(detections,result);
    detections = move(result);

    auto t_end = chrono::high_resolution_clock::now();
    float total = chrono::duration<float, milli>(t_end - t_start).count();
    mTensorIvlog->info("Time taken for nms is %f ms.", total);
    //    cout << "Time taken for nms is " << total << " ms." << endl;
}


vector<Bbox> postProcessImg(cv::Mat& img,vector<Detection>& detections,int classes)
{
    using namespace cv;

    int h = parser::getIntValue("H");   //net h
    int w = parser::getIntValue("W");   //net w

    //scale bbox to img
    int width = img.cols;
    int height = img.rows;
    float scale = min(float(w)/width,float(h)/height);
    float scaleSize[] = {width * scale,height * scale};

    //correct box
    for (auto& item : detections)
    {
        auto& bbox = item.bbox;
        bbox[0] = (bbox[0] * w - (w - scaleSize[0])/2.f) / scaleSize[0];
        bbox[1] = (bbox[1] * h - (h - scaleSize[1])/2.f) / scaleSize[1];
        bbox[2] /= scaleSize[0];
        bbox[3] /= scaleSize[1];
    }
    
    //nms
    float nmsThresh = parser::getFloatValue("nms");
    if(nmsThresh > 0)
        DoNms(detections,classes,nmsThresh);

    vector<Bbox> boxes;
    for(const auto& item : detections)
    {
        auto& b = item.bbox;
        Bbox bbox =
        {
            item.classId,   //classId
            max(int((b[0]-b[2]/2.)*width),0), //left
            min(int((b[0]+b[2]/2.)*width),width), //right
            max(int((b[1]-b[3]/2.)*height),0), //top
            min(int((b[1]+b[3]/2.)*height),height), //bot
            item.prob       //score
        };
        boxes.push_back(bbox);
    }

    return boxes;
}

vector<string> split(const string& str, char delim)
{
    stringstream ss(str);
    string token;
    vector<string> container;
    while (getline(ss, token, delim)) {
        container.push_back(token);
    }

    return container;
}


//std::vector<Obstacle> objvec,iv::lidar::objectarray & lidarobjvec
//txs
void GetCameraObj(std::vector<Bbox> objvec,iv::vision::cameraobjectarray &cameraobjvec,double pictime)
{
    mTensorIvlog->info("Bbox count: %d", objvec.size());
    //    std::cout<<"Bbox count: "<<objvec.size()<<std::endl;
    int i;
    for(i=0;i<objvec.size();i++){
        Bbox x;
        //iv::vision::cameraobject cameraobj;
        x = objvec.at(i);
        mTensorIvlog->info("left: %d\tright: %d", x.left, x.right);
        //        std::cout<<"left: " <<x.left<<"\tright: "<<x.right<<std::endl;

        iv::vision::cameraobject obj;
        MatrixXd SpeedInCamera(2,1);
        SpeedInCamera<<x.speedx, x.speedy;
        MatrixXd SpeedInWorld=TranslateToStandardLocation::translateCameraToStandardLocation(SpeedInCamera);
        obj.set_id(x.id);

        obj.set_type(cocoNames[x.classId]);
        obj.set_outpointx((x.left+x.right)/2);
        //2019-10-10 09-01AM
        //obj.set_outpointy((x.bot+x.top)/2);
        obj.set_outpointy(x.bot);
        obj.set_speedx(SpeedInWorld(0,0));
        obj.set_speedy(SpeedInWorld(1,0));
        obj.set_x((x.left+x.right)/2);
        obj.set_y((x.bot + x.top)/2);
        obj.set_w(x.right - x.left);
        obj.set_h(x.bot - x.top);
        obj.set_con(x.score);

        mTensorIvlog->info("raw output x: %f\traw output y: %f", obj.outpointx(), obj.outpointy());
        //        std::cout<<"raw output x: "<<obj.outpointx()<<"\traw output y:" <<obj.outpointy()<<std::endl;

        iv::vision::cameraobject * po = cameraobjvec.add_obj();
        po->CopyFrom(obj);
        cameraobjvec.set_mstime(pictime);
    }

}
//txs

int main( int argc, char* argv[] )
{
    parser::ADD_ARG_STRING("prototxt",Desc("input yolov3 deploy"),DefaultValue(INPUT_PROTOTXT),ValueDesc("file"));
    parser::ADD_ARG_STRING("caffemodel",Desc("input yolov3 caffemodel"),DefaultValue(INPUT_CAFFEMODEL),ValueDesc("file"));
    parser::ADD_ARG_INT("C",Desc("channel"),DefaultValue(to_string(INPUT_CHANNEL)));
    parser::ADD_ARG_INT("H",Desc("height"),DefaultValue(to_string(INPUT_HEIGHT)));
    parser::ADD_ARG_INT("W",Desc("width"),DefaultValue(to_string(INPUT_WIDTH)));
    parser::ADD_ARG_STRING("calib",Desc("calibration image List"),DefaultValue(CALIBRATION_LIST),ValueDesc("file"));
    parser::ADD_ARG_STRING("mode",Desc("runtime mode"),DefaultValue(MODE), ValueDesc("fp32/fp16/int8"));
    parser::ADD_ARG_STRING("outputs",Desc("output nodes name"),DefaultValue(OUTPUTS));
    parser::ADD_ARG_INT("class",Desc("num of classes"),DefaultValue(to_string(DETECT_CLASSES)));
    parser::ADD_ARG_FLOAT("nms",Desc("non-maximum suppression value"),DefaultValue(to_string(NMS_THRESH)));
    parser::ADD_ARG_INT("batchsize",Desc("batch size for input"),DefaultValue("1"));
    parser::ADD_ARG_STRING("enginefile",Desc("load from engine"),DefaultValue(""));

    //input
    parser::ADD_ARG_STRING("input",Desc("input image file"),DefaultValue(INPUT_IMAGE),ValueDesc("file"));
    parser::ADD_ARG_STRING("evallist",Desc("eval gt list"),DefaultValue(EVAL_LIST),ValueDesc("file"));

    gTensorFault = new iv::Ivfault("TensorRT_yolov3");
    mTensorIvlog = new iv::Ivlog("TensorRT_youlov3");

    int capture_width = 1920 ;
    int capture_height = 1080 ;
    int display_width = 1920 ;
    int display_height = 1080 ;
    int framerate = 60 ;
    int flip_method = 0 ;

    QString    strpath = "./vision_TensorRT_yolov3.xml";

    mTensorIvlog->info("%s", strpath.data());
    iv::xmlparam::Xmlparam xp(strpath.toStdString());

    std::string xmlmsgname = xp.GetParam("msgname","image00");
//    string engineName = "./../vision_TensorRT_yolov3/model/yolov3_fp16.engine";
    string engineName = "/home/nvidia/beiqi/models/vision/yolov3_fp16.engine";


    mTensorIvlog->info("msgname is %s",xmlmsgname.data());

    gpcamout = iv::modulecomm::RegisterSend("camera_output",1000000,1);//data

    //txs
    gdetec = iv::modulecomm::RegisterSend("camera_detect",10000000,1);//img
    std::unique_ptr<trtNet> net;

    int batchSize = 1;
    net.reset(new trtNet(engineName));
    assert(net->getBatchSize() == batchSize);

    RUN_MODE run_mode = RUN_MODE::FLOAT16;

    int outputCount = net->getOutputSize()/sizeof(float);
    unique_ptr<float[]> outputData(new float[outputCount]);


//    string inputFileName = "/home/nvidia/TensorRT-Yolov3/test.jpg";
//    fileNames.push_back(inputFileName);

    list<vector<Bbox>> outputs;

    int classNum = 80;
    int c = 608;
    int h = 608;
    int w = 3;
    int batchCount = 0;
    vector<float> inputData;
//    inputData.reserve(h*w*c*batchSize);
//    vector<cv::Mat> inputImgs;

//    auto iter = fileNames.begin();
    //txs
    cv::namedWindow("result",WINDOW_NORMAL);

    std::string pipeline = gstreamer_pipeline(capture_width,
        capture_height,
        display_width,
        display_height,
        framerate,
        flip_method);
        std::cout << "Using pipeline: \n\t" << pipeline << "\n";

        cv::VideoCapture cap(pipeline, cv::CAP_GSTREAMER);
        if(!cap.isOpened()) {
        std::cout<<"Failed to open camera."<<std::endl;
        return (-1);
        }

    //txs
    cv::Mat img2;
    int count = 0;
    bool m_isTrackerInitialized = false;
    TrackerSettings settings;
    CTracker tracker(settings);
    while(1)
    {
        if (!cap.read(img2)) {
                std::cout<<"Capture read error"<<std::endl;
                break;
            }
        pictime = QDateTime::currentMSecsSinceEpoch();
//        cv::imshow("imgs",img2);
//        cv::waitKey(10);
        imageBuffer->add(img2);
        imageBuffer->consume(img2);
        if(!img2.rows)
        {
            gTensorFault->SetFaultState(1, 0, "img2.rows state Error");
            break;
        }

        vector<float> curInput = prepareImage(img2);

        net->doInference(curInput.data(), outputData.get(),1);

        mTensorIvlog->info("run");

        //        std::cout<<"run"<<std::endl;

        auto output = outputData.get();

        int detCount = output[0];
        //later detect result
        vector<Detection> result;
        result.resize(detCount);
        memcpy(result.data(), &output[1], detCount*sizeof(Detection));

        outputs.clear();

        auto boxes = postProcessImg(img2,result,classNum);

        //---------------------------------------------  init tracker  -------------------------------------------------
        cv::UMat ufirst = img2.getUMat(cv::ACCESS_READ);
        if (!m_isTrackerInitialized)
        {
            m_isTrackerInitialized = InitTracker(ufirst,tracker);
            if (!m_isTrackerInitialized)
            {
                std::cerr << "Tracker initialize error!!!" << std::endl;
                break;
            }
        }
        Tracking(img2,boxes,tracker);
        //--------------------------------------------  end tracking  ---------------------------------------------------
        outputs.emplace_back(boxes);
//        vector<string> cocoNames = GetNamesFromFile("./../vision_TensorRT_yolov3/model/coco.names");
        cocoNames = GetNamesFromFile("/home/nvidia/beiqi/models/vision/coco.names");
        //        cout<<"cocoNames: "<<endl<<endl;
        mTensorIvlog->info("cocoName:");
//        for(int i=0; i<cocoNames.size(); i++)
//        {
//            std::cout<<cocoNames[i]<<std::endl;
//            //            cout<<cocoNames[i]<<endl;
//            mTensorIvlog->info("%s",cocoNames[i].data());
//        }
        //        cout<<"cocoNames end"<<endl<<endl;
        mTensorIvlog->info("cocoNames End");


        //txs
        iv::vision::cameraobjectarray cameraobjvec;
        GetCameraObj(boxes,cameraobjvec,pictime);
        std::cout<<"time     :"<<cameraobjvec.mstime()<<std::endl;
        boxes.clear();
        int ntlen;
        std::string out = cameraobjvec.SerializeAsString();

        iv::modulecomm::ModuleSendMsg(gpcamout,out.data(),out.length());
        mTensorIvlog->info("--------------");
        //        std::cout<<"------------------"<<std::endl;
        mTensorIvlog->info("lenth is %d",out.length());

        //txs


        auto bbox = *outputs.begin();
        for(const auto& item : bbox)
        {
            cv::rectangle(img2,cv::Point(item.left,item.top),cv::Point(item.right,item.bot),cv::Scalar(0,0,255),3,8,0);
            std::cout << "id=" << item.id << " class=" << item.classId << " prob=" << item.score*100 << std::endl;
            cout << "left=" << item.left << " right=" << item.right << " top=" << item.top << " bot=" << item.bot << endl;
            cout << "speedx=" << item.speedx << " speedy=" << item.speedy << endl;
//            cout<<cocoNames[item.classId]<<endl;
            //txs
            MatrixXd local(2,1);
            local<<(item.left +item.right)/2,(item.bot +item.bot)/2;
            //           local<<(item.bot +item.bot)/2,(item.left +item.right)/2;
            MatrixXd location=TranslateToStandardLocation::translateCameraToStandardLocation(local);
            //           string str = cocoNames[item.classId] +":" + std::to_string(item.score*100)+":"+std::to_string(location(0,0))+std::to_string(location(1,0));
            string str = std::to_string((int)(location(0,0)/100.0))+" " + std::to_string((int)(location(1,0)/100.0));
//            string str1=cocoNames[item.classId];

            cv::putText(img2,str, cv::Point(item.left, item.top), cv::FONT_HERSHEY_PLAIN, 1, cv::Scalar(255, 0, 0));
//            cv::putText(img2,str1, cv::Point(item.left, item.top-15), cv::FONT_HERSHEY_PLAIN, 1, cv::Scalar(0, 255, 0));
            //txs
        }


        //        ou.write(img2);
        //txs
        //        cv::imwrite("img/"+ std::to_string(count)+".jpg", img2);
        //txs

        iv::vision::rawpic pic;
        pic.set_time(pictime);
//        pic.set_index(gpicindex);
        pic.set_elemsize(img2.elemSize());
        pic.set_width(img2.cols);
        pic.set_height(img2.rows);
        pic.set_mattype(img2.type());

        std::vector<int> param = std::vector<int>(2);
        param[0] = CV_IMWRITE_JPEG_QUALITY;
        param[1] = 95; // default(95) 0-100
        std::vector<unsigned char> buff;
        cv::imencode(".jpg",img2,buff,param);
        pic.set_picdata(buff.data(),buff.size());
        buff.clear();
        pic.set_type(2);

        unsigned int nsize = pic.ByteSize();
        char * strser = new char[nsize];
        bool bser = pic.SerializeToArray(strser,nsize);
        if(bser)iv::modulecomm::ModuleSendMsg(gdetec,strser,nsize);

        delete[] strser;


        count += 1;
        cv::imshow("result",img2);


        cv::waitKey(1);
        curInput.clear();
        //        free(output_cpu);
        //        free(output2_cpu);
        //        img2.release();
        //        ou.release();

    }
    img2.release();
//    ou.release();
//    gTensorFault->SetFaultState(0, 0, "ok");

    return 0;
}