#include <fstream>
#include <algorithm>
#include "NCNNBYTETracker.h"
#include "NCNNObject.h"
namespace ByteTrackNCNN{
	BYTETracker::BYTETracker(int frame_rate, int track_buffer)
	{
		track_thresh = 0.5;
		high_thresh = 0.01; // Track everything
		match_thresh = 0.95;
		this->frame_id = 0;
		max_time_lost = std::max(5, int(frame_rate / 30.0 * track_buffer));
		track_id_count = 0;
		track_buffer_ = track_buffer;
		auto_frame_rate_ = false;
		estimated_fps_ = static_cast<float>(frame_rate);
		time_scale_factor_ = 1.0f;
		fps_sample_count_ = 0;
		has_last_update_time_ = false;
		tracked_stracks.clear();
		lost_stracks.clear();
		removed_stracks.clear();
	}

	void BYTETracker::update_parameters(int frameRate, int trackBuffer, double trackThreshold, double highThreshold, double matchThresold, bool autoFrameRate)
	{
		track_thresh = (float)trackThreshold;
		high_thresh = (float)highThreshold;
		match_thresh = (float)matchThresold;
		track_buffer_ = trackBuffer;
		auto_frame_rate_ = autoFrameRate;
		estimated_fps_ = static_cast<float>(frameRate);
		time_scale_factor_ = 1.0f;
		fps_sample_count_ = 0;
		has_last_update_time_ = false;
		max_time_lost = std::max(5, int(frameRate / 30.0 * trackBuffer));
		this->frame_id = 0;
		track_id_count = 0;
		tracked_stracks.clear();
		lost_stracks.clear();
		removed_stracks.clear();
	}

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

	void BYTETracker::estimateFrameRate() {
		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 || delta_sec > 5.0) {
			return;
		}

		float current_fps = static_cast<float>(1.0 / delta_sec);
		current_fps = std::max(1.0f, std::min(current_fps, 120.0f));

		fps_sample_count_++;

		float alpha = (fps_sample_count_ <= 10) ? 0.3f : 0.1f;
		estimated_fps_ = alpha * current_fps + (1.0f - alpha) * estimated_fps_;

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

