#include "NCNNSTrack.h"
namespace ByteTrackNCNN {
	
	STrack::STrack(std::vector<float> tlwh_, float _score,
				   int _class_id, float _left, float _top,
				   float _right, float _bottom, std::string _object_id)
	{
		this->_tlwh.resize(4);
		this->_tlwh.assign(tlwh_.begin(), tlwh_.end());

		this->is_activated = false;
		this->track_id = 0;
		this->state = TrackState::New;
		
		this->tlwh.resize(4);
		this->tlbr.resize(4);

		this->static_tlwh();
		this->static_tlbr();
		this->frame_id = 0;
		this->tracklet_len = 0;
		this->start_frame = 0;

		this->score = _score;
		this->class_id = _class_id;
		this->left = _left;
		this->top = _top;
		this->right = _right;
		this->bottom = _bottom;
		this->object_id = _object_id;
		this->reset_count();
		this->detection_count_ = 1;
		this->class_id_locked_ = false;
		this->class_id_scores_[_class_id] = _score;
	}

	STrack::~STrack()
	{
	}

	void STrack::voteClassId(int new_class_id, float score)
	{
		if (class_id_locked_) return;

		class_id_scores_[new_class_id] += score;
		detection_count_++;

		int best_id = class_id;
		float best_score = 0.0f;
		for (const auto& entry : class_id_scores_)
		{
			if (entry.second > best_score)
			{
				best_score = entry.second;
				best_id = entry.first;
			}
		}
		class_id = best_id;

		if (detection_count_ >= CLASS_ID_LOCK_FRAMES)
		{
			class_id_locked_ = true;
		}
	}
	void STrack::reset_count() {
		count = 0;
	}

	void STrack::activate(ByteTrackNCNN::ByteKalmanFilter &kalman_filter, int frame_id)
	{
		this->kalman_filter = kalman_filter;
		this->track_id = this->next_id();

		 std::vector<float> _tlwh_tmp(4);
		_tlwh_tmp[0] = this->_tlwh[0];
		_tlwh_tmp[1] = this->_tlwh[1];
		_tlwh_tmp[2] = this->_tlwh[2];
		_tlwh_tmp[3] = this->_tlwh[3];
		 std::vector<float> xyah = tlwh_to_xyah(_tlwh_tmp);
		DETECTBOX xyah_box;
		xyah_box[0] = xyah[0];
		xyah_box[1] = xyah[1];
		xyah_box[2] = xyah[2];
		xyah_box[3] = xyah[3];
		auto mc = this->kalman_filter.initiate(xyah_box);
		this->mean = mc.first;
		this->covariance = mc.second;

		static_tlwh();
		static_tlbr();

		this->tracklet_len = 0;
		this->state = TrackState::Tracked;
		if (frame_id == 1)
		{
			this->is_activated = true;
		}
		//this->is_activated = true;
		this->frame_id = frame_id;
		this->start_frame = frame_id;
	}
	void STrack::activate(ByteTrackNCNN::ByteKalmanFilter& kalman_filter, int frame_id, int track_id_count)
	{
		this->kalman_filter = kalman_filter;
		this->track_id = track_id_count;

		std::vector<float> _tlwh_tmp(4);
		_tlwh_tmp[0] = this->_tlwh[0];
		_tlwh_tmp[1] = this->_tlwh[1];
		_tlwh_tmp[2] = this->_tlwh[2];
		_tlwh_tmp[3] = this->_tlwh[3];
		std::vector<float> xyah = tlwh_to_xyah(_tlwh_tmp);
		DETECTBOX xyah_box;
		xyah_box[0] = xyah[0];
		xyah_box[1] = xyah[1];
		xyah_box[2] = xyah[2];
		xyah_box[3] = xyah[3];
		auto mc = this->kalman_filter.initiate(xyah_box);
		this->mean = mc.first;
		this->covariance = mc.second;

		static_tlwh();
		static_tlbr();

		this->tracklet_len = 0;
		this->state = TrackState::Tracked;
		if (frame_id == 1)
		{
			this->is_activated = true;
		}
		//this->is_activated = true;
		this->frame_id = frame_id;
		this->start_frame = frame_id;
	}

