package com.example.studentfaceregistry.face

import android.content.Context
import android.graphics.Bitmap
import android.util.Log
import com.google.ai.edge.litert.Accelerator
import com.google.ai.edge.litert.CompiledModel
import com.google.ai.edge.litert.TensorBuffer
import kotlin.math.sqrt

/**
 * 人脸特征提取器
 * 支持 FaceNet 和 ArcFace 模型
 *
 * ArcFace 模型推荐:
 * - 下载地址：https://github.com/onnela亭亭/ArcFace_TFLite
 * - 或使用 InsightFace 的 ArcFace 模型
 */
class FaceEmbedder(context: Context, private val modelType: ModelType = ModelType.ARCFACE) : AutoCloseable {
    enum class ModelType {
        FACENET,   // 归一化：(pixel - 127.5) / 127.5
        ARCFACE    // 归一化：pixel / 255.0，然后 (x - mean) / std
    }

    private val model: CompiledModel
    private val inputBuffer: TensorBuffer
    private val outputBuffer: TensorBuffer
    private val inputSize: Int
    private val outputSize: Int
    private val accelerator: Accelerator
    private val inputPixels: IntArray
    private val inputValues: FloatArray

    // ArcFace 标准化参数 (RGB)
    private val mean = floatArrayOf(127.5f, 127.5f, 127.5f)
    private val std = floatArrayOf(127.5f, 127.5f, 127.5f)

    init {
        val (compiledModel, selectedAccelerator) = createModel(context)
        model = compiledModel
        accelerator = selectedAccelerator
        inputBuffer = model.createInputBuffers().first()
        outputBuffer = model.createOutputBuffers().first()
        val inputElements = inputBuffer.readFloat().size
        inputSize = inferInputSize(inputElements)
        outputSize = outputBuffer.readFloat().size
        inputPixels = IntArray(inputSize * inputSize)
        inputValues = FloatArray(inputSize * inputSize * 3)
        Log.i(TAG, "FaceEmbedder initialized with accelerator=$accelerator")
    }

    /**
     * 提取人脸特征
     * @param face 人脸图片（已裁剪和对齐）
     * @return L2 归一化的特征向量
     */
    fun embed(face: Bitmap): FloatArray {
        // 缩放到模型输入尺寸
        val resized = Bitmap.createScaledBitmap(face, inputSize, inputSize, true)

        // 准备输入数据
        resized.getPixels(inputPixels, 0, inputSize, 0, 0, inputSize, inputSize)

        var index = 0
        when (modelType) {
            ModelType.FACENET -> {
                for (color in inputPixels) {
                    val r = ((color shr 16) and 0xFF).toFloat()
                    val g = ((color shr 8) and 0xFF).toFloat()
                    val b = (color and 0xFF).toFloat()
                    inputValues[index++] = (r - 127.5f) / 127.5f
                    inputValues[index++] = (g - 127.5f) / 127.5f
                    inputValues[index++] = (b - 127.5f) / 127.5f
                }
            }
            ModelType.ARCFACE -> {
                for (color in inputPixels) {
                    val r = ((color shr 16) and 0xFF).toFloat()
                    val g = ((color shr 8) and 0xFF).toFloat()
                    val b = (color and 0xFF).toFloat()
                    inputValues[index++] = (r - mean[0]) / std[0]
                    inputValues[index++] = (g - mean[1]) / std[1]
                    inputValues[index++] = (b - mean[2]) / std[2]
                }
            }
        }

        // 执行推理
        inputBuffer.writeFloat(inputValues)
        model.run(listOf(inputBuffer), listOf(outputBuffer))

        // L2 归一化输出特征
        return l2Normalize(outputBuffer.readFloat())
    }

    fun warmUp() {
        val bitmap = Bitmap.createBitmap(inputSize, inputSize, Bitmap.Config.ARGB_8888)
        embed(bitmap)
    }

    override fun close() {
        model.close()
    }

    /**
     * L2 归一化
     */
    private fun l2Normalize(values: FloatArray): FloatArray {
        var sum = 0f
        for (value in values) {
            sum += value * value
        }
        val norm = sqrt(sum.coerceAtLeast(1e-12f))
        val result = FloatArray(values.size)
        for (i in values.indices) {
            result[i] = values[i] / norm
        }
        return result
    }

    /**
     * 获取特征向量维度
     */
    fun getOutputSize(): Int = outputSize

    companion object {
        private const val TAG = "FaceEmbedder"
        const val MODEL_ASSET_NAME = "facenet.tflite"
    }

    private fun inferInputSize(floatCount: Int): Int {
        val pixelCount = floatCount / 3
        val size = kotlin.math.sqrt(pixelCount.toDouble()).toInt()
        require(size > 0 && size * size * 3 == floatCount) {
            "Unexpected face model input size: $floatCount floats"
        }
        return size
    }

    private fun createModel(context: Context): Pair<CompiledModel, Accelerator> {
        return runCatching {
            CompiledModel.create(
                context.assets,
                MODEL_ASSET_NAME,
                CompiledModel.Options(setOf(Accelerator.GPU))
            ) to Accelerator.GPU
        }.getOrElse { gpuError ->
            Log.w(TAG, "GPU initialization failed, falling back to CPU", gpuError)
            CompiledModel.create(
                context.assets,
                MODEL_ASSET_NAME,
                CompiledModel.Options(setOf(Accelerator.CPU))
            ) to Accelerator.CPU
        }
    }
}
