/* Copyright 2021 The TensorFlow Authors. 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.
==============================================================================
*/

package org.tensorflow.lite.examples.poseestimation.ml

import android.content.Context
import android.graphics.PointF
import androidx.test.ext.junit.runners.AndroidJUnit4
import androidx.test.platform.app.InstrumentationRegistry
import org.junit.Before
import org.junit.Test
import org.junit.runner.RunWith
import org.tensorflow.lite.examples.poseestimation.data.BodyPart
import org.tensorflow.lite.examples.poseestimation.data.Device

// 使用 AndroidJUnit4 测试框架运行测试
@RunWith(AndroidJUnit4::class)
class MovenetLightningTest {

    // 常量定义，用于测试
    companion object {
        private const val TEST_INPUT_IMAGE1 = "image1.png" // 测试用输入图像1
        private const val TEST_INPUT_IMAGE2 = "image2.jpg" // 测试用输入图像2
        private const val ACCEPTABLE_ERROR = 21f // 可接受的误差阈值
    }

    // 用于检测姿势的类实例
    private lateinit var poseDetector: PoseDetector
    // 应用上下文
    private lateinit var appContext: Context
    // 预期的检测结果
    private lateinit var expectedDetectionResult: List<Map<BodyPart, PointF>>

    // 在每个测试执行前进行初始化
    @Before
    fun setup() {
        // 获取应用的上下文
        appContext = InstrumentationRegistry.getInstrumentation().targetContext
        // 创建 PoseDetector 实例，使用 Movenet 模型的 Lightening 版本，并指定使用 CPU 设备
        poseDetector = MoveNet.create(appContext, Device.CPU, ModelType.Lightning)
        // 加载预期的检测结果（CSV 文件中的数据）
        expectedDetectionResult =
            EvaluationUtils.loadCSVAsset("pose_landmark_truth.csv")
    }

    // 测试第一个输入图像的姿势估计结果
    @Test
    fun testPoseEstimationResultWithImage1() {
        // 加载输入图像1
        val input = EvaluationUtils.loadBitmapAssetByName(TEST_INPUT_IMAGE1)

        // Movenet 使用前一帧的信息优化检测结果，因此多次运行相同图像以改善检测结果
        poseDetector.estimatePoses(input)
        poseDetector.estimatePoses(input)
        poseDetector.estimatePoses(input)
        // 获取姿势估计的第一个人的结果
        val person = poseDetector.estimatePoses(input)[0]
        // 校验检测结果与预期值的误差是否在可接受范围内
        EvaluationUtils.assertPoseDetectionResult(
            person,
            expectedDetectionResult[0],
            ACCEPTABLE_ERROR
        )
    }

    // 测试第二个输入图像的姿势估计结果
    @Test
    fun testPoseEstimationResultWithImage2() {
        // 加载输入图像2
        val input = EvaluationUtils.loadBitmapAssetByName(TEST_INPUT_IMAGE2)

        // Movenet 使用前一帧的信息优化检测结果，因此多次运行相同图像以改善检测结果
        poseDetector.estimatePoses(input)
        poseDetector.estimatePoses(input)
        poseDetector.estimatePoses(input)
        // 获取姿势估计的第一个人的结果
        val person = poseDetector.estimatePoses(input)[0]
        // 校验检测结果与预期值的误差是否在可接受范围内
        EvaluationUtils.assertPoseDetectionResult(
            person,
            expectedDetectionResult[1],
            ACCEPTABLE_ERROR
        )
    }
}
