#include "vp_track_node.h"
//#include "../objects/shapes/vp_rect.h"

namespace vp_nodes {
        
    vp_track_node::vp_track_node(std::string node_name, 
                                vp_track_for track_for): 
                                vp_node(node_name), 
                                track_for(track_for) {
    }
    
    vp_track_node::~vp_track_node() {
    }

    std::shared_ptr<vp_objects::vp_meta> vp_track_node::handle_control_meta(std::shared_ptr<vp_objects::vp_control_meta> meta) {
        return meta;
    }

    std::shared_ptr<vp_objects::vp_meta> vp_track_node::handle_frame_meta(std::shared_ptr<vp_objects::vp_frame_meta> meta) {
        // channel_index can be different each call
        auto channel_index = meta->channel_index;

        // data used for tracking
        std::vector<vp_objects::vp_rect> rects;      // rects of targets
        std::vector<std::vector<float>> embeddings;  // embeddings of targets
        std::vector<int> track_ids;                  // track ids of targets

        // step 1, collect data
        preprocess(meta, rects, embeddings);

        // step 2, track by channel
        track(channel_index, rects, embeddings, track_ids);

        // step 3, postprocess
        postprocess(meta, rects, embeddings, track_ids);

        return meta;
    }

    void vp_track_node::preprocess(std::shared_ptr<vp_objects::vp_frame_meta> frame_meta, 
                                std::vector<vp_objects::vp_rect>& target_rects, 
                                std::vector<std::vector<float>>& target_embeddings) {
        if (track_for == vp_track_for::NORMAL) {
            for(auto& i: frame_meta->targets) {
                target_rects.push_back(i->get_rect());      // rect fo target (via i variable)
                target_embeddings.push_back(i->embeddings); // embeddings of target (via i variable)
            }
        }

        if (track_for == vp_track_for::FACE) {
            for(auto& i: frame_meta->face_targets) {
                target_rects.push_back(i->get_rect());       // rect of face target (via i variable)
                target_embeddings.push_back(i->embeddings);  // embeddings of face target (via i variable)
            }
        }
        /* ... extend for more track for... */
    }

    // write track_ids back to frame meta
    // we can also cache history rects for each target, and then push them back to tracks field (such as vp_frame_target::tracks)
    void vp_track_node::postprocess(std::shared_ptr<vp_objects::vp_frame_meta> frame_meta, 
                    const std::vector<vp_objects::vp_rect>& target_rects, 
                    const std::vector<std::vector<float>>& target_embeddings, 
                    const std::vector<int>& track_ids) {

        if (track_ids.empty()) {
            return;
        }
        // assert for length of vectors since they are generated by step1 & step2 separately
        // assert(target_rects.size() == target_embeddings.size());
        assert(target_rects.size() == track_ids.size());

        // support multi channels
        auto& tracks_by_id = all_tracks_by_id[frame_meta->channel_index];
        auto& last_tracked_frame_indexes = all_last_tracked_frame_indexes[frame_meta->channel_index];

        if (track_for == vp_track_for::NORMAL) {
            //assert(target_rects.size() == frame_meta->targets.size());
            for (int i = 0; i < frame_meta->targets.size(); i++) {
                auto& target = frame_meta->targets[i];
                auto& rect = target_rects[i];
                auto& track_id = track_ids[i];

                // -1 means no track result returned yet
                if (track_id != -1) {
                    tracks_by_id[track_id].push_back(rect);                           // cache
                    last_tracked_frame_indexes[track_id] = frame_meta->frame_index;   // update stamp

                    target->track_id = track_id;               // write track_id back to target
                    target->tracks = tracks_by_id[track_id];   // write tracks back to target
                }
            }
        }

        if (track_for == vp_track_for::FACE) {
            // assert(target_rects.size() == frame_meta->face_targets.size());
            for (int i = 0; i < frame_meta->face_targets.size(); i++) {
                auto& face = frame_meta->face_targets[i];
                auto& rect = target_rects[i];
                auto& track_id = track_ids[i];

                // -1 means no track result returned yet
                if (track_id != -1) {
                    tracks_by_id[track_id].push_back(rect);                           // cache
                    last_tracked_frame_indexes[track_id] = frame_meta->frame_index;   // update stamp

                    face->track_id = track_id;                // write track_id back to face target
                    face->tracks = tracks_by_id[track_id];    // write tracks back to face target
                }
            }
        }
        /* ... extend for more track for... */

        // remove cache tracks if has been long time since last updated (maybe it disappeared already).
        for (auto i = last_tracked_frame_indexes.begin(); i != last_tracked_frame_indexes.end();) {
            if (frame_meta->frame_index - (i->second) > max_allowed_disappear_frames 
                || frame_meta->frame_index < i->second) {
                VP_DEBUG(vp_utils::string_format("[%s] [tracking] long time no update, so erase cache of tracks for track_id:`%d`, size of tracks is:`%d`", node_name.c_str(), i->first, tracks_by_id[i->first].size()));
                tracks_by_id.erase(i->first);              // erase tracks first
                i = last_tracked_frame_indexes.erase(i);   // erase stamp then
            }
            else {
                i++;
            }
        }
    }
}