import sys
sys.dont_write_bytecode = True

import unicodedata
import numpy
import sentencepiece

from helper import onnxSessionBuild

pathModel = "./"

rerankerBatchLength = 8
rerankerTokenMax = 512
rerankerPromptTokenMax = 128
rerankerBosId = 0
rerankerPadId = 1
rerankerEosId = 2
rerankerUnkId = 3
rerankerOffset = 1

sentencepieceReranker = sentencepiece.SentencePieceProcessor()
sentencepieceReranker.Load(f"{pathModel}sentencepiece.bpe.model")

onnxSessionReranker = onnxSessionBuild(f"{pathModel}onnx/model.onnx")

def rerankTokenize(text):
    resultList = []

    for spmId in sentencepieceReranker.encode(text, out_type=int):
        if spmId == 0:
            resultList.append(rerankerUnkId)
        else:
            resultList.append(spmId + rerankerOffset)

    return resultList

def rerank(prompt, textList):
    scoreList = []

    promptIdList = rerankTokenize(prompt)

    if len(promptIdList) > rerankerPromptTokenMax:
        promptIdList = promptIdList[0:rerankerPromptTokenMax]

    for a in range(0, len(textList), rerankerBatchLength):
        batchList = textList[a:a + rerankerBatchLength]

        sequenceList = []
        lengthMax = 0

        for b in range(len(batchList)):
            textIdList = rerankTokenize(batchList[b])

            lengthText = rerankerTokenMax - len(promptIdList) - 4

            if len(textIdList) > lengthText:
                textIdList = textIdList[0:lengthText]

            idList = [rerankerBosId] + promptIdList + [rerankerEosId, rerankerEosId] + textIdList + [rerankerEosId]

            if len(idList) > lengthMax:
                lengthMax = len(idList)

            sequenceList.append(idList)

        inputIds = numpy.full((len(sequenceList), lengthMax), rerankerPadId, dtype=numpy.int64)
        attentionMask = numpy.zeros((len(sequenceList), lengthMax), dtype=numpy.int64)

        for b in range(len(sequenceList)):
            inputIds[b, 0:len(sequenceList[b])] = sequenceList[b]
            attentionMask[b, 0:len(sequenceList[b])] = 1

        feedObject = {"input_ids": inputIds, "attention_mask": attentionMask}

        logits = onnxSessionReranker.run(["logits"], feedObject)[0]

        for b in range(len(logits)):
            scoreList.append(float(1.0 / (1.0 + numpy.exp(-logits[b][0]))))

    return scoreList

prompt = unicodedata.normalize("NFKC", "what is panda?")

textList = [
    "The giant panda (Ailuropoda melanoleuca), sometimes called a panda bear, is a bear species endemic to China.",
    "hi",
    "パンダはクマ科の哺乳類で、中国の固有種である。"
]

for a in range(len(textList)):
    textList[a] = unicodedata.normalize("NFKC", textList[a])

scoreList = rerank(prompt, textList)

for a in range(len(textList)):
    print(f"{scoreList[a]:.6f} | {textList[a]}")