package com.example.gesturerecognition

import android.content.Context
import org.tensorflow.lite.Interpreter
import java.io.FileInputStream
import java.nio.MappedByteBuffer
import java.nio.channels.FileChannel

class GestureClassifier(context: Context) {
    companion object {
        private const val MODEL_PATH = "modelo_gestos_uwave.tflite"

        // Labels dos gestos do dataset uWave (ajuste conforme seu modelo)
        private val GESTURE_LABELS = arrayOf(
            "Gesto 1", "Gesto 2", "Gesto 3", "Gesto 4",
            "Gesto 5", "Gesto 6", "Gesto 7", "Gesto 8"
        )
    }

    private var tflite: Interpreter? = null
    private var tfliteModel: MappedByteBuffer

    init {
        tfliteModel = loadModelFile(context)

        // Configurar opções do interpretador
        val options = Interpreter.Options().apply {
            setNumThreads(4) // Usar 4 threads para melhor performance
        }

        tflite = Interpreter(tfliteModel, options)

        // Log das dimensões do modelo para debug
        tflite?.let {
            val inputShape = it.getInputTensor(0).shape()
            val outputShape = it.getOutputTensor(0).shape()

            println("Input shape: ${inputShape.contentToString()}")
            println("Output shape: ${outputShape.contentToString()}")
        }
    }

    private fun loadModelFile(context: Context): MappedByteBuffer {
        val fileDescriptor = context.assets.openFd(MODEL_PATH)
        val inputStream = FileInputStream(fileDescriptor.fileDescriptor)
        val fileChannel = inputStream.channel
        val startOffset = fileDescriptor.startOffset
        val declaredLength = fileDescriptor.declaredLength
        return fileChannel.map(FileChannel.MapMode.READ_ONLY, startOffset, declaredLength)
    }

    fun classify(inputData: FloatArray): String {
        return tflite?.let { interpreter ->
            // Preparar entrada
            val input = arrayOf(inputData)

            // Preparar saída
            val output = Array(1) { FloatArray(GESTURE_LABELS.size) }

            // Executar inferência
            interpreter.run(input, output)

            // Encontrar a classe com maior probabilidade
            var predictedClass = 0
            var maxProbability = output[0][0]

            for (i in 1 until output[0].size) {
                if (output[0][i] > maxProbability) {
                    maxProbability = output[0][i]
                    predictedClass = i
                }
            }

            // Retornar resultado com confiança
            //val confidence = String.format("%.2f%%", maxProbability * 100)
            //"${GESTURE_LABELS[predictedClass]} ($confidence)"
            GESTURE_LABELS[predictedClass]
        } ?: "Erro: Modelo não carregado"
    }

    fun classifyWithAllProbabilities(inputData: FloatArray): Map<String, Float> {
        return tflite?.let { interpreter ->
            // Preparar entrada
            val input = arrayOf(inputData)

            // Preparar saída
            val output = Array(1) { FloatArray(GESTURE_LABELS.size) }

            // Executar inferência
            interpreter.run(input, output)

            // Criar mapa com todas as probabilidades
            val results = mutableMapOf<String, Float>()
            for (i in GESTURE_LABELS.indices) {
                results[GESTURE_LABELS[i]] = output[0][i]
            }

            results
        } ?: emptyMap()
    }

    fun close() {
        tflite?.close()
        tflite = null
    }
}