package com.example.livingai.domain.ml

import android.graphics.Bitmap
import android.graphics.Rect
import kotlin.math.max
import kotlin.math.min

class OrientationPixelEstimator(
    private val iouThreshold: Float = 0.60f
) {

    /**
     * Main function:
     *  - segmentationMaskBitmap: MLKit’s alpha mask (animal foreground)
     *  - silhouetteBitmap: template mask for EXPECTED orientation (e.g., LEFT)
     *  - bbox: detected bounding box from segmentation
     */
    fun analyze(
        segmentationMaskBitmap: Bitmap,
        silhouetteBitmap: Bitmap,
        bbox: Rect,
        frameWidth: Int,
        frameHeight: Int,
        medianDepthMeters: Float? = null
    ): OrientationPixelResult {

        // 1) Convert both masks → boolean
        val segFullMask = bitmapToBooleanMask(segmentationMaskBitmap)
        val silhouetteMask = bitmapToBooleanMask(silhouetteBitmap)

        // 2) Crop segmentation mask to bbox
        val croppedMask = cropMaskToBBox(segFullMask, frameWidth, frameHeight, bbox)

        // 3) Scale silhouette mask to bbox size
        val scaledSilhouette = scaleMask(
            silhouetteMask,
            silhouetteBitmap.width,
            silhouetteBitmap.height,
            bbox.width(),
            bbox.height()
        )

        // 4) Compute IoU
        val iou = computeIoU(croppedMask, scaledSilhouette)
        val orientationMatched = iou >= iouThreshold

        // 5) Pixel metrics extraction
        val metrics = computePixelMetrics(croppedMask, bbox, medianDepthMeters)

        return OrientationPixelResult(
            orientationMatched = orientationMatched,
            matchedOrientation = null,
            iouScore = iou,
            iouBestOther = 0f,
            pixelMetrics = metrics
        )
    }

    // -----------------------------
    // MASK HELPERS
    // -----------------------------

    private fun bitmapToBooleanMask(bitmap: Bitmap): BooleanArray {
        val w = bitmap.width
        val h = bitmap.height
        val pixels = IntArray(w * h)
        bitmap.getPixels(pixels, 0, w, 0, 0, w, h)

        val out = BooleanArray(w * h)
        for (i in pixels.indices) {
            val alpha = (pixels[i] ushr 24) and 0xFF
            out[i] = alpha > 0
        }
        return out
    }

    private fun cropMaskToBBox(
        fullMask: BooleanArray,
        frameW: Int,
        frameH: Int,
        bbox: Rect
    ): BooleanArray {

        val left = max(0, bbox.left)
        val top = max(0, bbox.top)
        val right = min(frameW - 1, bbox.right)
        val bottom = min(frameH - 1, bbox.bottom)

        val width = right - left + 1
        val height = bottom - top + 1

        val out = BooleanArray(width * height)
        var idx = 0

        for (y in top..bottom) {
            for (x in left..right) {
                out[idx++] = fullMask[y * frameW + x]
            }
        }

        return out
    }

    private fun scaleMask(
        src: BooleanArray,
        srcW: Int,
        srcH: Int,
        dstW: Int,
        dstH: Int
    ): BooleanArray {

        val out = BooleanArray(dstW * dstH)

        for (y in 0 until dstH) {
            val sy = ((y.toFloat() / dstH) * srcH).toInt().coerceIn(0, srcH - 1)
            for (x in 0 until dstW) {
                val sx = ((x.toFloat() / dstW) * srcW).toInt().coerceIn(0, srcW - 1)
                out[y * dstW + x] = src[sy * srcW + sx]
            }
        }

        return out
    }

    private fun computeIoU(a: BooleanArray, b: BooleanArray): Float {
        if (a.size != b.size) return 0f

        var inter = 0
        var union = 0

        for (i in a.indices) {
            val ai = a[i]
            val bi = b[i]
            if (ai || bi) union++
            if (ai && bi) inter++
        }

        return if (union == 0) 0f else inter.toFloat() / union
    }

    // -----------------------------
    // PIXEL METRICS
    // -----------------------------

    private fun computePixelMetrics(
        croppedMask: BooleanArray,
        bbox: Rect,
        medianDepthMeters: Float?
    ): PixelMetrics {

        val w = bbox.width()
        val h = bbox.height()

        var count = 0
        var sumX = 0L
        var sumY = 0L

        for (y in 0 until h) {
            for (x in 0 until w) {
                if (croppedMask[y * w + x]) {
                    count++
                    sumX += x
                    sumY += y
                }
            }
        }

        val centroidX = bbox.left + (sumX.toFloat() / max(1, count))
        val centroidY = bbox.top + (sumY.toFloat() / max(1, count))

        return PixelMetrics(
            widthPx = w,
            heightPx = h,
            areaPx = count,
            centroidX = centroidX,
            centroidY = centroidY,
            distanceProxyInvHeight = if (h > 0) 1f / h.toFloat() else Float.POSITIVE_INFINITY,
            heightPxFloat = h.toFloat(),
            medianDepthMeters = medianDepthMeters
        )
    }
}
