/**
 * Copyright (c) 2022 Samsung Electronics Co., Ltd All Rights Reserved
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 * http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */

#include <algorithm>
#include <iostream>
#include <string.h>

#include "gtest/gtest.h"

#include "ImageHelper.h"
#include "mv_facial_landmark_internal.h"
#include "mv_facial_landmark.h"
#include "mv_pose_landmark_internal.h"
#include "mv_pose_landmark.h"

#define IMG_FACE TEST_RES_PATH "/res/inference/images/faceLandmark.jpg"
#define IMG_POSE TEST_RES_PATH "/res/inference/images/poseLandmark.jpg"

using namespace testing;
using namespace std;

using namespace MediaVision::Common;

struct model_info {
	string model_name;
	string model_file;
	string meta_file;
	string label_file;
};

TEST(FacialLandmarkTest, GettingAvailableInferenceEnginesInfoShouldBeOk)
{
	mv_facial_landmark_h handle;

	int ret = mv_facial_landmark_create(&handle);
	ASSERT_EQ(ret, MEDIA_VISION_ERROR_NONE);

	unsigned int engine_count = 0;

	ret = mv_facial_landmark_get_engine_count(handle, &engine_count);
	ASSERT_EQ(ret, MEDIA_VISION_ERROR_NONE);

	cout << "Engine count = " << engine_count << endl;
	ASSERT_GE(engine_count, 1);

	for (unsigned int engine_idx = 0; engine_idx < engine_count; ++engine_idx) {
		char *engine_type = nullptr;

		ret = mv_facial_landmark_get_engine_type(handle, engine_idx, &engine_type);
		ASSERT_EQ(ret, MEDIA_VISION_ERROR_NONE);

		cout << "Engine type : " << engine_type << endl;

		unsigned int device_count = 0;

		ret = mv_facial_landmark_get_device_count(handle, engine_type, &device_count);
		ASSERT_EQ(ret, MEDIA_VISION_ERROR_NONE);

		cout << "Device count = " << device_count << endl;

		ASSERT_GE(engine_count, 1);

		for (unsigned int device_idx = 0; device_idx < device_count; ++device_idx) {
			char *device_type = nullptr;

			ret = mv_facial_landmark_get_device_type(handle, engine_type, device_idx, &device_type);
			ASSERT_EQ(ret, MEDIA_VISION_ERROR_NONE);

			cout << "Device type : " << device_type << endl;
		}
	}

	ret = mv_facial_landmark_destroy(handle);
	ASSERT_EQ(ret, MEDIA_VISION_ERROR_NONE);
}

TEST(FacialLandmarkTest, InferenceShouldBeOk)
{
	mv_facial_landmark_h handle;
	vector<model_info> test_models {
		{}, // If empty then default model will be used.
		{ "FLD_TWEAK_CNN", "fld_tweakcnn_128x128.tflite", "fld_tweakcnn_128x128.json", "" }
		// TODO.
	};

	const unsigned int coordinate_answers[][5] = { { 42, 87, 63, 48, 83 }, { 32, 31, 53, 75, 76 } };

	mv_source_h mv_source = NULL;
	int ret = mv_create_source(&mv_source);
	ASSERT_EQ(ret, MEDIA_VISION_ERROR_NONE);

	ret = ImageHelper::loadImageToSource(IMG_FACE, mv_source);
	ASSERT_EQ(ret, MEDIA_VISION_ERROR_NONE);

	for (auto &model : test_models) {
		ret = mv_facial_landmark_create(&handle);
		ASSERT_EQ(ret, MEDIA_VISION_ERROR_NONE);

		ret = mv_facial_landmark_set_model(handle, model.model_name.c_str(), model.model_file.c_str(),
										   model.meta_file.c_str(), model.label_file.c_str());
		ASSERT_EQ(ret, MEDIA_VISION_ERROR_NONE);

		ret = mv_facial_landmark_set_engine(handle, "tflite", "cpu");
		ASSERT_EQ(ret, MEDIA_VISION_ERROR_NONE);

		ret = mv_facial_landmark_configure(handle);
		ASSERT_EQ(ret, MEDIA_VISION_ERROR_NONE);

		ret = mv_facial_landmark_prepare(handle);
		ASSERT_EQ(ret, MEDIA_VISION_ERROR_NONE);

		ret = mv_facial_landmark_inference(handle, mv_source);
		ASSERT_EQ(ret, 0);

		unsigned long frame_number;
		unsigned int number_of_objects;

		ret = mv_facial_landmark_get_result_count(handle, &frame_number, &number_of_objects);
		ASSERT_EQ(ret, MEDIA_VISION_ERROR_NONE);

		for (unsigned int idx = 0; idx < number_of_objects; ++idx) {
			unsigned int pos_x, pos_y;

			ret = mv_facial_landmark_get_position(handle, idx, &pos_x, &pos_y);
			ASSERT_EQ(ret, MEDIA_VISION_ERROR_NONE);

			int distance_x = pos_x - coordinate_answers[0][idx];
			int distance_y = pos_y - coordinate_answers[1][idx];

			distance_x = distance_x < 0 ? distance_x * -1 : distance_x;
			distance_y = distance_y < 0 ? distance_y * -1 : distance_y;

			ASSERT_TRUE(distance_x <= 2 && distance_y <= 2);
		}

		ret = mv_facial_landmark_destroy(handle);
		ASSERT_EQ(ret, MEDIA_VISION_ERROR_NONE);
	}

	ret = mv_destroy_source(mv_source);
	ASSERT_EQ(ret, MEDIA_VISION_ERROR_NONE);
}