	void STrack::re_activate(STrack &new_track, int frame_id, bool new_id)
	{
		 std::vector<float> xyah = tlwh_to_xyah(new_track.tlwh);
		DETECTBOX xyah_box;
		xyah_box[0] = xyah[0];
		xyah_box[1] = xyah[1];
		xyah_box[2] = xyah[2];
		xyah_box[3] = xyah[3];
		auto mc = this->kalman_filter.update(this->mean, this->covariance, xyah_box);
		this->mean = mc.first;
		this->covariance = mc.second;

		static_tlwh();
		static_tlbr();

		this->tracklet_len = 0;
		this->state = TrackState::Tracked;
		this->is_activated = true;
		this->frame_id = frame_id;
		this->score = new_track.score;
		voteClassId(new_track.class_id, new_track.score);
		this->left = new_track.left;
		this->top = new_track.top;
		this->right = new_track.right;
		this->bottom = new_track.bottom;
		this->object_id = new_track.object_id;
		if (new_id)
			this->track_id = next_id();
	}

	void STrack::update(STrack &new_track, int frame_id)
	{
		this->frame_id = frame_id;
		this->tracklet_len++;

		 std::vector<float> xyah = tlwh_to_xyah(new_track.tlwh);
		DETECTBOX xyah_box;
		xyah_box[0] = xyah[0];
		xyah_box[1] = xyah[1];
		xyah_box[2] = xyah[2];
		xyah_box[3] = xyah[3];

		auto mc = this->kalman_filter.update(this->mean, this->covariance, xyah_box);
		this->mean = mc.first;
		this->covariance = mc.second;

		static_tlwh();
		static_tlbr();

		this->state = TrackState::Tracked;
		this->is_activated = true;
		this->score = new_track.score;
		voteClassId(new_track.class_id, new_track.score);
		this->left = new_track.left;
		this->top = new_track.top;
		this->right = new_track.right;
		this->bottom = new_track.bottom;
		this->object_id = new_track.object_id;
	}

	void STrack::static_tlwh()
	{
		if (this->state == TrackState::New)
		{
			tlwh[0] = _tlwh[0];
			tlwh[1] = _tlwh[1];
			tlwh[2] = _tlwh[2];
			tlwh[3] = _tlwh[3];
			return;
		}

		tlwh[0] = mean[0];
		tlwh[1] = mean[1];
		tlwh[2] = mean[2];
		tlwh[3] = mean[3];

		tlwh[2] *= tlwh[3];
		tlwh[0] -= tlwh[2] / 2;
		tlwh[1] -= tlwh[3] / 2;
	}

	void STrack::static_tlbr()
	{
		tlbr.clear();
		tlbr.assign(tlwh.begin(), tlwh.end());
		tlbr[2] += tlbr[0];
		tlbr[3] += tlbr[1];
	}

	 std::vector<float> STrack::tlwh_to_xyah( std::vector<float> tlwh_tmp)
	{
		 std::vector<float> tlwh_output = tlwh_tmp;
		tlwh_output[0] += tlwh_output[2] / 2;
		tlwh_output[1] += tlwh_output[3] / 2;
		if (tlwh_output[3] > 0)
			tlwh_output[2] /= tlwh_output[3];
		else
			tlwh_output[2] = 0;
		return tlwh_output;
	}

	 std::vector<float> STrack::to_xyah()
	{
		return tlwh_to_xyah(tlwh);
	}

	 std::vector<float> STrack::tlbr_to_tlwh( std::vector<float> &tlbr)
	{
		tlbr[2] -= tlbr[0];
		tlbr[3] -= tlbr[1];
		return tlbr;
	}

	void STrack::mark_lost()
	{
		state = TrackState::Lost;
	}

	void STrack::mark_removed()
	{
		state = TrackState::Removed;
	}

	int STrack::next_id()
	{
		count++;
		return count;
	}

	int STrack::end_frame()
	{
		return this->frame_id;
	}

	void STrack::multi_predict( std::vector<STrack*> &stracks, ByteTrackNCNN::ByteKalmanFilter &kalman_filter)
	{
		for (int i = 0; i < stracks.size(); i++)
		{
			if (stracks[i]->state != TrackState::Tracked)
			{
				stracks[i]->mean[7] = 0;
			}
			kalman_filter.predict(stracks[i]->mean, stracks[i]->covariance);
			stracks[i]->static_tlwh();
			stracks[i]->static_tlbr();
		}
	}
}