package com.example.studentfaceregistry.face

import android.graphics.Bitmap
import android.graphics.Canvas
import android.graphics.ColorMatrix
import android.graphics.ColorMatrixColorFilter
import android.graphics.Paint
import kotlin.math.abs
import kotlin.math.max
import kotlin.math.min

/**
 * 图像预处理器
 * 提供多种图像增强功能，提升人脸识别准确率
 */
class ImagePreprocessor {

    /**
     * 统一的人脸预处理流程。
     *
     * 识别和注册都走同一套轻量增强，减少两条链路的特征分布差异。
     */
    fun preprocess(bitmap: Bitmap): Bitmap {
        return normalizeForEmbedding(bitmap)
    }

    /**
     * 实时识别场景使用同一套轻量预处理，避免额外的帧间波动。
     */
    fun preprocessRealtime(bitmap: Bitmap): Bitmap {
        return normalizeForEmbedding(bitmap)
    }

    private fun normalizeForEmbedding(bitmap: Bitmap): Bitmap {
        return enhanceBrightnessContrast(bitmap, brightness = 0.06f, contrast = 1.08f)
    }

    /**
     * 亮度对比度增强
     * @param brightness 亮度调整 (-1 到 1，0 为不调整)
     * @param contrast 对比度调整 (0 到 2，1 为不调整)
     */
    fun enhanceBrightnessContrast(
        bitmap: Bitmap,
        brightness: Float = 0.1f,
        contrast: Float = 1.15f
    ): Bitmap {
        val colorMatrix = ColorMatrix().apply {
            set(floatArrayOf(
                contrast, 0f, 0f, 0f, brightness * 255f,
                0f, contrast, 0f, 0f, brightness * 255f,
                0f, 0f, contrast, 0f, brightness * 255f,
                0f, 0f, 0f, 1f, 0f
            ))
        }

        val result = Bitmap.createBitmap(bitmap.width, bitmap.height, bitmap.config ?: Bitmap.Config.ARGB_8888)
        val canvas = Canvas(result)
        val paint = Paint().apply {
            colorFilter = ColorMatrixColorFilter(colorMatrix)
        }
        canvas.drawBitmap(bitmap, 0f, 0f, paint)
        return result
    }

    /**
     * 直方图均衡化（简化版，对 RGB 各通道分别处理）
     */
    fun histogramEqualize(bitmap: Bitmap): Bitmap {
        val width = bitmap.width
        val height = bitmap.height
        val pixels = IntArray(width * height)
        bitmap.getPixels(pixels, 0, width, 0, 0, width, height)

        // 对每个通道进行均衡化
        val lookupR = buildHistogramLookup(pixels, { color -> (color shr 16) and 0xFF })
        val lookupG = buildHistogramLookup(pixels, { color -> (color shr 8) and 0xFF })
        val lookupB = buildHistogramLookup(pixels, { color -> color and 0xFF })

        // 应用查找表
        for (i in pixels.indices) {
            val color = pixels[i]
            val r = lookupR[(color shr 16) and 0xFF]
            val g = lookupG[(color shr 8) and 0xFF]
            val b = lookupB[color and 0xFF]
            pixels[i] = -0x1000000 or (r shl 16) or (g shl 8) or b
        }

        val result = Bitmap.createBitmap(width, height, bitmap.config ?: Bitmap.Config.ARGB_8888)
        result.setPixels(pixels, 0, width, 0, 0, width, height)
        return result
    }

    /**
     * 构建直方图均衡化的查找表
     */
    private fun buildHistogramLookup(pixels: IntArray, channelExtractor: (Int) -> Int): IntArray {
        val histogram = IntArray(256)
        val total = pixels.size

        // 统计直方图
        for (pixel in pixels) {
            histogram[channelExtractor(pixel)]++
        }

        // 计算累积分布函数
        val cumulative = IntArray(256)
        var sum = 0
        for (i in 0..255) {
            sum += histogram[i]
            cumulative[i] = sum
        }

        // 构建查找表
        val lookup = IntArray(256)
        val minNonZero = (0..255).firstOrNull { histogram[it] > 0 } ?: 0
        val minCumulative = cumulative[minNonZero]

        for (i in 0..255) {
            lookup[i] = (((cumulative[i] - minCumulative) * 255f) / (total - minCumulative)).toInt()
                .coerceIn(0, 255)
        }

        return lookup
    }

    /**
     * 自适应亮度调整：根据图片平均亮度自动调整
     */
    fun adaptiveBrightness(bitmap: Bitmap): Bitmap {
        val avgBrightness = calculateAverageBrightness(bitmap)

        // 目标亮度约为 128（中间值）
        val adjustment = (128f - avgBrightness) / 255f * 0.5f

        return if (abs(adjustment) > 0.02f) {
            enhanceBrightnessContrast(bitmap, brightness = adjustment, contrast = 1f)
        } else {
            bitmap
        }
    }

    /**
     * 计算图片平均亮度
     */
    private fun calculateAverageBrightness(bitmap: Bitmap): Float {
        val width = bitmap.width
        val height = bitmap.height
        val pixels = IntArray(width * height)
        bitmap.getPixels(pixels, 0, width, 0, 0, width, height)

        var totalBrightness = 0f
        for (pixel in pixels) {
            // 使用人眼对 RGB 的敏感度权重计算亮度
            val r = (pixel shr 16) and 0xFF
            val g = (pixel shr 8) and 0xFF
            val b = pixel and 0xFF
            totalBrightness += 0.299f * r + 0.587f * g + 0.114f * b
        }
        return totalBrightness / pixels.size
    }

    /**
     * 边缘柔和裁剪：在裁剪区域边缘添加渐变遮罩
     */
    fun applySoftEdge(bitmap: Bitmap, edgeWidth: Int? = null): Bitmap {
        val width = bitmap.width
        val height = bitmap.height
        val edge = edgeWidth ?: min(width, height) / 8

        val result = Bitmap.createBitmap(width, height, Bitmap.Config.ARGB_8888)
        val pixels = IntArray(width * height)
        bitmap.getPixels(pixels, 0, width, 0, 0, width, height)

        for (y in 0 until height) {
            for (x in 0 until width) {
                val i = y * width + x
                val pixel = pixels[i]

                // 计算到边缘的距离
                val distToLeft = x
                val distToRight = width - 1 - x
                val distToTop = y
                val distToBottom = height - 1 - y
                val minDist = minOf(distToLeft, distToRight, distToTop, distToBottom)

                if (minDist < edge) {
                    // 应用渐变透明度
                    val alpha = (minDist.toFloat() / edge * 0.7f + 0.3f * 255).toInt()
                    pixels[i] = (pixel and 0x00FFFFFF) or (alpha shl 24)
                }
            }
        }

        result.setPixels(pixels, 0, width, 0, 0, width, height)
        return result
    }
}
