/* 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 org.tensorflow.lite.Interpreter
import org.tensorflow.lite.examples.poseestimation.data.Person
import org.tensorflow.lite.support.common.FileUtil

// 定义PoseClassifier类，用于进行姿势分类
class PoseClassifier(
    private val interpreter: Interpreter, // TensorFlow Lite解释器实例
    private val labels: List<String> // 标签列表，用于模型输出后转换为具体的类别名称
) {
    private val input = interpreter.getInputTensor(0).shape() // 获取输入张量的形状
    private val output = interpreter.getOutputTensor(0).shape() // 获取输出张量的形状

    companion object {
        private const val MODEL_FILENAME = "classifier.tflite" // 模型文件名
        private const val LABELS_FILENAME = "labels.txt" // 标签文件名
        private const val CPU_NUM_THREADS = 4 // 设置使用的线程数

        // 创建PoseClassifier的工厂方法
        fun create(context: Context): PoseClassifier {
            // 设置解释器的选项，包括线程数
            val options = Interpreter.Options().apply {
                setNumThreads(CPU_NUM_THREADS)
            }
            // 返回一个新的PoseClassifier实例
            return PoseClassifier(
                Interpreter(
                    FileUtil.loadMappedFile(
                        context, MODEL_FILENAME // 加载模型文件
                    ), options
                ),
                FileUtil.loadLabels(context, LABELS_FILENAME) // 加载标签文件
            )
        }
    }

    // 姿势分类方法
    fun classify(person: Person?): List<Pair<String, Float>> {
        // 预处理姿势估计结果，转换为一维数组
        val inputVector = FloatArray(input[1]) // 创建输入向量
        person?.keyPoints?.forEachIndexed { index, keyPoint ->
            // 将每个关键点的坐标（y, x）和得分（score）填充到输入向量中
            inputVector[index * 3] = keyPoint.coordinate.y
            inputVector[index * 3 + 1] = keyPoint.coordinate.x
            inputVector[index * 3 + 2] = keyPoint.score
        }

        // 处理模型输出，将结果转换为可读的类别名称
        val outputTensor = FloatArray(output[1]) // 创建输出张量
        interpreter.run(arrayOf(inputVector), arrayOf(outputTensor)) // 运行模型进行推理
        val output = mutableListOf<Pair<String, Float>>() // 存储输出结果的列表
        outputTensor.forEachIndexed { index, score ->
            // 将每个输出结果与标签配对
            output.add(Pair(labels[index], score))
        }
        return output // 返回输出结果
    }

    // 关闭解释器
    fun close() {
        interpreter.close()
    }
}
