package com.example.livingai.data.ml

import android.content.Context
import android.graphics.Bitmap
import android.graphics.Rect
import org.tensorflow.lite.Interpreter
import org.tensorflow.lite.support.common.FileUtil
import org.tensorflow.lite.support.common.ops.NormalizeOp
import org.tensorflow.lite.support.image.ImageProcessor
import org.tensorflow.lite.support.image.TensorImage
import org.tensorflow.lite.support.image.ops.ResizeOp
import java.nio.ByteBuffer
import java.nio.ByteOrder

data class MidasDepthResult(
    val relativeDepth: Float,
    val absoluteDistanceMeters: Float?
)

class MidasDepthEstimator(private val context: Context) {

    private var interpreter: Interpreter? = null

    companion object {
        private const val MODEL_NAME = ""
        private const val INPUT_SIZE = 256

        private val NORM_MEAN = floatArrayOf(123.675f, 116.28f, 103.53f)
        private val NORM_STD  = floatArrayOf(58.395f, 57.12f, 57.375f)
    }

    init {
        setupInterpreter()
    }

    private fun setupInterpreter() {
        try {
            val files = context.assets.list("") ?: emptyArray()
            if (!files.contains(MODEL_NAME)) return

            val model = FileUtil.loadMappedFile(context, MODEL_NAME)
            interpreter = Interpreter(model, Interpreter.Options().apply { setNumThreads(4) })
        } catch (e: Exception) {
            e.printStackTrace()
        }
    }

    fun analyzeObject(
        bitmap: Bitmap,
        bbox: Rect,
        realObjectHeightMeters: Float?,
        focalLengthPixels: Float?
    ): MidasDepthResult? {
        val interp = interpreter ?: return null

        try {
            // 1. Preprocess
            var tensorImage = TensorImage(org.tensorflow.lite.DataType.FLOAT32)
            tensorImage.load(bitmap)

            val processor = ImageProcessor.Builder()
                .add(ResizeOp(INPUT_SIZE, INPUT_SIZE, ResizeOp.ResizeMethod.BILINEAR))
                .add(NormalizeOp(NORM_MEAN, NORM_STD))
                .build()

            tensorImage = processor.process(tensorImage)

            // 2. Output Buffer
            val outShape = interp.getOutputTensor(0).shape()
            val size = outShape[1] * outShape[2]
            val output = ByteBuffer.allocateDirect(size * 4).order(ByteOrder.nativeOrder())

            // 3. Run MiDaS
            interp.run(tensorImage.buffer, output)

            output.rewind()
            val depthArray = FloatArray(size)
            output.asFloatBuffer().get(depthArray)
            
            // Calculate median relative depth (inverse depth) from the BBOX region only?
            // Usually MiDaS runs on full frame.
            // If we want depth of the object, we should look at pixels corresponding to the bbox.
            // But mapping bbox to 256x256 map requires scaling.
            
            // For now, let's keep it simple: Median of WHOLE FRAME (as relative depth context)
            // OR median of the center?
            // The previous implementation used median of whole frame.
            // Let's refine it: Use median of the whole frame as 'relative depth'
            // OR if you want object depth, we need to crop.
            // Given the user wants "relative depth", median of frame is a common proxy for scene depth.
            // But "distance to object" -> usually means object depth.
            // Let's sample the center of the bbox in the depth map.
            
            // Map BBox center to 256x256
            val cx = bbox.centerX()
            val cy = bbox.centerY()
            val mapX = (cx * INPUT_SIZE) / bitmap.width
            val mapY = (cy * INPUT_SIZE) / bitmap.height
            
            // Clamp
            val safeX = mapX.coerceIn(0, INPUT_SIZE - 1)
            val safeY = mapY.coerceIn(0, INPUT_SIZE - 1)
            
            val depthIndex = safeY * INPUT_SIZE + safeX
            val objectRelativeDepth = depthArray[depthIndex] 
            // Note: MiDaS output is inverse depth (disparity). 
            // Higher value = Closer.
            
            // 4. Absolute Distance (Pinhole)
            val hPx = bbox.height().toFloat()
            val absDistance = if (realObjectHeightMeters != null && focalLengthPixels != null && hPx > 0) {
                (focalLengthPixels * realObjectHeightMeters) / hPx
            } else {
                null
            }

            return MidasDepthResult(
                relativeDepth = objectRelativeDepth,
                absoluteDistanceMeters = absDistance
            )

        } catch (e: Exception) {
            e.printStackTrace()
            return null
        }
    }
    
    // Kept for compatibility if needed, but analyzeObject is the new main entry
    fun estimateDepth(bitmap: Bitmap): Float? {
        // Fallback or simpler version
        return analyzeObject(bitmap, Rect(0,0,bitmap.width, bitmap.height), null, null)?.relativeDepth
    }
}
