#include <set>
#include "detect_obstacle.h"

namespace od{
// Computes IOU between two bounding boxes
double GetIOU(Rect_<float> bb_test, Rect_<float> bb_gt)
{
    float in = (bb_test & bb_gt).area();
    float un = bb_test.area() + bb_gt.area() - in;

    if (un < DBL_EPSILON)
        return 0;

    return (double)(in / un);
}


//tracking obstacle
bool TrackObstacle(int frame_count,vector<KalmanTracker> &trackers,vector<bbox_t> &outs,vector<od::TrackingBox> &track_result)
{
    // variables used in the for-loop
    vector<Rect_<float>> predictedBoxes;
    vector<vector<double>> iouMatrix;
    vector<int> assignment;
    set<int> unmatchedDetections;
    set<int> unmatchedTrajectories;
    set<int> allItems;
    set<int> matchedItems;
    vector<cv::Point> matchedPairs;
    unsigned int trkNum = 0;
    unsigned int detNum = 0;
    vector<od::DetectBox> detect_outs;
    //bbox_t to Detect_box
    for(unsigned int i=0;i<outs.size();i++)
    {
        od::DetectBox detect_temp;
        detect_temp.class_id = outs[i].obj_id;
        detect_temp.prob = outs[i].prob;
        float tpx = outs[i].x;
        float tpy = outs[i].y;
        float tpw = outs[i].w;
        float tph = outs[i].h;
        //detect_temp.box = Rect_<float>(Point_<float>(tpx, tpy),Point_<float>(tpx + tpw, tpy + tph));
        detect_temp.box = Rect_<float>(tpx,tpy,tpw,tph);
        detect_outs.push_back(detect_temp);
    }
    //tracking
    if (trackers.size() == 0) // the first frame met
    {
        // initialize kalman trackers using first detections.
        for (unsigned int i = 0; i < outs.size(); i++)
        {
            KalmanTracker trk = KalmanTracker(detect_outs[i].box,
                                              detect_outs[i].class_id,
                                              detect_outs[i].prob);
            trackers.push_back(trk);
        }
        return false;
    }
    ///////////////////////////////////////
    // 3.1. get predicted locations from existing trackers.
    predictedBoxes.clear();

    for (auto it = trackers.begin(); it != trackers.end();)
    {
        Rect_<float> pBox = (*it).predict();
        if (pBox.x >= 0 && pBox.y >= 0)
        {
            predictedBoxes.push_back(pBox);
            it++;
        }
        else
        {

            cerr << "Box invalid at frame: " << frame_count <<" id "<<(*it).m_id+1<<endl;
            it = trackers.erase(it);

        }
    }

    if (trackers.size() == 0 || detect_outs.size() == 0) return false;

    ///////////////////////////////////////
    // 3.2. associate detections to tracked object (both represented as bounding boxes)
    // dets : detFrameData[fi]
    trkNum = predictedBoxes.size();
    detNum = outs.size();

    iouMatrix.clear();
    iouMatrix.resize(trkNum, vector<double>(detNum, 0));

    for (unsigned int i = 0; i < trkNum; i++) // compute iou matrix as a distance matrix
    {
        for (unsigned int j = 0; j < detNum; j++)
        {
            // use 1-iou because the hungarian algorithm computes a minimum-cost assignment.
            iouMatrix[i][j] = 1 - GetIOU(predictedBoxes[i], detect_outs[j].box);
        }
    }
    // solve the assignment problem using hungarian algorithm.
    // the resulting assignment is [track(prediction) : detection], with len=preNum
    HungarianAlgorithm HungAlgo;
    assignment.clear();
    HungAlgo.Solve(iouMatrix, assignment);

    // find matches, unmatched_detections and unmatched_predictions
    unmatchedTrajectories.clear();
    unmatchedDetections.clear();
    allItems.clear();
    matchedItems.clear();
    if (detNum > trkNum) //	there are unmatched detections
    {
        for (unsigned int n = 0; n < detNum; n++)
            allItems.insert(n);

        for (unsigned int i = 0; i < trkNum; ++i)
            matchedItems.insert(assignment[i]);

        set_difference(allItems.begin(), allItems.end(),
                       matchedItems.begin(), matchedItems.end(),
                       insert_iterator<set<int>>(unmatchedDetections, unmatchedDetections.begin()));
    }
    else if (detNum < trkNum) // there are unmatched trajectory/predictions
    {
        for (unsigned int i = 0; i < trkNum; ++i)
            if (assignment[i] == -1) // unassigned label will be set as -1 in the assignment algorithm
                unmatchedTrajectories.insert(i);
    }
    // filter out matched with low IOU
    matchedPairs.clear();
    for (unsigned int i = 0; i < trkNum; ++i)
    {
        if (assignment[i] == -1) // pass over invalid values
            continue;
        if (1 - iouMatrix[i][assignment[i]] < od::iouThreshold)
        {
            unmatchedTrajectories.insert(i);
            unmatchedDetections.insert(assignment[i]);
        }
        else
            matchedPairs.push_back(cv::Point(i, assignment[i]));
    }
    ///////////////////////////////////////
    // 3.3. updating trackers
    // update matched trackers with assigned detections.
    // each prediction is corresponding to a tracker
    int detIdx, trkIdx;
    for (unsigned int i = 0; i < matchedPairs.size(); i++)
    {
        trkIdx = matchedPairs[i].x;
        detIdx = matchedPairs[i].y;
        trackers[trkIdx].update(detect_outs[detIdx].box,
                                detect_outs[detIdx].class_id,
                                detect_outs[detIdx].prob);
    }
    // create and initialise new trackers for unmatched detections
    for (auto umd : unmatchedDetections)
    {
        KalmanTracker tracker = KalmanTracker(detect_outs[umd].box,
                                              detect_outs[umd].class_id,
                                              detect_outs[umd].prob);
        trackers.push_back(tracker);
    }

#if 0
    //get unique trackers,merg same trackers
    unsigned int trackers_num = trackers.size();
    iouMatrix.clear();
    iouMatrix.resize(trackers_num, vector<double>(trackers_num, 0));
    for (unsigned int i = 0; i < trackers_num; i++) // compute iou matrix as a distance matrix
    {
        for (unsigned int j = 0; j < trackers_num; j++)
        {
            // use 1-iou because the hungarian algorithm computes a minimum-cost assignment.
            if(j==i)
                iouMatrix[i][j] = 1;
            else
                iouMatrix[i][j] = 1 - GetIOU(trackers[i].get_state(), trackers[j].get_state());
        }
    }
    // solve the assignment problem using hungarian algorithm.
    // the resulting assignment is [track(prediction) : detection], with len=preNum
    assignment.clear();
    HungAlgo.Solve(iouMatrix, assignment);
    // filter out matched with low IOU
    matchedPairs.clear();
    for (unsigned int i = 0; i < trackers_num; ++i)
    {
        if (assignment[i] == -1) // pass over invalid values
            continue;
        if (iouMatrix[i][assignment[i]] < od::iouThreshold)
        {
            matchedPairs.push_back(cv::Point(i, assignment[i]));
        }
    }
    int index1,index2;
    vector<int> delete_index;
    for (unsigned int i = 0; i < matchedPairs.size(); i++)
    {
        index1 = matchedPairs[i].x;
        index2 = matchedPairs[i].y;
        if(index1 >= index2)
            continue;
        if((trackers[index1].m_id > trackers[index2].m_id) && (trackers[index1].m_class_history.size()>0))
        {
            trackers[index1].m_id = trackers[index2].m_id;
            trackers[index1].m_class_history.insert(trackers[index1].m_class_history.begin(),
                                                    trackers[index2].m_class_history.begin(),trackers[index2].m_class_history.end());
            delete_index.push_back(index2);
        }
        else if((trackers[index2].m_id > trackers[index1].m_id) && (trackers[index2].m_class_history.size()>0))
        {
            trackers[index2].m_id = trackers[index1].m_id;
            trackers[index2].m_class_history.insert(trackers[index2].m_class_history.begin(),
                                                    trackers[index1].m_class_history.begin(),trackers[index1].m_class_history.end());
            delete_index.push_back(index1);

        }
    }
    for(unsigned int i = 0; i < delete_index.size(); i++)
    {
        int idx = delete_index[i] - i;
        trackers.erase(trackers.begin() + idx);
    }
#endif

    // get trackers' output
    track_result.clear();
    for (auto it = trackers.begin(); it != trackers.end();)
    {
        if (((*it).m_time_since_update <= od::max_age) &&
                ((*it).m_hit_streak >= od::min_hits || frame_count <= od::min_hits))
        {
            od::TrackingBox res;
            res.box = (*it).get_state();
            res.id = (*it).m_id + 1;
            res.frame = frame_count;
            res.class_id = (*it).m_class_id;
            res.prob = (*it).m_prob;
            res.class_history = (*it).m_class_history;
            track_result.push_back(res);
            it++;
        }
        else
            it ++;
        //remove dead tracklet
        if(it != trackers.end() && (*it).m_time_since_update > od::max_age)
        {
            it = trackers.erase(it);
        }
    }
    if(track_result.size()>0)
        return true;
    else return false;
}