package com.example.studentfaceregistry.face

import android.graphics.Bitmap
import android.graphics.Rect
import com.google.mlkit.vision.common.InputImage
import com.google.mlkit.vision.face.Face
import com.google.mlkit.vision.face.FaceDetection
import com.google.mlkit.vision.face.FaceDetectorOptions
import com.google.mlkit.vision.face.FaceLandmark
import kotlinx.coroutines.tasks.await
import kotlin.math.PI
import kotlin.math.abs
import kotlin.math.atan2
import kotlin.math.cos
import kotlin.math.sin
import kotlin.math.sqrt

/**
 * 视频人脸注册器
 * 从视频中提取多角度人脸帧，融合生成更鲁棒的特征向量
 *
 * 使用场景：
 * - 学生录制视频时转头、点头
 * - 自动选择质量好的帧
 * - 融合多帧特征，提升识别准确率
 */
class VideoEnrollment(private val context: android.content.Context) : AutoCloseable {

    private val detector = FaceDetection.getClient(
        FaceDetectorOptions.Builder()
            .setPerformanceMode(FaceDetectorOptions.PERFORMANCE_MODE_ACCURATE)
            .setLandmarkMode(FaceDetectorOptions.LANDMARK_MODE_ALL)
            .setClassificationMode(FaceDetectorOptions.CLASSIFICATION_MODE_ALL)
            .build()
    )

    private val embedder = FaceEmbedder(context, FaceEmbedder.ModelType.ARCFACE)

    /**
     * 从视频帧序列中提取人脸特征
     * @param frames 视频帧列表（Bitmap）
     * @param minFrames 最少需要的人脸帧数
     * @param maxFrames 最多使用的人脸帧数（用于特征融合）
     * @return 融合后的特征向量
     */
    suspend fun enrollFromFrames(
        frames: List<Bitmap>,
        minFrames: Int = 5,
        maxFrames: Int = 20
    ): VideoEnrollmentResult {
        val faceFrames = mutableListOf<FaceFrame>()

        // 检测所有帧中的人脸
        for ((index, frame) in frames.withIndex()) {
            val faces = detector.process(InputImage.fromBitmap(frame, 0)).await()
            for (face in faces) {
                val quality = calculateFaceQuality(face)
                if (quality >= MIN_QUALITY_THRESHOLD) {
                    val yaw = calculateYawAngle(face)
                    val pitch = calculatePitchAngle(face)
                    faceFrames.add(
                        FaceFrame(
                            bitmap = frame,
                            face = face,
                            quality = quality,
                            yaw = yaw,
                            pitch = pitch,
                            frameIndex = index
                        )
                    )
                }
            }
        }

        if (faceFrames.size < minFrames) {
            return VideoEnrollmentResult.Failure(
                reason = "检测到 ${faceFrames.size} 帧有效人脸，需要至少 $minFrames 帧。请确保视频中人脸清晰且有多角度展示。"
            )
        }

        // 按角度分组，确保多角度覆盖
        val groupedFrames = groupFramesByAngle(faceFrames)

        // 从每组选择质量最好的帧
        val selectedFrames = selectBestFramesFromGroups(groupedFrames, maxFrames)

        // 提取每帧的特征
        val embeddings = selectedFrames.map { frame ->
            extractSingleEmbedding(frame)
        }

        // 融合特征（平均 + L2 归一化）
        val fusedEmbedding = fuseEmbeddings(embeddings)

        return VideoEnrollmentResult.Success(
            embedding = fusedEmbedding,
            frameCount = selectedFrames.size,
            totalDetected = faceFrames.size,
            angleCoverage = calculateAngleCoverage(selectedFrames)
        )
    }

    /**
     * 计算人脸质量分数（0-1）
     */
    private fun calculateFaceQuality(face: Face): Float {
        var score = 1.0f

        // 可见度分数
        val leftEyeVisible = face.getLandmark(FaceLandmark.LEFT_EYE) != null
        val rightEyeVisible = face.getLandmark(FaceLandmark.RIGHT_EYE) != null
        if (!leftEyeVisible || !rightEyeVisible) score *= 0.7f

        // 置信度分数
        val confidence = face.trackingId?.let { 1.0f } ?: 0.9f
        score *= confidence

        // 人脸大小分数（越大越好）
        val faceArea = face.boundingBox.width() * face.boundingBox.height()
        val sizeScore = (faceArea / 100000f).coerceIn(0.5f, 1.0f)
        score *= sizeScore

        return score
    }

