#include <opencv2/opencv.hpp>
#include <Eigen/Eigen>
#include <Eigen/Dense>
#include <cmath>
#include <iostream>
#include <algorithm>
#include "UCMClapjv.h"
#include "UCMCTracker.h"

namespace UCMC 
{
   
    Mapper::Mapper(std::vector<double>& _Ki, std::vector<double>& _Ko) {
        Eigen::Map<Eigen::Matrix<double, 4, 3>> KiT(_Ki.data());
        Eigen::Map<Eigen::Matrix4d> KoT(_Ko.data());

        Eigen::Matrix<double, 3, 4> Ki = KiT.transpose();
        Eigen::Matrix4d Ko = KoT.transpose();

        A.setZero();
        InvA.setZero();
        KiKo = Ki * Ko;

        for (int now_row = 0; now_row < 3; now_row++)
            for (int now_col = 0; now_col < 3; now_col++) {
                A(now_col, now_row) = KiKo(now_col, now_row);
                if (now_row == 2) {
                    A(now_col, now_row) = KiKo(now_col, 3);
                }
            }

        InvA = A.inverse();

    }

    std::vector<double> Mapper::uvError(cv::Rect box) {
        std::vector<double> uv;
        uv.resize(2);
        uv[0] = MAX(MIN(13., 0.05 * box.width), 2.);
        uv[1] = MAX(MIN(10., 0.05 * box.height), 2.);
        return uv;
    }

    void Mapper::uv2xy(Eigen::MatrixXd uv, Eigen::MatrixXd sigma_uv,
        Eigen::Matrix<double, 2, 1>& xy, Eigen::Matrix2d& sigma_xy) {
        Eigen::Matrix<double, 3, 1> uv1;
        uv1(0, 0) = uv(0, 0);
        uv1(1, 0) = uv(1, 0);
        uv1(2, 0) = 1.;


        Eigen::MatrixXd b = InvA * uv1;

        if (std::abs(b(2, 0)) < 1e-10) {
            xy.setZero();
            sigma_xy.setZero();
            return;
        }
        double gamma = 1. / b(2, 0);

        Eigen::Matrix2d C = gamma * InvA.block(0, 0, 2, 2)
            - (gamma * gamma) * b.block(0, 0, 2, 1) * InvA.block(2, 0, 1, 2);

        // INFO << "C:\n" << C << ENDL;
        // INFO << "sigma_uv:\n" << sigma_uv << ENDL;

        xy = b.block(0, 0, 2, 1) * gamma;
        sigma_xy = (C * sigma_uv) * C.transpose();

    }


    void Mapper::map_to(cv::Rect box, Eigen::Matrix<double, 2, 1>& y, Eigen::Matrix2d& R) {
        Eigen::Matrix<double, 2, 1> uv;
        uv(0, 0) = box.x + 0.5 * box.width;
        uv(1, 0) = box.y + box.height;
        std::vector<double> uv_err = uvError(box);
        Eigen::MatrixXd sigma_uv = Eigen::Matrix2d::Identity();
        sigma_uv(0, 0) = uv_err[1] * uv_err[1];
        sigma_uv(1, 1) = uv_err[1] * uv_err[1];
        uv2xy(uv, sigma_uv, y, R);
    }

    //  Tracker
    void Tracker::update_parameters(double a1, double a2,
        double wx, double wy,
        double vmax, double max_age,
        double high_score, double conf_threshold,
        double dt, const std::vector<double>& Ki,
        const std::vector<double>& Ko,
        bool autoFrameRate) {
            this->params.a1 = a1;
            this->params.a2 = a2;
            this->params.vmax = vmax;
            this->params.max_age = std::max(5.0, max_age);
            this->params.high_score = high_score;
            this->params.conf_threshold = conf_threshold;
            this->params.dt = dt;
            this->params.Ki = Ki;
            this->params.Ko = Ko;
            mapper = Mapper(params.Ki, params.Ko);

            // 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 Tracker::getEstimatedFps() const {
        return estimated_fps_;
    }