TEST(PoseLandmarkTest, GettingAvailableInferenceEnginesInfoShouldBeOk)
{
	mv_pose_landmark_h handle;

	int ret = mv_pose_landmark_create(&handle);
	ASSERT_EQ(ret, MEDIA_VISION_ERROR_NONE);

	unsigned int engine_count = 0;

	ret = mv_pose_landmark_get_engine_count(handle, &engine_count);
	ASSERT_EQ(ret, MEDIA_VISION_ERROR_NONE);

	cout << "Engine count = " << engine_count << endl;
	ASSERT_GE(engine_count, 1);

	for (unsigned int engine_idx = 0; engine_idx < engine_count; ++engine_idx) {
		char *engine_type = nullptr;

		ret = mv_pose_landmark_get_engine_type(handle, engine_idx, &engine_type);
		ASSERT_EQ(ret, MEDIA_VISION_ERROR_NONE);

		cout << "Engine type : " << engine_type << endl;

		unsigned int device_count = 0;

		ret = mv_pose_landmark_get_device_count(handle, engine_type, &device_count);
		ASSERT_EQ(ret, MEDIA_VISION_ERROR_NONE);

		cout << "Device count = " << device_count << endl;

		ASSERT_GE(engine_count, 1);

		for (unsigned int device_idx = 0; device_idx < device_count; ++device_idx) {
			char *device_type = nullptr;

			ret = mv_pose_landmark_get_device_type(handle, engine_type, device_idx, &device_type);
			ASSERT_EQ(ret, MEDIA_VISION_ERROR_NONE);

			cout << "Device type : " << device_type << endl;
		}
	}

	ret = mv_pose_landmark_destroy(handle);
	ASSERT_EQ(ret, MEDIA_VISION_ERROR_NONE);
}

TEST(PoseLandmarkTest, InferenceShouldBeOk)
{
	mv_pose_landmark_h handle;
	vector<model_info> test_models {
		{}, // If empty then default model will be used.
		{ "PLD_CPM", "pld_cpm_192x192.tflite", "pld_cpm_192x192.json", "" }
		// TODO.
	};
	const unsigned int coordinate_answers[][10] = { { 300, 300, 275, 250, 325, 325, 225, 225, 350, 375 },
													{ 50, 87, 100, 137, 100, 137, 187, 250, 187, 250 } };

	mv_source_h mv_source = NULL;
	int ret = mv_create_source(&mv_source);
	ASSERT_EQ(ret, MEDIA_VISION_ERROR_NONE);

	ret = ImageHelper::loadImageToSource(IMG_POSE, mv_source);
	ASSERT_EQ(ret, MEDIA_VISION_ERROR_NONE);

	for (auto &model : test_models) {
		ret = mv_pose_landmark_create(&handle);
		ASSERT_EQ(ret, MEDIA_VISION_ERROR_NONE);

		ret = mv_pose_landmark_set_model(handle, model.model_name.c_str(), model.model_file.c_str(),
										 model.meta_file.c_str(), model.label_file.c_str());
		ASSERT_EQ(ret, MEDIA_VISION_ERROR_NONE);

		ret = mv_pose_landmark_set_engine(handle, "tflite", "cpu");
		ASSERT_EQ(ret, MEDIA_VISION_ERROR_NONE);

		ret = mv_pose_landmark_configure(handle);
		ASSERT_EQ(ret, MEDIA_VISION_ERROR_NONE);

		ret = mv_pose_landmark_prepare(handle);
		ASSERT_EQ(ret, MEDIA_VISION_ERROR_NONE);

		ret = mv_pose_landmark_inference(handle, mv_source);
		ASSERT_EQ(ret, 0);

		unsigned long frame_number;
		unsigned int number_of_objects;

		ret = mv_pose_landmark_get_result_count(handle, &frame_number, &number_of_objects);
		ASSERT_EQ(ret, MEDIA_VISION_ERROR_NONE);

		for (unsigned int idx = 0; idx < number_of_objects; ++idx) {
			unsigned int pos_x, pos_y;

			ret = mv_pose_landmark_get_position(handle, idx, &pos_x, &pos_y);
			ASSERT_EQ(ret, MEDIA_VISION_ERROR_NONE);

			int distance_x = pos_x - coordinate_answers[0][idx];
			int distance_y = pos_y - coordinate_answers[1][idx];

			distance_x = distance_x < 0 ? distance_x * -1 : distance_x;
			distance_y = distance_y < 0 ? distance_y * -1 : distance_y;

			ASSERT_TRUE(distance_x <= 3 && distance_y <= 3);
		}

		ret = mv_pose_landmark_destroy(handle);
		ASSERT_EQ(ret, MEDIA_VISION_ERROR_NONE);
	}

	ret = mv_destroy_source(mv_source);
	ASSERT_EQ(ret, MEDIA_VISION_ERROR_NONE);
}