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.Tensor
import java.io.BufferedReader
import java.io.IOException
import java.io.InputStreamReader

import java.util.regex.Pattern

class MainVm : ViewModel() {

    private var torchscriptDemo: TorchscriptDemo? = null

    private var mTokenIdMap = mutableMapOf<String, Long?>()
    private var mIdTokenMap = mutableMapOf<Long, String>()

    private val MODEL_INPUT_LENGTH = 360

    private val EXTRA_ID_NUM = 3
    private val CLS = "[CLS]"
    private val SEP = "[SEP]"
    private val PAD = "[PAD]"
    private val START_LOGITS = "start_logits"
    private val END_LOGITS = "end_logits"

    private var vocab = mutableMapOf<String, Long>()
//    private var maxLength=0


    fun initialize(context: Context) {
        torchscriptDemo = TorchscriptDemo(context)
        viewModelScope.launch {
            torchscriptDemo!!.initializeInterpreter()
        }
    }

    fun ask(str: String) {
        try {
//            var mMainVmJava = MainVmJava()
//            mMainVmJava.ask(str,torchscriptDemo)
//            var inputTensor = strToken(str)
//            var outTensors = torchscriptDemo?.ask(inputTensor)
//            val result = outTensors!!.dataAsFloatArray.get(0)
            answer(str,"")
            var b = 0
        } catch (e: Exception) {
            e.printStackTrace()
        }

    }

    fun strToken(text: String): Tensor {
//        var inTensorBuffer = Tensor.allocateLongBuffer(MODEL_INPUT_LENGTH)
        val tokens = text.split(" ")
//        val indices = LongArray(MODEL_INPUT_LENGTH)
//        for (i in 0..<Math.min(tokens.count(), MODEL_INPUT_LENGTH)) {
////            indices[i] = vocab.getOrDefault(tokens[i], 0)
//            indices[i]=
//        }
        val indices= tokens.map { it.toLong() }.toLongArray()
        val inTensor = Tensor.fromBlob(indices, longArrayOf(1, MODEL_INPUT_LENGTH.toLong()))
        return inTensor
    }

    fun readVocab(context: Context) {
        viewModelScope.launch {
            try {
                withContext(Dispatchers.IO) {
                    val br = BufferedReader(InputStreamReader(context.assets.open("vocab.txt")))
                    var line: String
                    var count = 0L
                    while (true) {
                        val line = br.readLine()
                        if (line != null) {
                            mTokenIdMap[line] = count
                            mIdTokenMap[count] = line
                            count++
                        } else break
                    }
                }
            } catch (e: IOException) {
                e.printStackTrace()
            }
        }

    }


    private fun answer(question: String, text: String): String? {

        try {
            val tokenIds = tokenizer(question, text)
            val inTensorBuffer = Tensor.allocateLongBuffer(MODEL_INPUT_LENGTH)
            for (n in tokenIds) inTensorBuffer.put(n.toLong())
            for (i in 0 until MODEL_INPUT_LENGTH - tokenIds.size) mTokenIdMap!![PAD]?.let { inTensorBuffer.put(it) }

            val inTensor = Tensor.fromBlob(inTensorBuffer, longArrayOf(1, MODEL_INPUT_LENGTH.toLong()))
//            val eValue1: EValue = EValue.from(inTensor)
//            val outTensors = torchscriptDemo!!.module!!.forward(eValue1).get(0).toTensor()

            var a=0
//            val startTensor = outTensors[START_LOGITS]!!.toTensor()
//            val endTensor = outTensors[END_LOGITS]!!.toTensor()
//
//            val starts = startTensor.dataAsFloatArray
//            val ends = endTensor.dataAsFloatArray
//            val answerTokens: MutableList<String?> = ArrayList()
//            val start = argmax(starts)
//            val end = argmax(ends)
//            for (i in start until end + 1) answerTokens.add(mIdTokenMap!![tokenIds[i]])

//            return java.lang.String.join(" ", answerTokens).replace(" ##".toRegex(), "").replace("\\s+(?=\\p{Punct})".toRegex(), "")
        } catch (e: Exception) {
            var a=0
//            runOnUiThread { mTextViewAnswer!!.text = e.message }
        }
        return null
    }



    private fun tokenizer(question: String, text: String): LongArray {
        val tokenIdsQuestion = wordPieceTokenizer(question)
        if (tokenIdsQuestion.size >= MODEL_INPUT_LENGTH) throw RuntimeException("Question too long")
        val tokenIdsText = wordPieceTokenizer(text)
        val inputLength = tokenIdsQuestion.size + tokenIdsText.size + EXTRA_ID_NUM
        val ids = LongArray(Math.min(MODEL_INPUT_LENGTH, inputLength))
        ids[0] = mTokenIdMap!![CLS]!!

        for (i in tokenIdsQuestion.indices) ids[i + 1] = tokenIdsQuestion[i]!!.toLong()
        ids[tokenIdsQuestion.size + 1] = mTokenIdMap!![SEP]!!
        val maxTextLength =
            Math.min(tokenIdsText.size, MODEL_INPUT_LENGTH - tokenIdsQuestion.size - EXTRA_ID_NUM)

        for (i in 0 until maxTextLength) {
            ids[tokenIdsQuestion.size + i + 2] = tokenIdsText[i]!!.toLong()
        }

        ids[tokenIdsQuestion.size + maxTextLength + 2] = mTokenIdMap!![SEP]!!
        return ids
    }


    private fun wordPieceTokenizer(questionOrText: String): List<Long?> {
        // for each token, if it's in the vocab.txt (a key in mTokenIdMap), return its Id
        // else do: a. find the largest sub-token (at least the first letter) that exists in vocab;
        // b. add "##" to the rest (even if the rest is a valid token) and get the largest sub-token "##..." that exists in vocab;
        // and c. repeat b.
        val tokenIds = mutableListOf<Long?>()
        val p = Pattern.compile("\\w+|\\S")
        val m = p.matcher(questionOrText)
        while (m.find()) {
            val token = m.group().lowercase()
            if (mTokenIdMap!!.containsKey(token)) tokenIds.add(mTokenIdMap!![token]) else {
                for (i in 0 until token.length) {
                    if (mTokenIdMap!!.containsKey(token.substring(0, token.length - i - 1))) {
                        tokenIds.add(mTokenIdMap!![token.substring(0, token.length - i - 1)])
                        var subToken = token.substring(token.length - i - 1)
                        var j = 0

                        while (j < subToken.length) {
                            if (mTokenIdMap!!.containsKey(
                                    "##" + subToken.substring(
                                        0,
                                        subToken.length - j
                                    )
                                )
                            ) {
                                tokenIds.add(
                                    mTokenIdMap!!["##" + subToken.substring(
                                        0,
                                        subToken.length - j
                                    )]
                                )
                                subToken = subToken.substring(subToken.length - j)
                                j = subToken.length - j
                            } else if (j == subToken.length - 1) {
                                tokenIds.add(mTokenIdMap!!["##$subToken"])
                                break
                            } else j++
                        }
                        break
                    }
                }
            }
        }
        return tokenIds
    }


    fun clear() {
        viewModelScope.launch {
            torchscriptDemo?.close()
        }
    }
}