package com.example.studentfaceregistry.face

import android.graphics.Bitmap
import android.graphics.Canvas
import android.graphics.Matrix
import android.graphics.Paint
import kotlin.math.absoluteValue

/**
 * 人脸对齐器
 * 通过人脸关键点（左右眼）进行旋转对齐，提升识别准确率
 */
class FaceAligner {

    /**
     * 根据左右眼位置对齐人脸
     * @param bitmap 原始人脸图片（已裁剪）
     * @param leftEyeX 左眼 X 坐标（相对于原始图片）
     * @param leftEyeY 左眼 Y 坐标（相对于原始图片）
     * @param rightEyeX 右眼 X 坐标（相对于原始图片）
     * @param rightEyeY 右眼 Y 坐标（相对于原始图片）
     * @param faceLeft 人脸矩形左边（相对于原始图片）
     * @param faceTop 人脸矩形上边（相对于原始图片）
     * @return 对齐后的人脸图片
     */
    fun align(
        bitmap: Bitmap,
        leftEyeX: Float,
        leftEyeY: Float,
        rightEyeX: Float,
        rightEyeY: Float,
        faceLeft: Int,
        faceTop: Int
    ): Bitmap {
        // 计算眼睛在裁剪后图片中的相对位置
        var leftEyeRelX = leftEyeX - faceLeft
        var leftEyeRelY = leftEyeY - faceTop
        var rightEyeRelX = rightEyeX - faceLeft
        var rightEyeRelY = rightEyeY - faceTop
        if (leftEyeRelX > rightEyeRelX) {
            val oldLeftX = leftEyeRelX
            val oldLeftY = leftEyeRelY
            leftEyeRelX = rightEyeRelX
            leftEyeRelY = rightEyeRelY
            rightEyeRelX = oldLeftX
            rightEyeRelY = oldLeftY
        }

        // 计算两眼连线的旋转角度（目标是让连线水平）
        val angle = calculateRotationAngle(leftEyeRelX, leftEyeRelY, rightEyeRelX, rightEyeRelY)

        // 如果角度很小（小于 2 度），直接返回原图
        if (angle.absoluteValue < 2f) {
            return bitmap
        }

        // 计算旋转中心（两眼中心点）
        val centerX = bitmap.width / 2f
        val centerY = bitmap.height / 2f

        // 创建旋转矩阵
        val matrix = Matrix()
        matrix.postRotate(-angle, centerX, centerY)

        return Bitmap.createBitmap(bitmap, 0, 0, bitmap.width, bitmap.height, matrix, true)
    }

    fun alignFivePoint(
        bitmap: Bitmap,
        leftEyeX: Float,
        leftEyeY: Float,
        rightEyeX: Float,
        rightEyeY: Float,
        noseX: Float,
        noseY: Float,
        leftMouthX: Float,
        leftMouthY: Float,
        rightMouthX: Float,
        rightMouthY: Float,
        outputSize: Int = ARC_FACE_SIZE
    ): Bitmap {
        val (imageLeftEye, imageRightEye) = sortedByX(leftEyeX, leftEyeY, rightEyeX, rightEyeY)
        val (imageLeftMouth, imageRightMouth) = sortedByX(leftMouthX, leftMouthY, rightMouthX, rightMouthY)
        val src = floatArrayOf(
            imageLeftEye.x, imageLeftEye.y,
            imageRightEye.x, imageRightEye.y,
            noseX, noseY,
            imageLeftMouth.x, imageLeftMouth.y,
            imageRightMouth.x, imageRightMouth.y
        )
        val dst = arcFaceTemplate(outputSize)
        val matrix = similarityTransform(src, dst)
        val result = Bitmap.createBitmap(outputSize, outputSize, Bitmap.Config.ARGB_8888)
        Canvas(result).drawBitmap(
            bitmap,
            matrix,
            Paint(Paint.ANTI_ALIAS_FLAG or Paint.FILTER_BITMAP_FLAG)
        )
        return result
    }