    /**
     * 计算人脸偏航角（左右转头，-90 到 90 度）
     */
    private fun calculateYawAngle(face: Face): Float {
        val leftCheek = face.getLandmark(FaceLandmark.LEFT_CHEEK) ?: return 0f
        val rightCheek = face.getLandmark(FaceLandmark.RIGHT_CHEEK) ?: return 0f

        val dx = rightCheek.position.x - leftCheek.position.x
        val dy = rightCheek.position.y - leftCheek.position.y

        // 计算角度
        val angle = atan2(dy.toDouble(), dx.toDouble()).toFloat() * 180f / PI.toFloat()

        return angle.coerceIn(-90f, 90f)
    }

    /**
     * 计算人脸俯仰角（上下点头，-90 到 90 度）
     */
    private fun calculatePitchAngle(face: Face): Float {
        val nose = face.getLandmark(FaceLandmark.NOSE_BASE) ?: return 0f
        val leftEye = face.getLandmark(FaceLandmark.LEFT_EYE) ?: return 0f
        val rightEye = face.getLandmark(FaceLandmark.RIGHT_EYE) ?: return 0f

        // 眼睛中心点
        val eyeCenterY = (leftEye.position.y + rightEye.position.y) / 2f

        val dx = nose.position.x - (leftEye.position.x + rightEye.position.x) / 2f
        val dy = nose.position.y - eyeCenterY

        val angle = atan2(dy.toDouble(), dx.toDouble()).toFloat() * 180f / PI.toFloat()

        return angle.coerceIn(-90f, 90f)
    }

    /**
     * 按角度分组人脸帧
     */
    private fun groupFramesByAngle(frames: List<FaceFrame>): Map<VideoAngleBucket, List<FaceFrame>> {
        val buckets = mutableMapOf<VideoAngleBucket, MutableList<FaceFrame>>()

        for (frame in frames) {
            val bucket = getAngleBucket(frame.yaw, frame.pitch)
            if (!buckets.containsKey(bucket)) {
                buckets[bucket] = mutableListOf()
            }
            buckets[bucket]!!.add(frame)
        }

        return buckets
    }

    /**
     * 获取角度分组
     */
    private fun getAngleBucket(yaw: Float, pitch: Float): VideoAngleBucket {
        val yawBucket = when {
            yaw < -30 -> -2  // 左
            yaw < -10 -> -1  // 左前
            yaw < 10 -> 0    // 正
            yaw < 30 -> 1    // 右前
            else -> 2        // 右
        }

        val pitchBucket = when {
            pitch < -15 -> -1  // 上
            pitch < 15 -> 0    // 中
            else -> 1          // 下
        }

        return VideoAngleBucket(yawBucket, pitchBucket)
    }

    /**
     * 从每组中选择质量最好的帧
     */
    private fun selectBestFramesFromGroups(
        groups: Map<VideoAngleBucket, List<FaceFrame>>,
        maxFrames: Int
    ): List<FaceFrame> {
        val selected = mutableListOf<FaceFrame>()

        // 先按组排序（组内按质量降序）
        val sortedGroups = groups.entries.sortedByDescending { entry ->
            entry.value.maxOfOrNull { it.quality } ?: 0f
        }

        for ((_, frames) in sortedGroups) {
            val sortedFrames = frames.sortedByDescending { it.quality }
            for (frame in sortedFrames) {
                if (selected.size >= maxFrames) break
                // 避免选择同一帧多次
                if (!selected.any { it.frameIndex == frame.frameIndex }) {
                    selected.add(frame)
                }
            }
        }

        return selected
    }

    /**
     * 提取单帧人脸特征
     */
    private fun extractSingleEmbedding(frame: FaceFrame): FloatArray {
        val bitmap = frame.bitmap
        val face = frame.face

        // 裁剪人脸
        val cropped = cropFace(bitmap, face.boundingBox)

        // 对齐
        val aligned = alignFace(cropped, face)

        // 预处理
        val preprocessor = ImagePreprocessor()
        val preprocessed = preprocessor.preprocess(aligned)

        // 提取特征
        return embedder.embed(preprocessed)
    }

