package com.sd.hmct_small_model

import android.content.Context
import androidx.lifecycle.ViewModel
import androidx.lifecycle.viewModelScope
import kotlinx.coroutines.Dispatchers
import kotlinx.coroutines.launch
import kotlinx.coroutines.withContext
import org.pytorch.IValue
import org.pytorch.Module
import org.pytorch.Tensor
import java.io.File
import java.io.FileOutputStream
import java.io.IOException
import java.io.OutputStream


class PyTorchVm : ViewModel() {

    var mModule: Module? = null

    val TOP_K: Int = 3

//    private var mModuleClasses: Array<String>? = null

    fun initialize(context: Context) {
        viewModelScope.launch {
            var fPath = assetFilePath(context, "slot_model_torchscript.pt")
            if (fPath.isEmpty())
                return@launch
            try {
                mModule = Module.load(fPath)
            } catch (e: Exception) {
                e.printStackTrace()
            }
        }
    }

    suspend fun assetFilePath(context: Context, assetName: String): String {
        return withContext(Dispatchers.IO) {
            val file = File(context.filesDir, assetName)
            if (file.exists() && file.length() > 0) {
                return@withContext file.absolutePath
            }
            try {
                var inputs = context.assets.open(assetName)
                var os: OutputStream = FileOutputStream(file)
                val buffer = ByteArray(4 * 1024)
                var read: Int
                while ((inputs.read(buffer).also { read = it }) != -1) {
                    os.write(buffer, 0, read)
                }
                os.flush()
                return@withContext file.absolutePath
            } catch (e: IOException) {
                return@withContext ""
            }
        }
    }


    fun analyzeText(text: String) {
        viewModelScope.launch {
//            withContext(Dispatchers.Default) {
//                val getClassesOutput = mModule!!.runMethod("get_classes")
//
//                val classesListIValue = getClassesOutput.toList()
//                val moduleClasses = arrayOfNulls<String>(classesListIValue.size)
//
//                var i = 0
//                for (iv in classesListIValue) {
//                    moduleClasses[i++] = iv.toStr()
//                }
//            mModuleClasses = moduleClasses;

            try {
                val bytes: ByteArray = text.encodeToByteArray()
                val shape = longArrayOf(1, bytes.size.toLong())
                val inputTensor = Tensor.fromBlobUnsigned(bytes, shape)

                val outputTensor: Tensor = mModule!!.forward(IValue.from(inputTensor)).toTensor()
                val scores = outputTensor.getDataAsFloatArray()
                val ixs: IntArray? = Utils.topK(scores, TOP_K)
            } catch (e: Exception) {
                e.printStackTrace()
            }


            var a = 0
//            val topKClassNames = arrayOfNulls<String>(TOP_K)
//            val topKScores = FloatArray(TOP_K)
//            for (i in 0..<TOP_K) {
//                val ix = ixs!![i]
//                topKClassNames[i] = mModuleClasses!![ix]
//                topKScores[i] = scores[ix]
//            }
//            }
        }

    }


    fun clean() {
        mModule?.destroy()
    }


}