#pragma once
#include <chrono>
#include <stdexcept>
#include <stdlib.h>
#include <memory>
#include <vector>
#include <sstream>
#include "EigenKalmanFilter.h"
#include "EigenKalmanBBoxTrack.h"
#include "EigenBaseTrack.h"
#include "EigenBoundingBoxIoUMatching.h"
#include "EigenLinearAssignment.h"
#include "EigenBoundingBoxTrackUtils.h"

namespace ByteTrackEigen {
    /**
     * @brief BYTETracker class for tracking objects in video frames using Kalman Filters.
     *
     * This class implements the BYTETrack algorithm, which combines the use of Kalman Filters
     * and Hungarian algorithm for object tracking in video streams.
     *
     * Key functionalities include processing frame detections, updating track states, and
     * maintaining lists of active, lost, and removed tracks.
     */
    class BYTETracker {
    public:
        /**
         * @brief Constructor for BYTETracker.
         *
         * Initializes the tracker with specific thresholds and parameters for tracking.
         *
         * @param track_thresh Threshold for track confidence.
         * @param track_buffer Size of the buffer to store track history.
         * @param match_thresh Threshold for matching detections to tracks.
         * @param frame_rate Frame rate of the video being processed.
         */
        BYTETracker(float track_thresh = 0.25, int track_buffer = 30, float match_thresh = 0.8, int frame_rate = 10);
        void update_parameters(int frameRate = 10, int trackBuffer = 30, double trackThreshold = 0.5, double highThreshold = 0.6, double matchThresold = 0.8, bool autoFrameRate = false);
        float getEstimatedFps() const;
        std::vector<KalmanBBoxTrack> update(const Eigen::MatrixXf& output_results, const std::vector<std::string> obj_ids);

    private:
        // Internal constants and thresholds
        const float BASE_FRAME_RATE = 30.0;
        const float MIN_KEEP_THRESH = 0.1f;
        const float LOWER_CONFIDENCE_MATCHING_THRESHOLD = 0.5;
        const float ACTIVATION_MATCHING_THRESHOLD = 0.7;

        void estimateFrameRate();

        // Member variables for tracking settings and state
        float track_thresh; // Threshold for determining track validity.
        float match_thresh; // Threshold for matching detections to existing tracks.
        int frame_id;       // Current frame identifier.
        float det_thresh;   // Detection threshold for filtering weak detections.
        int buffer_size;    // Size of the history buffer for tracks.
        int max_time_lost;  // Maximum time a track can be lost before removal.
        int track_buffer_;  // Stored track buffer for recalculating max_time_lost.

        // Frame rate auto-estimation
        bool auto_frame_rate_;
        float estimated_fps_;
        float time_scale_factor_;
        size_t fps_sample_count_;
        std::chrono::steady_clock::time_point last_update_time_;
        bool has_last_update_time_;

        // Kalman filter and track lists
        KalmanFilter kalman_filter; // Kalman filter for state estimation.
        std::vector<std::shared_ptr<KalmanBBoxTrack>> tracked_tracks; // Active tracks.
        std::vector<std::shared_ptr<KalmanBBoxTrack>> lost_tracks;    // Tracks that are currently lost.
        std::vector<std::shared_ptr<KalmanBBoxTrack>> removed_tracks; // Tracks that are removed.

        LinearAssignment linear_assignment; // For solving assignment problems.

        // Internal methods for processing detections and tracks
        std::vector<KalmanBBoxTrack>extract_kalman_bbox_tracks(const Eigen::MatrixXf dets, const Eigen::VectorXf scores_keep, const Eigen::VectorXf clss_ids, const std::vector<std::string> obj_ids );
        Eigen::MatrixXf select_matrix_rows_by_indices(const Eigen::MatrixXf matrix, const std::vector<int> indices);
        std::vector<std::string> select_vector_by_indices(const std::vector<std::string> obj_ids, const std::vector<int> indices);

        std::pair<std::vector<KalmanBBoxTrack>, std::vector<KalmanBBoxTrack>> filter_and_partition_detections(const Eigen::MatrixXf& output_results, const std::vector<std::string> obj_ids);
        std::pair<std::vector<std::shared_ptr<KalmanBBoxTrack>>, std::vector<std::shared_ptr<KalmanBBoxTrack>>> partition_tracks_by_activation();
        std::tuple<std::vector<std::pair<int, int>>, std::set<int>, std::set<int>> assign_tracks_to_detections(
            const std::vector<std::shared_ptr<KalmanBBoxTrack>> tracks,
            const std::vector<KalmanBBoxTrack> detections,
            double thresh
        );
        void update_tracks_from_detections(
            std::vector<std::shared_ptr<KalmanBBoxTrack>>& tracks,
            const std::vector<KalmanBBoxTrack> detections,
            const std::vector<std::pair<int, int>> track_detection_pair_indices,
            std::vector<std::shared_ptr<KalmanBBoxTrack>>& reacquired_tracked_tracks,
            std::vector<std::shared_ptr<KalmanBBoxTrack>>& activated_tracks
        );
        std::vector<std::shared_ptr<KalmanBBoxTrack>> extract_active_tracks(
            const std::vector<std::shared_ptr<KalmanBBoxTrack>>& tracks,
            std::set<int> unpaired_track_indices
        );
        void flag_unpaired_tracks_as_lost(
            std::vector<std::shared_ptr<KalmanBBoxTrack>>& currently_tracked_tracks,
            std::vector<std::shared_ptr<KalmanBBoxTrack>>& lost_tracks,
            std::set<int> unpaired_track_indices
        );
        void prune_and_merge_tracked_tracks(
            std::vector<std::shared_ptr<KalmanBBoxTrack>>& reacquired_tracked_tracks,
            std::vector<std::shared_ptr<KalmanBBoxTrack>>& activated_tracks
        );
        void handle_lost_and_removed_tracks(
            std::vector<std::shared_ptr<KalmanBBoxTrack>>& removed_tracks,
            std::vector<std::shared_ptr<KalmanBBoxTrack>>& lost_tracks
        );
    };

}