#include "DetUtils.h"
#include "YoloCrowdDetector.h"

#include "Logger.h"
#include "OSTime.h"
#include "SysUtils.h"

#include "ilogger.hpp"
#include "trt_infer.hpp"

#include <opencv2/opencv.hpp>

NAMESPACE_MAS_BEGIN

NAMESPACE_YOLOCROWD_BEGIN

YoloCrowdDetector::YoloCrowdDetector(const std::string& key,
	const std::string& name):
	Detector(key, name), m_yoloCrowd(nullptr)
{

}

YoloCrowdDetector::~YoloCrowdDetector()
{
	this->Dispose();
}

TZ_INT YoloCrowdDetector::Initialize(const std::string& initParam)
{
	if (m_inited)
	{
		TZLogInfo("YoloCrowdDetector has been initialized~~~");
		return MEC_OK;
	}

	m_yoloCrowd = YoloCrowdApp::Instance();

	m_yoloCrowd->Initialize(initParam);

	m_targets.target_class = 0;
	m_targets.target_threshold = 0.7;

	this->Start();

	m_inited = TRUE;

	TZLogInfo("YoloCrowdDetector initialized~~~");

	return MEC_OK;
}

TZ_INT YoloCrowdDetector::Dispose()
{
	if (!m_inited)
	{
		return MEC_NOT_INITED;
	}

	this->StopAndWait();

	m_yoloCrowd->Dispose();

	m_inited = FALSE;

	return MEC_OK;
}

TZ_INT YoloCrowdDetector::TurnOnGPU()
{
	m_useGPU = TRUE;
	return MEC_OK;
}

TZ_INT YoloCrowdDetector::TurnOffGPU()
{
	m_useGPU = FALSE;
	return MEC_OK;
}

TZ_INT YoloCrowdDetector::DoDetect(
	SPtr<StreamInfo>& media,
	fDetCallback callback, void* ctx)
{
	if (!m_inited || !callback || !media)
	{
		TZLogWarn("detector state is not active!!!");
		return MEC_FAILED;
	}

	if (m_sema.Count() > 3000)
	{
		TZLogInfo("sema is more than 3000~~~");
		return MEC_FAILED;
	}

	m_cntLock.Lock();
	++m_detCount;
	if (m_detCount != m_detSample)
	{
		m_cntLock.Unlock();
		this->skipFrame(media, callback, ctx);
		return MEC_OK;
	}

	m_detCount = 0;
	m_cntLock.Unlock();

	m_listLock.Lock();
	m_waitList.push_back({media, callback, ctx});
	if (m_waitList.size() > MAX_WAIT_LIST_LEN)
	{
		TZLogInfo("m_waitList.size = %d", m_waitList.size());
	}

	m_listLock.Unlock();

	m_sema.Signal();

	return MEC_OK;
}

TZ_INT YoloCrowdDetector::SetDetectCfg(const std::string& conf)
{
	if (!m_inited) return MEC_NOT_INITED;

	YoloCrowdCfg cfg;
	YoloCrowdCfg::fromJson(conf, cfg);

	m_detSample = detutils::GetDetSample(cfg.freq);
	m_focusArea = cfg.focusArea;
	m_ignoreArea = cfg.ignoreArea;

	TZLogInfo("target class %d, target_threshold %f", cfg.target.target_class, cfg.target.target_threshold);
	m_targets.target_class = cfg.target.target_class;
	m_targets.target_threshold = cfg.target.target_threshold;

	return MEC_OK;
}

void YoloCrowdDetector::Entry()
{
	while (!this->IsStop())
	{
		if (!m_sema.Wait(YOLO_WAIT_MSECOND)) continue;

		m_listLock.Lock();
		TZ_INT skipCnt = m_waitList.size() - MAX_WAIT_LIST_LEN;
		if (skipCnt > 0)
		{
			TZLogInfo("m_waitList.size=%d, will skip %d frames", m_waitList.size(), skipCnt);
		}
		else
		{
			skipCnt = 0;
		}
		
		m_listLock.Unlock();

		while (skipCnt > 0)
		{
			m_listLock.Lock();

			auto front = m_waitList.front();
			m_waitList.pop_front();
			m_listLock.Unlock();

			SPtr<StreamInfo> media = std::get<0>(front);
			fDetCallback callback = std::get<1>(front);
			void* ctx = std::get<2>(front);

			this->skipFrame(media, callback, ctx);

			--skipCnt;
		}

		m_listLock.Lock();
		if (m_waitList.empty())
		{
			m_listLock.Unlock();
			continue;
		}

		auto front = m_waitList.front();
		m_waitList.pop_front();
		m_listLock.Unlock();

		SPtr<StreamInfo> media = std::get<0>(front);
		fDetCallback callback = std::get<1>(front);
		void* ctx = std::get<2>(front);

		SPtr<Media> data = media->GetMediaRsc();
		if (!data || !data->Mem)
		{
			TZLogWarn("MediaRsc is nullptr!!!");
			this->skipFrame(media, callback, ctx);
			continue;
		}

		cv::Mat image(cv::Mat(data->Height, data->Width, CV_8UC3, data->Mem));

		std::vector<int> objClass;
		std::vector<std::vector<float>> objPos;

		m_yoloCrowd->DoDetect(image, m_targets, objClass, objPos);

		YoloCrowdDetectResult info;

		for (size_t i = 0; i < objClass.size(); ++i)
		{
			DetObj tmpObj;
			tmpObj.obj_class = objClass[i];
			tmpObj.left = objPos[i][0];
			tmpObj.top = objPos[i][1];
			tmpObj.right = objPos[i][2];
			tmpObj.bottom = objPos[i][3];
			tmpObj.confidence = objPos[i][4];

			if (detutils::IsInRoi(tmpObj.left / data->Width, tmpObj.top / data->Height,
								  tmpObj.right / data->Width, tmpObj.bottom / data->Height,
								  m_focusArea, m_ignoreArea, m_roiThreshold))
			{
				info.objs.push_back(tmpObj);
			}
		}

		m_rstLock.Lock();

		m_cacheJson = info.toJson();
		m_cacheDraw.clear();

		this->draw(info, m_cacheDraw, data->Height, data->Width);
		media->AddDetRst(m_key, m_cacheJson, data, m_cacheDraw);

		m_rstLock.Unlock();

		callback(media, ctx);
	}
}


void YoloCrowdDetector::skipFrame(
	SPtr<StreamInfo> & media,
	fDetCallback callback,
	void * ctx)
{
	m_rstLock.Lock();

	media->AddDetRst(m_key, m_cacheJson, media->GetMediaRsc(), m_cacheDraw);

	m_rstLock.Unlock();

	callback(media, ctx);
}

void YoloCrowdDetector::draw(
	YoloCrowdDetectResult & yoloCrowdRst, DrawInfo & draw,
	TZ_INT height, TZ_INT width)
{
	TextInfo text;

	for (auto obj : yoloCrowdRst.objs) {
		draw.Rects.emplace_back(
			obj.left / width,
			obj.top / height,
			obj.right / width,
			obj.bottom / height,
			"#00FF00",
			2,
			text
		);
	}
}

NAMESPACE_YOLOCROWD_END

NAMESPACE_MAS_END