    /**
     * 融合多个特征向量
     */
    private fun fuseEmbeddings(embeddings: List<FloatArray>): FloatArray {
        if (embeddings.isEmpty()) throw IllegalArgumentException("embeddings is empty")
        if (embeddings.size == 1) return embeddings[0]

        val size = embeddings[0].size
        val sum = FloatArray(size)

        for (embedding in embeddings) {
            for (i in 0 until size) {
                sum[i] += embedding[i]
            }
        }

        // 平均
        for (i in 0 until size) {
            sum[i] /= embeddings.size
        }

        // L2 归一化
        return l2Normalize(sum)
    }

    /**
     * 计算角度覆盖范围
     */
    private fun calculateAngleCoverage(frames: List<FaceFrame>): AngleCoverage {
        var minYaw = Float.MAX_VALUE
        var maxYaw = -Float.MAX_VALUE
        var minPitch = Float.MAX_VALUE
        var maxPitch = -Float.MAX_VALUE

        for (frame in frames) {
            minYaw = minOf(minYaw, frame.yaw)
            maxYaw = maxOf(maxYaw, frame.yaw)
            minPitch = minOf(minPitch, frame.pitch)
            maxPitch = maxOf(maxPitch, frame.pitch)
        }

        return AngleCoverage(
            yawRange = maxYaw - minYaw,
            pitchRange = maxPitch - minPitch,
            minYaw = minYaw,
            maxYaw = maxYaw,
            minPitch = minPitch,
            maxPitch = maxPitch
        )
    }

    /**
     * 裁剪人脸
     */
    private fun cropFace(bitmap: Bitmap, box: Rect): Bitmap {
        val padding = (maxOf(box.width(), box.height()) * 0.25f).toInt()
        val left = (box.left - padding).coerceIn(0, bitmap.width)
        val top = (box.top - padding).coerceIn(0, bitmap.height)
        val right = (box.right + padding).coerceIn(0, bitmap.width)
        val bottom = (box.bottom + padding).coerceIn(0, bitmap.height)
        return Bitmap.createBitmap(bitmap, left, top, right - left, bottom - top)
    }

    /**
     * 对齐人脸
     */
    private fun alignFace(bitmap: Bitmap, face: Face): Bitmap {
        val leftEye = face.getLandmark(FaceLandmark.LEFT_EYE) ?: return bitmap
        val rightEye = face.getLandmark(FaceLandmark.RIGHT_EYE) ?: return bitmap

        val aligner = FaceAligner()
        return aligner.align(
            bitmap = bitmap,
            leftEyeX = leftEye.position.x,
            leftEyeY = leftEye.position.y,
            rightEyeX = rightEye.position.x,
            rightEyeY = rightEye.position.y,
            faceLeft = 0,
            faceTop = 0
        )
    }

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

    override fun close() {
        detector.close()
        embedder.close()
    }

    companion object {
        private const val MIN_QUALITY_THRESHOLD = 0.5f
    }
}

/**
 * 人脸帧信息
 */
data class FaceFrame(
    val bitmap: Bitmap,
    val face: Face,
    val quality: Float,
    val yaw: Float,      // 偏航角（左右转头）
    val pitch: Float,    // 俯仰角（上下点头）
    val frameIndex: Int
)

/**
 * 角度分组
 */
data class VideoAngleBucket(val yaw: Int, val pitch: Int)

/**
 * 角度覆盖范围
 */
data class AngleCoverage(
    val yawRange: Float,
    val pitchRange: Float,
    val minYaw: Float,
    val maxYaw: Float,
    val minPitch: Float,
    val maxPitch: Float
)

/**
 * 注册结果
 */
sealed class VideoEnrollmentResult {
    data class Success(
        val embedding: FloatArray,
        val frameCount: Int,
        val totalDetected: Int,
        val angleCoverage: AngleCoverage
    ) : VideoEnrollmentResult()

    data class Failure(val reason: String) : VideoEnrollmentResult()
}