		if (fps_sample_count_ >= 10) {
			int new_max_time_lost = std::max(5, int(estimated_fps_ / 30.0 * track_buffer_));

			double ratio = static_cast<double>(new_max_time_lost) / static_cast<double>(max_time_lost);
			if (ratio > 1.1 || ratio < 0.9) {
				max_time_lost = new_max_time_lost;
			}
		}
	}
	BYTETracker::~BYTETracker()
	{
	}
	 std::vector<ByteTrackNCNN::STrack> BYTETracker::update(const  std::vector<ByteTrackNCNN::Object>& objects)
	{
		// Auto-estimate frame rate from update() call timing
		if (auto_frame_rate_) {
			estimateFrameRate();
		}

		////////////////// Step 1: Get detections //////////////////
		this->frame_id++;
		 std::vector<STrack> activated_stracks;
		 std::vector<STrack> refind_stracks;
		 std::vector<STrack> removed_stracks;
		 std::vector<STrack> lost_stracks;
		 std::vector<STrack> detections;
		 std::vector<STrack> detections_low;

		 std::vector<STrack> detections_cp;
		 std::vector<STrack> tracked_stracks_swap;
		 std::vector<STrack> resa, resb;
		 std::vector<STrack> output_stracks;

		 std::vector<STrack*> unconfirmed;
		 std::vector<STrack*> tracked_stracks;
		 std::vector<STrack*> strack_pool;
		 std::vector<STrack*> r_tracked_stracks;

		if (objects.size() > 0)
		{
			for (int i = 0; i < objects.size(); i++)
			{
				std::vector<float> tlbr_;
				tlbr_.resize(4);
				tlbr_[0] = objects[i].box.x();
				tlbr_[1] = objects[i].box.y();
				tlbr_[2] = objects[i].box.x() + objects[i].box.width();
				tlbr_[3] = objects[i].box.y() + objects[i].box.height();
				float score = objects[i].confidence;
				int class_id = objects[i].label;
				float left = objects[i].left;
				float top = objects[i].top;
				float right = objects[i].right;
				float bottom = objects[i].bottom;
				std::string object_id = objects[i].object_id;

				STrack strack(STrack::tlbr_to_tlwh(tlbr_), score, class_id,left,top,right,bottom, object_id);
				if (score >= track_thresh)
				{
					detections.push_back(strack);
				}
				else
				{
					detections_low.push_back(strack);
				}
			}
		}

		// Add newly detected tracklets to tracked_stracks
		for (int i = 0; i < this->tracked_stracks.size(); i++)
		{
			if (!this->tracked_stracks[i].is_activated)
				unconfirmed.push_back(&this->tracked_stracks[i]);
			else
				tracked_stracks.push_back(&this->tracked_stracks[i]);
		}

		////////////////// Step 2: First association, with IoU //////////////////
		strack_pool = joint_stracks(tracked_stracks, this->lost_stracks);

		// Multi-predict: call multi_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++) {
			STrack::multi_predict(strack_pool, this->kalman_filter);
		}

		// Adaptive matching: relax threshold during frame skips
		float effective_match_thresh = match_thresh;
		if (num_predicts > 1) {
			effective_match_thresh = std::min(match_thresh + 0.005f * (num_predicts - 1), 0.99f);
		}

		 std::vector< std::vector<float> > dists;
		int dist_size = 0, dist_size_size = 0;
		dists = iou_distance(strack_pool, detections, dist_size, dist_size_size);

		 std::vector< std::vector<int> > matches;
		 std::vector<int> u_track, u_detection;
		linear_assignment(dists, dist_size, dist_size_size, effective_match_thresh, matches, u_track, u_detection);

		for (int i = 0; i < matches.size(); i++)
		{
			STrack *track = strack_pool[matches[i][0]];
			STrack *det = &detections[matches[i][1]];
			if (track->state == TrackState::Tracked)
			{
				track->update(*det, this->frame_id);
				activated_stracks.push_back(*track);
			}
			else
			{
				track->re_activate(*det, this->frame_id, false);
				refind_stracks.push_back(*track);
			}
		}

		////////////////// Step 3: Second association, using low score dets //////////////////
		for (int i = 0; i < u_detection.size(); i++)
		{
			detections_cp.push_back(detections[u_detection[i]]);
		}
		detections.clear();
		detections.assign(detections_low.begin(), detections_low.end());
	
		for (int i = 0; i < u_track.size(); i++)
		{
			if (strack_pool[u_track[i]]->state == TrackState::Tracked)
			{
				r_tracked_stracks.push_back(strack_pool[u_track[i]]);
			}
		}

		dists.clear();
		dists = iou_distance(r_tracked_stracks, detections, dist_size, dist_size_size);

		matches.clear();
		u_track.clear();
		u_detection.clear();
		linear_assignment(dists, dist_size, dist_size_size, 0.5, matches, u_track, u_detection);

		for (int i = 0; i < matches.size(); i++)
		{
			STrack *track = r_tracked_stracks[matches[i][0]];
			STrack *det = &detections[matches[i][1]];
			if (track->state == TrackState::Tracked)
			{
				track->update(*det, this->frame_id);
				activated_stracks.push_back(*track);
			}
			else
			{
				track->re_activate(*det, this->frame_id, false);
				refind_stracks.push_back(*track);
			}
		}

		for (int i = 0; i < u_track.size(); i++)
		{
			STrack *track = r_tracked_stracks[u_track[i]];
			if (track->state != TrackState::Lost)
			{
				track->mark_lost();
				lost_stracks.push_back(*track);
			}
		}

		// Deal with unconfirmed tracks, usually tracks with only one beginning frame
		detections.clear();
		detections.assign(detections_cp.begin(), detections_cp.end());

		dists.clear();
		dists = iou_distance(unconfirmed, detections, dist_size, dist_size_size);

		matches.clear();
		std::vector<int> u_unconfirmed;
		u_detection.clear();
		linear_assignment(dists, dist_size, dist_size_size, 0.7, matches, u_unconfirmed, u_detection);

		for (int i = 0; i < matches.size(); i++)
		{
			unconfirmed[matches[i][0]]->update(detections[matches[i][1]], this->frame_id);
			activated_stracks.push_back(*unconfirmed[matches[i][0]]);
		}

		for (int i = 0; i < u_unconfirmed.size(); i++)
		{
			STrack *track = unconfirmed[u_unconfirmed[i]];
			track->mark_removed();
			removed_stracks.push_back(*track);
		}

		////////////////// Step 4: Init new stracks //////////////////
		for (int i = 0; i < u_detection.size(); i++)
		{
			STrack *track = &detections[u_detection[i]];
			if (track->score < this->high_thresh)
				continue;
			track_id_count++;
			track->activate(this->kalman_filter, this->frame_id, track_id_count);
			activated_stracks.push_back(*track);
		}

		////////////////// Step 5: Update state //////////////////
		for (int i = 0; i < this->lost_stracks.size(); i++)
		{
			if (this->frame_id - this->lost_stracks[i].end_frame() > this->max_time_lost)
			{
				this->lost_stracks[i].mark_removed();
				removed_stracks.push_back(this->lost_stracks[i]);
			}
		}
	
		for (int i = 0; i < this->tracked_stracks.size(); i++)
		{
			if (this->tracked_stracks[i].state == TrackState::Tracked)
			{
				tracked_stracks_swap.push_back(this->tracked_stracks[i]);
			}
		}
		this->tracked_stracks.clear();
		this->tracked_stracks.assign(tracked_stracks_swap.begin(), tracked_stracks_swap.end());

		this->tracked_stracks = joint_stracks(this->tracked_stracks, activated_stracks);
		this->tracked_stracks = joint_stracks(this->tracked_stracks, refind_stracks);

		//std::cout << activated_stracks.size() << std::endl;

		this->lost_stracks = sub_stracks(this->lost_stracks, this->tracked_stracks);
		for (int i = 0; i < lost_stracks.size(); i++)
		{
			this->lost_stracks.push_back(lost_stracks[i]);
		}

		this->lost_stracks = sub_stracks(this->lost_stracks, this->removed_stracks);
		for (int i = 0; i < removed_stracks.size(); i++)
		{
			this->removed_stracks.push_back(removed_stracks[i]);
		}
	
		remove_duplicate_stracks(resa, resb, this->tracked_stracks, this->lost_stracks);

		this->tracked_stracks.clear();
		this->tracked_stracks.assign(resa.begin(), resa.end());
		this->lost_stracks.clear();
		this->lost_stracks.assign(resb.begin(), resb.end());
	
		for (int i = 0; i < this->tracked_stracks.size(); i++)
		{
			output_stracks.push_back(this->tracked_stracks[i]);
	/*		if (this->tracked_stracks[i].is_activated)
			{
				output_stracks.push_back(this->tracked_stracks[i]);
			}*/
		}
		return output_stracks;
	}	
}
