#include "OCSort.h"
#include "iomanip"
#include <utility>

namespace ANSOCSort {
    template<typename Matrix>
    std::ostream& operator<<(std::ostream& os, const std::vector<Matrix>& v) {
        os << "{";
        for (auto it = v.begin(); it != v.end(); ++it) {
            os << "(" << *it << ")\n";
            if (it != v.end() - 1) os << ",";
        }
        os << "}\n";
        return os;
    }
    OCSort::OCSort() {
        update_parameters(0, 50, 1, 0.22136877277096445, 1, "giou", 0.3941737016672115, true);
    }

    void OCSort::update_parameters(float det_thresh_, int max_age_, int min_hits_, float iou_threshold_, int delta_t_, std::string asso_func_, float inertia_, bool use_byte_, bool autoFrameRate) {
        max_age = max_age_;
        min_hits = min_hits_;
        iou_threshold = iou_threshold_;
        trackers.clear();
        frame_count = 0;
        det_thresh = det_thresh_;
        delta_t = delta_t_;
        std::unordered_map<std::string, std::function<Eigen::MatrixXf(const Eigen::MatrixXf&, const Eigen::MatrixXf&)>> ASSO_FUNCS{
            {"iou", iou_batch},
            { "giou", giou_batch }};
        ;
        std::function<Eigen::MatrixXf(const Eigen::MatrixXf&, const Eigen::MatrixXf&)> asso_func = ASSO_FUNCS[asso_func_];
        inertia = inertia_;
        use_byte = use_byte_;
        KalmanBoxTracker::count = 0;

        // Frame rate auto-estimation
        auto_frame_rate_ = autoFrameRate;
        estimated_fps_ = 30.0f;
        time_scale_factor_ = 1.0f;
        fps_sample_count_ = 0;
        has_last_update_time_ = false;
    }

    float OCSort::getEstimatedFps() const {
        return estimated_fps_;
    }

    void OCSort::estimateFrameRate() {
        if (!auto_frame_rate_) return;

        auto now = std::chrono::steady_clock::now();
        if (!has_last_update_time_) {
            last_update_time_ = now;
            has_last_update_time_ = true;
            return;
        }

        double delta_sec = std::chrono::duration<double>(now - last_update_time_).count();
        last_update_time_ = now;

        if (delta_sec <= 0.001) return; // Skip near-zero intervals

        float instant_fps = static_cast<float>(1.0 / delta_sec);

        // EMA smoothing: alpha=0.3 for warmup (first 10 samples), alpha=0.1 steady state
        fps_sample_count_++;
        float alpha = (fps_sample_count_ <= 10) ? 0.3f : 0.1f;
        estimated_fps_ = alpha * instant_fps + (1.0f - alpha) * estimated_fps_;

        // Clamp estimated fps to a reasonable range
        estimated_fps_ = std::max(1.0f, std::min(estimated_fps_, 120.0f));

        // Compute time scale factor: ratio of actual interval to expected interval
        float expected_dt = 1.0f / estimated_fps_;
        time_scale_factor_ = static_cast<float>(delta_sec) / expected_dt;
        time_scale_factor_ = std::max(0.5f, std::min(time_scale_factor_, 10.0f));
    }
    std::ostream& precision(std::ostream& os) {
        os << std::fixed << std::setprecision(2);
        return os;
    }