    void Tracker::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));
    }

    Tracker::Tracker(Params& params)
    {
        this->params = params;
        mapper = Mapper(params.Ki, params.Ko);
    }

    std::vector<Obj> Tracker::update(std::vector<Object>& det_results) {
        // Estimate frame rate from timing of update() calls
        estimateFrameRate();

        std::vector<Obj> dets;
        int id = 0;
        for (Object obj : det_results) {
            Obj this_obj;
            this_obj.id = id++;
            this_obj.obj = obj;
            mapper.map_to(obj.rect, this_obj.y, this_obj.R);

            dets.push_back(this_obj);
        }
        this->update(dets, frame_idx++);
        return dets;
    }

    void Tracker::update(std::vector<Obj>& dets, int frame_id) {
        this->data_association(dets, frame_id);
        this->associate_tentative(dets);
        this->initial_tentative(dets);
        this->delete_old_trackers();
        this->update_status(dets);
    }

    void Tracker::data_association(std::vector<Obj>& dets, int frame_id) {
        std::vector<int> detidx_high, detidx_low;
        for (size_t i = 0; i < dets.size(); ++i) {
            if (dets[i].obj.prob >= params.high_score) {
                detidx_high.push_back(i);
            }
            else {
                detidx_low.push_back(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);
        }
        for (int p = 0; p < num_predicts; p++) {
            for (auto& track : trackers) {
                track.predict();
            }
        }

        std::vector<int> trackidx_remain;
        detidx_remain.clear();

        // INFO << "try high score mapping" << ENDL;
        std::vector<int> trackidx = confirmed_idx;
        trackidx.insert(trackidx.end(), coasted_idx.begin(), coasted_idx.end());

        int num_det = detidx_high.size();
        int num_trk = trackidx.size();

        for (auto& track : trackers) {
            track.detidx = -1;
        }

        // Adaptive association thresholds: relax during frame skips
        // Multi-predict already increases covariance (making Mahalanobis more lenient),
        // but we add a small relaxation for extra tolerance
        double effective_a1 = params.a1;
        double effective_a2 = params.a2;
        if (num_predicts > 1) {
            effective_a1 = params.a1 * (1.0 + 0.1 * (num_predicts - 1));
            effective_a2 = params.a2 * (1.0 + 0.1 * (num_predicts - 1));
        }

        if (num_det * num_trk > 0) {
            Eigen::MatrixXd cost_matrix(num_det, num_trk);
            cost_matrix.setZero();

            for (int i = 0; i < num_det; ++i) {
                int det_idx = detidx_high[i];
                for (int j = 0; j < num_trk; ++j) {
                    cost_matrix(i, j) = trackers[trackidx[j]].distance(
                        dets[det_idx].y,
                        dets[det_idx].R,
                        i + j == 0 && debug
                    );
                }
            }

            // INFO << "det[0].y:\n" << dets[0].y << ENDL;
            // INFO << "det[0].R:\n" << dets[0].R << ENDL;
            // INFO << "cost matrix: " << cost_matrix << ENDL;

            std::vector<std::vector<int>> matched_indices;
            std::vector<int> unmatched_a;
            std::vector<int> unmatched_b;
            UCMCKalman::linearAssignment(cost_matrix, effective_a1,
                matched_indices, unmatched_a, unmatched_b);

            // INFO << "high: " << detidx_remain.size() << " " << unmatched_a.size() << " "
            //      << unmatched_b.size() << ENDL;

            for (int i : unmatched_a) {
                detidx_remain.push_back(detidx_high[i]);
            }
            for (int i : unmatched_b) {
                trackidx_remain.push_back(trackidx[i]);
            }

            for (size_t i = 0; i < matched_indices.size(); ++i) {
                int det_idx = detidx_high[matched_indices[i][0]];
                int trk_idx = trackidx[matched_indices[i][1]];

                trackers[trk_idx].update(dets[det_idx].y, dets[det_idx].R);
                trackers[trk_idx].death_count = 0;
                trackers[trk_idx].detidx = det_idx;
                trackers[trk_idx].status = UCMCKalman::TrackStatus::Confirmed;
                trackers[trk_idx].voteClassId(dets[det_idx].obj.label, dets[det_idx].obj.prob);
                dets[det_idx].obj.label = trackers[trk_idx].class_id_;

                dets[det_idx].track_idx = trackers[trk_idx].id;

            }
        }
        else {
            detidx_remain = detidx_high;
            trackidx_remain = trackidx;
        }


        int num_det_low = detidx_low.size();
        int num_trk_remain = trackidx_remain.size();

        // INFO << "try low score mapping" << ENDL;
        if (num_det_low * num_trk_remain > 0) {
            Eigen::MatrixXd cost_matrix(num_det_low, num_trk_remain);
            cost_matrix.setZero();

            for (int i = 0; i < num_det_low; ++i) {
                int det_idx = detidx_low[i];
                for (int j = 0; j < num_trk_remain; ++j) {
                    int trk_idx = trackidx_remain[j];
                    cost_matrix(i, j) = trackers[trk_idx].distance(
                        dets[det_idx].y,
                        dets[det_idx].R,
                        i + j == 0 && debug
                    );
                }
            }

            // INFO << "det[0].y:\n" << dets[0].y << ENDL;
            // INFO << "det[0].R:\n" << dets[0].R << ENDL;

            std::vector<std::vector<int>> matched_indices;
            std::vector<int> unmatched_a;
            std::vector<int> unmatched_b;
            UCMCKalman::linearAssignment(cost_matrix, effective_a2,
                matched_indices, unmatched_a, unmatched_b);

            for (int i : unmatched_b) {
                int trk_idx = trackidx_remain[i];
                trackers[trk_idx].status = UCMCKalman::TrackStatus::Coasted;
                // trackers[trk_idx].death_count += 1; 
                trackers[trk_idx].detidx = -1;
            }

            for (size_t i = 0; i < matched_indices.size(); ++i) {
                int det_idx = detidx_low[matched_indices[i][0]];
                int trk_idx = trackidx_remain[matched_indices[i][1]];

                trackers[trk_idx].update(dets[det_idx].y, dets[det_idx].R);
                trackers[trk_idx].death_count = 0;
                trackers[trk_idx].detidx = det_idx;
                trackers[trk_idx].status = UCMCKalman::TrackStatus::Confirmed;
                trackers[trk_idx].voteClassId(dets[det_idx].obj.label, dets[det_idx].obj.prob);
                dets[det_idx].obj.label = trackers[trk_idx].class_id_;
                dets[det_idx].track_idx = trackers[trk_idx].id;
            }
        }

    }

    void Tracker::associate_tentative(std::vector<Obj>& dets) {
        // INFO << "try associate_tentative mapping" << ENDL;
        size_t num_det = detidx_remain.size();
        size_t num_trk = tentative_idx.size();

        Eigen::MatrixXd cost_matrix(num_det, num_trk);
        cost_matrix.setZero();

        for (int i = 0; i < num_det; ++i) {
            int det_idx = detidx_remain[i];
            for (int j = 0; j < num_trk; ++j) {
                int trk_idx = tentative_idx[j];
                cost_matrix(i, j) = trackers[trk_idx].distance(
                    dets[det_idx].y,
                    dets[det_idx].R,
                    i + j == 0 && debug
                );
            }
        }

        // INFO << "det[0].y:\n" << dets[0].y << ENDL;
        // INFO << "det[0].R:\n" << dets[0].R << ENDL;

        std::vector<std::vector<int>> matched_indices;
        std::vector<int> unmatched_a;
        std::vector<int> unmatched_b;
        UCMCKalman::linearAssignment(cost_matrix, this->params.a1,
            matched_indices, unmatched_a, unmatched_b);

        // INFO << detidx_remain.size() << " " << unmatched_a.size() << " "
        //      << unmatched_b.size() << ENDL;

        for (size_t i = 0; i < matched_indices.size(); ++i) {
            int det_idx = detidx_remain[matched_indices[i][0]];
            int trk_idx = tentative_idx[matched_indices[i][1]];

            trackers[trk_idx].update(dets[det_idx].y, dets[det_idx].R);
            trackers[trk_idx].death_count = 0;
            trackers[trk_idx].birth_count++;
            trackers[trk_idx].detidx = det_idx;
            trackers[trk_idx].voteClassId(dets[det_idx].obj.label, dets[det_idx].obj.prob);
            dets[det_idx].obj.label = trackers[trk_idx].class_id_;
            dets[det_idx].track_idx = trackers[trk_idx].id;
            if (trackers[trk_idx].birth_count >= 2) {
                trackers[trk_idx].birth_count = 0;
                trackers[trk_idx].status = UCMCKalman::TrackStatus::Confirmed;
            }
        }

        for (int i : unmatched_b) {
            int trk_idx = tentative_idx[i];
            trackers[trk_idx].detidx--;
        }

        std::vector<int> unmatched_detidx;
        for (int i : unmatched_a) {
            unmatched_detidx.push_back(detidx_remain[i]);
        }
        detidx_remain = unmatched_detidx;
    }

    void Tracker::initial_tentative(std::vector<Obj>& dets) {
        // INFO << detidx_remain.size() << ENDL;
        for (int i : detidx_remain) {
            // INFO << "detidx:" << i << "" << dets[i].obj.rect << ENDL;
            UCMCKalman::Tracker new_obj(
                dets[i].y,
                dets[i].R,
                this->params.wx,
                this->params.wy,
                this->params.vmax,
                dets[i].obj.rect.width,
                dets[i].obj.rect.height,
                this->params.dt,
                ++this->tracker_count
            );
            new_obj.status = UCMCKalman::TrackStatus::Tentative;
            new_obj.detidx = i;
            new_obj.class_id_ = dets[i].obj.label;
            new_obj.voteClassId(dets[i].obj.label, dets[i].obj.prob);
            trackers.push_back(new_obj);
        }
        detidx_remain.clear();
    }

    void Tracker::delete_old_trackers() {
        std::vector<int> idx_reserve;
        std::vector<UCMCKalman::Tracker> reserve_trackers;
        for (int trk_idx = 0; trk_idx < trackers.size(); trk_idx++) {
            trackers[trk_idx].death_count++;
            if (!((trackers[trk_idx].status == UCMCKalman::TrackStatus::Coasted &&
                trackers[trk_idx].death_count >= this->params.max_age) ||
                (trackers[trk_idx].status == UCMCKalman::TrackStatus::Tentative &&
                    trackers[trk_idx].death_count >= 2))) {
                reserve_trackers.push_back(trackers.at(trk_idx));
            }
        }
        trackers = reserve_trackers;
    }

    void Tracker::update_status(std::vector<Obj>& dets) {
        confirmed_idx.clear();
        coasted_idx.clear();
        tentative_idx.clear();

        for (int i = 0; i < trackers.size(); i++) {
            int detidx = trackers[i].detidx;

            if (detidx >= 0 && detidx < dets.size()) {
                trackers[i].h = dets[detidx].obj.rect.height;
                trackers[i].w = dets[detidx].obj.rect.width;
            }

            switch (trackers[i].status)
            {
            case UCMCKalman::TrackStatus::Confirmed:
                confirmed_idx.push_back(i);
                break;
            case UCMCKalman::TrackStatus::Coasted:
                coasted_idx.push_back(i);
                break;
            case UCMCKalman::TrackStatus::Tentative:
                tentative_idx.push_back(i);
                break;
            default:
                break;
            }

        }
    }

    Tracker::~Tracker() {
        trackers.clear();
        confirmed_idx.clear();
        coasted_idx.clear();
        tentative_idx.clear();
        detidx_remain.clear();
    }



    static cv::Scalar get_color(int index) {
        index %= 20;
        return cv::Scalar(color_list[index][2], color_list[index][1], color_list[index][0]);
    }
    static cv::Mat draw_boxes(cv::Mat image,
        std::vector<UCMC::Obj>& objects,
        std::vector<std::string>& names,
        int draw_size,
        bool draw_label) {
        cv::Mat d_img = image.clone();
        cv::Scalar color;
        cv::Scalar txt_color;
        cv::Scalar txt_bk_color;
        cv::Size label_size;
        int baseLine = 0;
        int x, y, out_point_y, w, h;
        int line_thickness = static_cast<int>(std::round(static_cast<double>(draw_size) / 10.0));

        for (int k = 0; k < objects.size(); k++) {

            if (!objects.at(k).track_idx) continue;
            color = get_color(objects.at(k).obj.label);

            x = objects.at(k).obj.rect.x;
            y = objects.at(k).obj.rect.y;
            w = objects.at(k).obj.rect.width;
            h = objects.at(k).obj.rect.height;

            cv::rectangle(d_img,
                objects.at(k).obj.rect,
                color,
                line_thickness);

            if (draw_label) {
                txt_color = (cv::mean(color)[0] > 127) ? cv::Scalar(0, 0, 0) : cv::Scalar(255, 255, 255);
                std::string label = std::to_string(objects.at(k).track_idx) + " " + names.at(objects.at(k).obj.label) + " " + std::to_string(objects.at(k).obj.prob).substr(0, 4);
                label_size = cv::getTextSize(label.c_str(), cv::LINE_AA, static_cast<double>(draw_size) / 30.0, (line_thickness > 1) ? line_thickness - 1 : 1, &baseLine);
                txt_bk_color = color; // * 0.7;
                y = (y > d_img.rows) ? d_img.rows : y + 1;
                out_point_y = y - label_size.height - baseLine;
                if (out_point_y >= 0) y = out_point_y;
                cv::rectangle(d_img, cv::Rect(cv::Point(x - (line_thickness - 1), y), cv::Size(label_size.width, label_size.height + baseLine)),
                    txt_bk_color, -1);
                cv::putText(d_img, label, cv::Point(x, y + label_size.height),
                    cv::LINE_AA, static_cast<double>(draw_size) / 30.0, txt_color, (line_thickness > 1) ? line_thickness - 1 : 1);
            }

        }
        return d_img;
    }

}