    /**
     * 计算两眼连线的旋转角度
     */
    private fun calculateRotationAngle(x1: Float, y1: Float, x2: Float, y2: Float): Float {
        val dx = x2 - x1
        val dy = y2 - y1
        return Math.atan2(dy.toDouble(), dx.toDouble()).toFloat() * 180f / Math.PI.toFloat()
    }

    private fun sortedByX(x1: Float, y1: Float, x2: Float, y2: Float): Pair<Point, Point> {
        val first = Point(x1, y1)
        val second = Point(x2, y2)
        return if (x1 <= x2) first to second else second to first
    }

    private fun arcFaceTemplate(outputSize: Int): FloatArray {
        val scale = outputSize / ARC_FACE_SIZE.toFloat()
        return FloatArray(ARC_FACE_TEMPLATE.size) { index -> ARC_FACE_TEMPLATE[index] * scale }
    }

    private fun similarityTransform(src: FloatArray, dst: FloatArray): Matrix {
        val count = src.size / 2
        var srcCenterX = 0f
        var srcCenterY = 0f
        var dstCenterX = 0f
        var dstCenterY = 0f
        for (index in 0 until count) {
            srcCenterX += src[index * 2]
            srcCenterY += src[index * 2 + 1]
            dstCenterX += dst[index * 2]
            dstCenterY += dst[index * 2 + 1]
        }
        srcCenterX /= count
        srcCenterY /= count
        dstCenterX /= count
        dstCenterY /= count

        var numeratorA = 0f
        var numeratorB = 0f
        var denominator = 0f
        for (index in 0 until count) {
            val srcX = src[index * 2] - srcCenterX
            val srcY = src[index * 2 + 1] - srcCenterY
            val dstX = dst[index * 2] - dstCenterX
            val dstY = dst[index * 2 + 1] - dstCenterY
            numeratorA += srcX * dstX + srcY * dstY
            numeratorB += srcX * dstY - srcY * dstX
            denominator += srcX * srcX + srcY * srcY
        }

        val a = numeratorA / denominator.coerceAtLeast(1e-6f)
        val b = numeratorB / denominator.coerceAtLeast(1e-6f)
        val translateX = dstCenterX - a * srcCenterX + b * srcCenterY
        val translateY = dstCenterY - b * srcCenterX - a * srcCenterY

        return Matrix().apply {
            setValues(
                floatArrayOf(
                    a, -b, translateX,
                    b, a, translateY,
                    0f, 0f, 1f
                )
            )
        }
    }

    /**
     * 简单版本：仅根据估计的眼睛位置进行对齐
     * 适用于没有启用关键点检测的场景
     */
    fun alignWithEstimation(bitmap: Bitmap, faceLeft: Int, faceTop: Int, headAngle: Float? = null): Bitmap {
        // 估计眼睛位置（基于人脸矩形的经验比例）
        val faceWidth = bitmap.width
        val faceHeight = bitmap.height

        // 眼睛通常位于人脸上部约 35-40% 的位置
        val eyeY = faceHeight * 0.38f

        // 左右眼大约在人脸宽度的 35% 和 65% 位置
        val leftEyeX = faceWidth * 0.35f
        val rightEyeX = faceWidth * 0.65f

        // 如果有头部角度信息，使用它来调整
        return if (headAngle != null && headAngle.absoluteValue > 2f) {
            val centerX = faceWidth / 2f
            val centerY = faceHeight / 2f
            val matrix = Matrix()
            matrix.postRotate(-headAngle, centerX, centerY)
            Bitmap.createBitmap(bitmap, 0, 0, faceWidth, faceHeight, matrix, true)
        } else {
            align(bitmap, leftEyeX, eyeY, rightEyeX, eyeY, faceLeft, faceTop)
        }
    }

    private data class Point(val x: Float, val y: Float)

    companion object {
        private const val ARC_FACE_SIZE = 112
        private val ARC_FACE_TEMPLATE = floatArrayOf(
            38.2946f, 51.6963f,
            73.5318f, 51.5014f,
            56.0252f, 71.7366f,
            41.5493f, 92.3655f,
            70.7299f, 92.2041f
        )
    }
}