    std::vector<Eigen::RowVectorXf> OCSort::update(Eigen::MatrixXf dets) {
        // Estimate frame rate from timing of update() calls
        estimateFrameRate();

        frame_count += 1;
        std::vector<Eigen::RowVectorXf> ret;
        try {
            Eigen::Matrix<float, 1, Eigen::Dynamic> confs = dets.col(4);
            Eigen::MatrixXf output_results = dets;

            auto inds_low = confs.array() > 0.1;
            auto inds_high = confs.array() < det_thresh;
            auto inds_second = inds_low && inds_high;
            Eigen::Matrix<float, Eigen::Dynamic, 9> dets_second;
            Eigen::Matrix<bool, 1, Eigen::Dynamic> remain_inds = (confs.array() > det_thresh);
            Eigen::Matrix<float, Eigen::Dynamic, 9> dets_first;
            for (int i = 0; i < output_results.rows(); i++) {
                if (true == inds_second(i)) {
                    dets_second.conservativeResize(dets_second.rows() + 1, Eigen::NoChange);
                    dets_second.row(dets_second.rows() - 1) = output_results.row(i);
                }
                if (true == remain_inds(i)) {
                    dets_first.conservativeResize(dets_first.rows() + 1, Eigen::NoChange);
                    dets_first.row(dets_first.rows() - 1) = output_results.row(i);
                }
            }

            // Multi-predict: call predict() multiple times when frames are skipped
            int num_predicts = 1;
            if (auto_frame_rate_ && time_scale_factor_ > 1.5f) {
                num_predicts = std::min(static_cast<int>(std::round(time_scale_factor_)), 10);
            }

            Eigen::MatrixXf trks = Eigen::MatrixXf::Zero(trackers.size(), 5);
            std::vector<int> to_del;
            for (int i = 0; i < trks.rows(); i++) {
                // First predict call captures the return value for trks
                Eigen::RowVectorXf pos = trackers[i].predict();
                // Additional predict calls for frame skip compensation
                for (int p = 1; p < num_predicts; p++) {
                    pos = trackers[i].predict();
                }
                trks.row(i) << pos(0), pos(1), pos(2), pos(3), 0;
            }

            // Adaptive matching: relax iou_threshold during frame skips
            float effective_iou_threshold = iou_threshold;
            if (num_predicts > 1) {
                effective_iou_threshold = std::max(iou_threshold - 0.02f * (num_predicts - 1), 0.05f);
            }
            Eigen::MatrixXf velocities = Eigen::MatrixXf::Zero(trackers.size(), 2);
            Eigen::MatrixXf last_boxes = Eigen::MatrixXf::Zero(trackers.size(), 9);
            Eigen::MatrixXf k_observations = Eigen::MatrixXf::Zero(trackers.size(), 9);
            for (int i = 0; i < trackers.size(); i++) {
                velocities.row(i) = trackers[i].velocity;
                last_boxes.row(i) = trackers[i].last_observation;
                k_observations.row(i) = k_previous_obs(trackers[i].observations, trackers[i].age, delta_t);
            }

            std::vector<Eigen::Matrix<int, 1, 2>> matched;
            std::vector<int> unmatched_dets;
            std::vector<int> unmatched_trks;
            auto result = associate(dets_first, trks, effective_iou_threshold, velocities, k_observations, inertia);
            matched = std::get<0>(result);
            unmatched_dets = std::get<1>(result);
            unmatched_trks = std::get<2>(result);
            for (auto m : matched) {
                Eigen::Matrix<float, 9, 1> tmp_bbox;
                tmp_bbox = dets_first.block<1, 9>(m(0), 0);
                trackers[m(1)].update(&(tmp_bbox));
            }

            if (true == use_byte && dets_second.rows() > 0 && unmatched_trks.size() > 0) {
                Eigen::MatrixXf u_trks(unmatched_trks.size(), trks.cols());
                int index_for_u_trks = 0;
                for (auto i : unmatched_trks) {
                    u_trks.row(index_for_u_trks++) = trks.row(i);
                }
                Eigen::MatrixXf iou_left = giou_batch(dets_second, u_trks);
                if (iou_left.maxCoeff() > effective_iou_threshold) {
                    std::vector<std::vector<float>> iou_matrix(iou_left.rows(), std::vector<float>(iou_left.cols()));
                    for (int i = 0; i < iou_left.rows(); i++) {
                        for (int j = 0; j < iou_left.cols(); j++) {
                            iou_matrix[i][j] = -iou_left(i, j);
                        }
                    }
                    std::vector<int> rowsol, colsol;
                    float MIN_cost = execLapjv(iou_matrix, rowsol, colsol, true, 0.01, true);
                    std::vector<std::vector<int>> matched_indices;
                    for (int i = 0; i < rowsol.size(); i++) {
                        if (rowsol.at(i) >= 0) {
                            matched_indices.push_back({ colsol.at(rowsol.at(i)), rowsol.at(i) });
                        }
                    }

                    std::vector<int> to_remove_trk_indices;
                    for (auto m : matched_indices) {
                        int det_ind = m[0];
                        int trk_ind = unmatched_trks[m[1]];
                        if (iou_left(m[0], m[1]) < effective_iou_threshold) continue;

                        Eigen::Matrix<float, 9, 1> tmp_box;
                        tmp_box = dets_second.block<1, 9>(det_ind, 0);
                        trackers[trk_ind].update(&tmp_box);
                        to_remove_trk_indices.push_back(trk_ind);
                    }
                    std::vector<int> tmp_res1(unmatched_trks.size());
                    sort(unmatched_trks.begin(), unmatched_trks.end());
                    sort(to_remove_trk_indices.begin(), to_remove_trk_indices.end());
                    auto end1 = set_difference(unmatched_trks.begin(), unmatched_trks.end(),
                        to_remove_trk_indices.begin(), to_remove_trk_indices.end(),
                        tmp_res1.begin());
                    tmp_res1.resize(end1 - tmp_res1.begin());
                    unmatched_trks = tmp_res1;
                }
            }


            if (unmatched_dets.size() > 0 && unmatched_trks.size() > 0) {
                Eigen::MatrixXf left_dets(unmatched_dets.size(), 9);//6
                int inx_for_dets = 0;
                for (auto i : unmatched_dets) {
                    left_dets.row(inx_for_dets++) = dets_first.row(i);
                }
                Eigen::MatrixXf left_trks(unmatched_trks.size(), last_boxes.cols());
                int indx_for_trk = 0;
                for (auto i : unmatched_trks) {
                    left_trks.row(indx_for_trk++) = last_boxes.row(i);
                }
                Eigen::MatrixXf iou_left = giou_batch(left_dets, left_trks);
                if (iou_left.maxCoeff() > effective_iou_threshold) {
                    std::vector<std::vector<float>> iou_matrix(iou_left.rows(), std::vector<float>(iou_left.cols()));
                    for (int i = 0; i < iou_left.rows(); i++) {
                        for (int j = 0; j < iou_left.cols(); j++) {
                            iou_matrix[i][j] = -iou_left(i, j);
                        }
                    }
                    std::vector<int> rowsol, colsol;
                    float MIN_cost = execLapjv(iou_matrix, rowsol, colsol, true, 0.01, true);
                    std::vector<std::vector<int>> rematched_indices;
                    for (int i = 0; i < rowsol.size(); i++) {
                        if (rowsol.at(i) >= 0) {
                            rematched_indices.push_back({ colsol.at(rowsol.at(i)), rowsol.at(i) });
                        }
                    }
                    std::vector<int> to_remove_det_indices;
                    std::vector<int> to_remove_trk_indices;
                    for (auto i : rematched_indices) {
                        int det_ind = unmatched_dets[i.at(0)];
                        int trk_ind = unmatched_trks[i.at(1)];
                        if (iou_left(i.at(0), i.at(1)) < effective_iou_threshold) {
                            continue;
                        }
                        Eigen::Matrix<float, 9, 1> tmp_bbox;
                        tmp_bbox = dets_first.block<1, 9>(det_ind, 0);
                        trackers.at(trk_ind).update(&tmp_bbox);
                        to_remove_det_indices.push_back(det_ind);
                        to_remove_trk_indices.push_back(trk_ind);
                    }
                    std::vector<int> tmp_res(unmatched_dets.size());
                    sort(unmatched_dets.begin(), unmatched_dets.end());
                    sort(to_remove_det_indices.begin(), to_remove_det_indices.end());
                    auto end = set_difference(unmatched_dets.begin(), unmatched_dets.end(),
                        to_remove_det_indices.begin(), to_remove_det_indices.end(),
                        tmp_res.begin());
                    tmp_res.resize(end - tmp_res.begin());
                    unmatched_dets = tmp_res;
                    std::vector<int> tmp_res1(unmatched_trks.size());
                    sort(unmatched_trks.begin(), unmatched_trks.end());
                    sort(to_remove_trk_indices.begin(), to_remove_trk_indices.end());
                    auto end1 = set_difference(unmatched_trks.begin(), unmatched_trks.end(),
                        to_remove_trk_indices.begin(), to_remove_trk_indices.end(),
                        tmp_res1.begin());
                    tmp_res1.resize(end1 - tmp_res1.begin());
                    unmatched_trks = tmp_res1;
                }
            }

            for (auto m : unmatched_trks) {
                trackers.at(m).update(nullptr);
            }
            for (int i : unmatched_dets) {
                Eigen::RowVectorXf tmp_bbox = dets_first.block(i, 0, 1, 9);
                KalmanBoxTracker trk(tmp_bbox, delta_t);
                trackers.push_back(std::move(trk));
            }
            int tmp_i = trackers.size();
            for (int i = trackers.size() - 1; i >= 0; i--) {
                Eigen::Matrix<float, 1, 4> d;
                int last_observation_sum = trackers.at(i).last_observation.sum();
                if (last_observation_sum < 0) {
                    d = trackers.at(i).get_state();
                }
                else {
                    d = trackers.at(i).last_observation.block(0, 0, 1, 4);
                }
                if (trackers.at(i).time_since_update < 1 && ((trackers.at(i).hit_streak >= min_hits) | (frame_count <= min_hits))) {
                    Eigen::RowVectorXf tracking_res(10);
                    tracking_res << d(0), d(1), d(2), d(3),
                        trackers.at(i).id + 1,              // Tracker Id
                        trackers.at(i).GetClassId(),        // Class Id
                        trackers.at(i).GetScore(),          // Score
                        trackers.at(i).GetDGId(),           // DG Id
                        trackers.at(i).GetCamId(),          // Camera Id
                        trackers.at(i).GetModelId();        // Model Id
                    ret.push_back(tracking_res);
                }

                if (trackers.at(i).time_since_update > max_age) {
                    trackers.erase(trackers.begin() + i);
                }
            }
            return ret;
        }
        catch (const std::exception& e) {
            return ret;
        }
    }
}// 