import sys
sys.dont_write_bytecode = True

import unicodedata
import numpy
import sentencepiece

from helper import onnxSessionBuild

pathModel = "./"

embeddingTokenMax = 2048

sentencepieceEmbedding = sentencepiece.SentencePieceProcessor()
sentencepieceEmbedding.Load(f"{pathModel}tokenizer.model")

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

def embedding(mode, text):
    inputList = text if isinstance(text, list) else [text]

    inputPrefixList = []

    for a in range(len(inputList)):
        if mode == "document":
            inputPrefixList.append(f"title: none | text: {inputList[a]}")
        else:
            inputPrefixList.append(f"task: search result | query: {inputList[a]}")

    tokenList = []
    lengthMax = 0

    for a in range(len(inputPrefixList)):
        idList = sentencepieceEmbedding.EncodeAsIds(inputPrefixList[a])

        if len(idList) > embeddingTokenMax - 2:
            idList = idList[0:embeddingTokenMax - 2]

        idList = [sentencepieceEmbedding.bos_id()] + idList + [sentencepieceEmbedding.eos_id()]

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

        tokenList.append(idList)

    inputIds = numpy.full((len(tokenList), lengthMax), sentencepieceEmbedding.pad_id(), dtype=numpy.int64)
    attentionMask = numpy.zeros((len(tokenList), lengthMax), dtype=numpy.int64)

    for a in range(len(tokenList)):
        inputIds[a, 0:len(tokenList[a])] = tokenList[a]
        attentionMask[a, 0:len(tokenList[a])] = 1

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

    return onnxSessionEmbedding.run(["sentence_embedding"], feedObject)[0]

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])

promptVector = embedding("query", prompt)[0]
vectorList = embedding("document", textList)

for a in range(len(textList)):
    score = float(numpy.dot(promptVector, vectorList[a]) / (numpy.linalg.norm(promptVector) * numpy.linalg.norm(vectorList[a])))

    print(f"{score:.6f} | {textList[a]}")