#pragma once
#include "NCNNSTrack.h"
#include <NCNNObject.h>
#include <chrono>
namespace ByteTrackNCNN {
	class BYTETracker
	{
	public:
		BYTETracker(int frame_rate = 10, int track_buffer = 30);
		~BYTETracker();

		std::vector<STrack> update(const  std::vector<ByteTrackNCNN::Object>& objects);
		//cv::Scalar get_color(int idx);
		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;
	private:
		std::vector<STrack*> joint_stracks(std::vector<STrack*>& tlista, std::vector<STrack>& tlistb);
		std::vector<STrack> joint_stracks(std::vector<STrack>& tlista, std::vector<STrack>& tlistb);

		std::vector<STrack> sub_stracks(std::vector<STrack>& tlista, std::vector<STrack>& tlistb);
		void remove_duplicate_stracks(std::vector<STrack>& resa, std::vector<STrack>& resb, std::vector<STrack>& stracksa, std::vector<STrack>& stracksb);

		void linear_assignment(std::vector< std::vector<float> >& cost_matrix, int cost_matrix_size, int cost_matrix_size_size, float thresh,
			std::vector< std::vector<int> >& matches, std::vector<int>& unmatched_a, std::vector<int>& unmatched_b);
		std::vector< std::vector<float> > iou_distance(std::vector<STrack*>& atracks, std::vector<STrack>& btracks, int& dist_size, int& dist_size_size);
		std::vector< std::vector<float> > iou_distance(std::vector<STrack>& atracks, std::vector<STrack>& btracks);
		std::vector< std::vector<float> > ious(std::vector< std::vector<float> >& atlbrs, std::vector< std::vector<float> >& btlbrs);

		double lapjv(const  std::vector< std::vector<float> >& cost, std::vector<int>& rowsol, std::vector<int>& colsol,
			bool extend_cost = false, float cost_limit = LONG_MAX, bool return_cost = true);

		void estimateFrameRate();

	private:

		float track_thresh;
		float high_thresh;
		float match_thresh;
		int frame_id;
		int max_time_lost;
		int track_id_count;
		int track_buffer_;

		// 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_;

		std::vector<STrack> tracked_stracks;
		std::vector<STrack> lost_stracks;
		std::vector<STrack> removed_stracks;
		ByteTrackNCNN::ByteKalmanFilter kalman_filter;
	};

}