/* 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.Bitmap
import android.graphics.PointF
import android.os.SystemClock
import android.util.Log
import org.tensorflow.lite.DataType
import org.tensorflow.lite.Interpreter
import org.tensorflow.lite.examples.poseestimation.data.BodyPart
import org.tensorflow.lite.examples.poseestimation.data.Device
import org.tensorflow.lite.examples.poseestimation.data.KeyPoint
import org.tensorflow.lite.examples.poseestimation.data.Person
import org.tensorflow.lite.gpu.GpuDelegate
import org.tensorflow.lite.support.common.FileUtil
import org.tensorflow.lite.support.common.ops.NormalizeOp
import org.tensorflow.lite.support.image.ImageProcessor
import org.tensorflow.lite.support.image.TensorImage
import org.tensorflow.lite.support.image.ops.ResizeOp
import org.tensorflow.lite.support.image.ops.ResizeWithCropOrPadOp
import kotlin.math.acos
import kotlin.math.atan2
import kotlin.math.exp
import kotlin.math.sqrt

// 定义PoseNet类，继承自PoseDetector接口，负责姿势检测
class PoseNet(private val interpreter: Interpreter, private var gpuDelegate: GpuDelegate?) : PoseDetector {

    companion object {
        // 设置CPU线程数
        private const val CPU_NUM_THREADS = 4
        // 图像标准化时的均值和标准差
        private const val MEAN = 127.5f
        private const val STD = 127.5f
        private const val TAG = "Posenet"
        // 模型文件名称
        private const val MODEL_FILENAME = "posenet.tflite"

        // 创建PoseNet实例的方法，根据设备类型选择相应的计算方式（CPU、GPU或NNAPI）
        fun create(context: Context, device: Device): PoseNet {
            val options = Interpreter.Options()
            var gpuDelegate: GpuDelegate? = null
            options.setNumThreads(CPU_NUM_THREADS)
            when (device) {
                Device.CPU -> {
                }
                Device.GPU -> {
                    gpuDelegate = GpuDelegate()  // 使用GPU加速
                    options.addDelegate(gpuDelegate)
                }
                Device.NNAPI -> options.setUseNNAPI(true)  // 使用NNAPI加速
            }
            return PoseNet(
                Interpreter(
                    FileUtil.loadMappedFile(context, MODEL_FILENAME), options
                ),
                gpuDelegate
            )
        }
    }

    // 保存上一次推理的时间
    private var lastInferenceTimeNanos: Long = -1
    // 获取输入张量的宽度和高度
    private val inputWidth = interpreter.getInputTensor(0).shape()[1]
    private val inputHeight = interpreter.getInputTensor(0).shape()[2]
    private var cropHeight = 0f
    private var cropWidth = 0f
    private var cropSize = 0

    // 重写estimatePoses方法，用于估算图像中的姿势
    @Suppress("UNCHECKED_CAST")
    override fun estimatePoses(bitmap: Bitmap): List<Person> {
        // 记录估算开始时间
        val estimationStartTimeNanos = SystemClock.elapsedRealtimeNanos()
        // 处理输入图像并转换为TensorBuffer
        val inputArray = arrayOf(processInputImage(bitmap).tensorBuffer.buffer)
        Log.i(
            TAG,
            String.format(
                "Scaling to [-1,1] took %.2f ms",
                (SystemClock.elapsedRealtimeNanos() - estimationStartTimeNanos) / 1_000_000f
            )
        )

        // 初始化输出map
        val outputMap = initOutputMap(interpreter)

        // 记录推理开始时间
        val inferenceStartTimeNanos = SystemClock.elapsedRealtimeNanos()
        // 运行模型进行推理
        interpreter.runForMultipleInputsOutputs(inputArray, outputMap)
        // 记录推理时间
        lastInferenceTimeNanos = SystemClock.elapsedRealtimeNanos() - inferenceStartTimeNanos
        Log.i(
            TAG,
            String.format("Interpreter took %.2f ms", 1.0f * lastInferenceTimeNanos / 1_000_000)
        )

        // 获取热图和偏移量输出
        val heatmaps = outputMap[0] as Array<Array<Array<FloatArray>>>
        val offsets = outputMap[1] as Array<Array<Array<FloatArray>>>

        // 记录后处理开始时间
        val postProcessingStartTimeNanos = SystemClock.elapsedRealtimeNanos()
        // 进行后处理，提取关键点
        val person = postProcessModelOuputs(heatmaps, offsets)
        Log.i(
            TAG,
            String.format(
                "Postprocessing took %.2f ms",
                (SystemClock.elapsedRealtimeNanos() - postProcessingStartTimeNanos) / 1_000_000f
            )
        )

        // 返回结果，封装成一个Person对象
        return listOf(person)
    }

    /**
     * 将Posenet的输出（热图和偏移量）转换为关键点列表
     */
    //函数1: 计算两点连线倾斜程度
    fun calculateAngle(x1: Float, y1: Float, x2: Float, y2: Float): Float {
        val Y = y2 - y1
        val X = x2 - x1
        return Math.toDegrees(atan2(Y.toDouble(),X.toDouble())).toFloat()
    }
    private fun postProcessModelOuputs(
        heatmaps: Array<Array<Array<FloatArray>>>,
        offsets: Array<Array<Array<FloatArray>>>
    ): Person {
        val height = heatmaps[0].size
        val width = heatmaps[0][0].size
        val numKeypoints = heatmaps[0][0][0].size

        // 寻找每个关键点最可能出现的（行，列）位置
        val keypointPositions = Array(numKeypoints) { Pair(0, 0) }
        for (keypoint in 0 until numKeypoints) {
            var maxVal = heatmaps[0][0][0][keypoint]
            var maxRow = 0
            var maxCol = 0
            for (row in 0 until height) {
                for (col in 0 until width) {
                    if (heatmaps[0][row][col][keypoint] > maxVal) {
                        maxVal = heatmaps[0][row][col][keypoint]
                        maxRow = row
                        maxCol = col
                    }
                }
            }
            keypointPositions[keypoint] = Pair(maxRow, maxCol)
        }

        // 计算关键点的x和y坐标并进行偏移调整
        val xCoords = IntArray(numKeypoints)
        val yCoords = IntArray(numKeypoints)
        val confidenceScores = FloatArray(numKeypoints)
        keypointPositions.forEachIndexed { idx, position ->
            val positionY = keypointPositions[idx].first
            val positionX = keypointPositions[idx].second

            val inputImageCoordinateY =
                position.first / (height - 1).toFloat() * inputHeight + offsets[0][positionY][positionX][idx]
            val ratioHeight = cropSize.toFloat() / inputHeight
            val paddingHeight = cropHeight / 2
            yCoords[idx] = (inputImageCoordinateY * ratioHeight - paddingHeight).toInt()

            val inputImageCoordinateX =
                position.second / (width - 1).toFloat() * inputWidth + offsets[0][positionY][positionX][idx + numKeypoints]
            val ratioWidth = cropSize.toFloat() / inputWidth
            val paddingWidth = cropWidth / 2
            xCoords[idx] = (inputImageCoordinateX * ratioWidth - paddingWidth).toInt()

            confidenceScores[idx] = sigmoid(heatmaps[0][positionY][positionX][idx])
        }

        // 创建关键点列表并计算总分数
        val keypointList = mutableListOf<KeyPoint>()
        var totalScore = 0.0f
        enumValues<BodyPart>().forEachIndexed { idx, it ->
            keypointList.add(
                KeyPoint(
                    it,
                    PointF(xCoords[idx].toFloat(), yCoords[idx].toFloat()),
                    confidenceScores[idx]
                )
            )
            totalScore += confidenceScores[idx]
        }
        //计算并得到关键点之间的角度关系
        val keypoint = KeypointAngles()
        keypoint.csAngles(keypointList.toList())

        // 返回一个包含关键点和总分数的Person对象
        return Person(keyPoints = keypointList.toList(), score = totalScore / numKeypoints)
    }

    class KeypointAngles {
        // 用来存储角度的列表
        val angles = mutableListOf<Float>()

        // 计算三点之间的夹角
        fun abttPoints(A: KeyPoint, B: KeyPoint, C: KeyPoint): Float {
            val BAx = A.coordinate.x - B.coordinate.x
            val BAy = A.coordinate.y - B.coordinate.y
            val BCx = C.coordinate.x - B.coordinate.x
            val BCy = C.coordinate.y - B.coordinate.y

            val dotProduct = BAx * BCx + BAy * BCy
            val magBA = sqrt(BAx * BAx + BAy * BAy)
            val magBC = sqrt(BCx * BCx + BCy * BCy)

            val cosAngle = dotProduct / (magBA * magBC)
            return Math.toDegrees(acos(cosAngle.toDouble())).toFloat()
        }

        // 添加角度并控制角度列表长度
        fun addAngle(angle: Float) {
            angles.add(angle)
            if (angles.size > 8) {
                angles.removeAt(0)
            }
        }

        // 通用方法：根据三点计算并添加角度
        fun calculateAndAddAngle(A: KeyPoint, B: KeyPoint, C: KeyPoint) {
            val angle = abttPoints(A, B, C)
            addAngle(angle)
        }

        // 遍历keypointList并计算夹角
        fun csAngles(keypointList: List<KeyPoint>) {
            keypointList.forEach { keypoint ->
                when (keypoint.bodyPart) {
                    BodyPart.LEFT_SHOULDER -> {
                        val leftShoulder = keypoint
                        val leftElbow = keypointList.first { it.bodyPart == BodyPart.LEFT_ELBOW }
                        val leftWrist = keypointList.first { it.bodyPart == BodyPart.LEFT_WRIST }
                        calculateAndAddAngle(leftShoulder, leftElbow, leftWrist)

                        val rightShoulder = keypointList.first { it.bodyPart == BodyPart.RIGHT_SHOULDER }
                        val rightElbow = keypointList.first { it.bodyPart == BodyPart.RIGHT_ELBOW }
                        calculateAndAddAngle(leftShoulder, rightShoulder, rightElbow)
                    }
                    BodyPart.RIGHT_SHOULDER -> {
                        val rightShoulder = keypoint
                        val rightElbow = keypointList.first { it.bodyPart == BodyPart.RIGHT_ELBOW }
                        val rightWrist = keypointList.first { it.bodyPart == BodyPart.RIGHT_WRIST }
                        calculateAndAddAngle(rightShoulder, rightElbow, rightWrist)

                        val leftShoulder = keypointList.first { it.bodyPart == BodyPart.LEFT_SHOULDER }
                        val leftElbow = keypointList.first { it.bodyPart == BodyPart.LEFT_ELBOW }
                        calculateAndAddAngle(rightShoulder, leftShoulder, leftElbow)
                    }
                    BodyPart.LEFT_HIP -> {
                        val leftHip = keypoint
                        val rightHip = keypointList.first { it.bodyPart == BodyPart.RIGHT_HIP }
                        val leftKnee = keypointList.first { it.bodyPart == BodyPart.LEFT_KNEE }
                        calculateAndAddAngle(rightHip, leftHip, leftKnee)

                        val rightKnee = keypointList.first { it.bodyPart == BodyPart.RIGHT_KNEE }
                        calculateAndAddAngle(leftHip, rightHip, rightKnee)

                        val leftAnkle = keypointList.first { it.bodyPart == BodyPart.LEFT_ANKLE }
                        calculateAndAddAngle(leftHip, leftKnee, leftAnkle)
                    }
                    BodyPart.RIGHT_HIP -> {
                        val rightHip = keypoint
                        val leftHip = keypointList.first { it.bodyPart == BodyPart.LEFT_HIP }
                        val rightKnee = keypointList.first { it.bodyPart == BodyPart.RIGHT_KNEE }
                        calculateAndAddAngle(rightHip, leftHip, rightKnee)

                        val rightAnkle = keypointList.first { it.bodyPart == BodyPart.RIGHT_ANKLE }
                        calculateAndAddAngle(rightHip, rightKnee, rightAnkle)
                    }
                    // 这里可以继续添加其他 BodyPart 的处理分支
                    else -> {
                        // 如果遇到其他没有显式列出的bodyPart，进行处理
                        println("Unknown body part: ${keypoint.bodyPart}")
                    }
                }
            }
        }
    }

    // 返回上一次推理的时间
    override fun lastInferenceTimeNanos(): Long = lastInferenceTimeNanos

    // 关闭PoseNet模型和GPU代理
    override fun close() {
        gpuDelegate?.close()
        interpreter.close()
    }

    /**
     * 将输入的位图图像进行缩放和裁剪，转换为TensorImage
     */
    private fun processInputImage(bitmap: Bitmap): TensorImage {
        // 重置裁剪的宽度和高度
        cropWidth = 0f
        cropHeight = 0f
        cropSize = if (bitmap.width > bitmap.height) {
            cropHeight = (bitmap.width - bitmap.height).toFloat()
            bitmap.width
        } else {
            cropWidth = (bitmap.height - bitmap.width).toFloat()
            bitmap.height
        }

        // 设置图像处理器，执行缩放、裁剪和标准化操作
        val imageProcessor = ImageProcessor.Builder().apply {
            add(ResizeWithCropOrPadOp(cropSize, cropSize))
            add(ResizeOp(inputWidth, inputHeight, ResizeOp.ResizeMethod.BILINEAR))
            add(NormalizeOp(MEAN, STD))
        }.build()

        // 创建TensorImage并加载图像
        val tensorImage = TensorImage(DataType.FLOAT32)
        tensorImage.load(bitmap)
        // 处理图像并返回
        return imageProcessor.process(tensorImage)
    }
}


    /**
     * Initializes an outputMap of 1 * x * y * z FloatArrays for the model processing to populate.
     */
    private fun initOutputMap(interpreter: Interpreter): HashMap<Int, Any> {
        val outputMap = HashMap<Int, Any>()

        // 1 * 9 * 9 * 17 contains heatmaps
        val heatmapsShape = interpreter.getOutputTensor(0).shape()
        outputMap[0] = Array(heatmapsShape[0]) {
            Array(heatmapsShape[1]) {
                Array(heatmapsShape[2]) { FloatArray(heatmapsShape[3]) }
            }
        }

        // 1 * 9 * 9 * 34 contains offsets
        val offsetsShape = interpreter.getOutputTensor(1).shape()
        outputMap[1] = Array(offsetsShape[0]) {
            Array(offsetsShape[1]) { Array(offsetsShape[2]) { FloatArray(offsetsShape[3]) } }
        }

        // 1 * 9 * 9 * 32 contains forward displacements
        val displacementsFwdShape = interpreter.getOutputTensor(2).shape()
        outputMap[2] = Array(offsetsShape[0]) {
            Array(displacementsFwdShape[1]) {
                Array(displacementsFwdShape[2]) { FloatArray(displacementsFwdShape[3]) }
            }
        }

        // 1 * 9 * 9 * 32 contains backward displacements
        val displacementsBwdShape = interpreter.getOutputTensor(3).shape()
        outputMap[3] = Array(displacementsBwdShape[0]) {
            Array(displacementsBwdShape[1]) {
                Array(displacementsBwdShape[2]) { FloatArray(displacementsBwdShape[3]) }
            }
        }

        return outputMap
    }

    /** Returns value within [0,1].   */
    private fun sigmoid(x: Float): Float {
        return (1.0f / (1.0f + exp(-x)))
    }
