import sys
sys.dont_write_bytecode = True

import re
import unicodedata
import numpy
import sentencepiece
from sentencepiece import sentencepiece_model_pb2 as sentencepieceModel

from helper import onnxSessionBuild

pathModel = "./"

glinerTypeAllowList = ["person", "organization", "place", "category", "event"]
glinerScoreMin = 0.4
glinerEntId = 250103
glinerSepId = 250104
glinerClsId = 1
glinerEosId = 2
glinerMaxWidth = 12
glinerMaxLength = 384

proto = sentencepieceModel.ModelProto()

with open(f"{pathModel}spm.model", "rb") as file:
    proto.ParseFromString(file.read())

proto.normalizer_spec.add_dummy_prefix = False

sentencepieceGliner = sentencepiece.SentencePieceProcessor()
sentencepieceGliner.LoadFromSerializedProto(proto.SerializeToString())

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

def wideCheck(character):
    return character != "" and unicodedata.east_asian_width(character) in ("W", "F")

def tokenizeWord(text):
    resultList = []

    for segment in re.findall(r"\d|\D+", text):
        for tokenId in sentencepieceGliner.encode(segment, out_type=int):
            resultList.append(tokenId)

    return resultList

def wordSplit(text):
    resultList = []

    for match in re.finditer(r"\w+(?:[-_]\w+)*|[^\w\s]", text):
        word = match.group(0)
        start = match.start()

        segment = ""
        segmentStart = start

        for a in range(len(word)):
            if wideCheck(word[a]):
                if segment != "":
                    resultList.append({"text": segment, "start": segmentStart, "end": segmentStart + len(segment)})

                    segment = ""

                resultList.append({"text": word[a], "start": start + a, "end": start + a + 1})

                segmentStart = start + a + 1
            else:
                if segment == "":
                    segmentStart = start + a

                segment += word[a]

        if segment != "":
            resultList.append({"text": segment, "start": segmentStart, "end": segmentStart + len(segment)})

    return resultList

def predict(text):
    resultList = []

    wordList = wordSplit(text)

    if len(wordList) > glinerMaxLength:
        wordList = wordList[0:glinerMaxLength]

    numWords = len(wordList)

    if numWords == 0:
        return resultList

    inputIdList = [glinerClsId]
    wordsMaskList = [0]

    for a in range(len(glinerTypeAllowList)):
        inputIdList.append(glinerEntId)
        wordsMaskList.append(0)

        for tokenId in tokenizeWord(glinerTypeAllowList[a]):
            inputIdList.append(tokenId)
            wordsMaskList.append(0)

    inputIdList.append(glinerSepId)
    wordsMaskList.append(0)

    for a in range(numWords):
        subwordList = tokenizeWord(wordList[a]["text"])

        for b in range(len(subwordList)):
            inputIdList.append(subwordList[b])
            wordsMaskList.append(a + 1 if b == 0 else 0)

    inputIdList.append(glinerEosId)
    wordsMaskList.append(0)

    inputIds = numpy.array([inputIdList], dtype=numpy.int64)
    attentionMask = numpy.ones((1, len(inputIdList)), dtype=numpy.int64)
    wordsMask = numpy.array([wordsMaskList], dtype=numpy.int64)
    textLengths = numpy.array([[numWords]], dtype=numpy.int64)

    spanIdxList = []
    spanMaskList = []

    for a in range(numWords):
        for b in range(glinerMaxWidth):
            end = a + b

            spanIdxList.append([a, end])
            spanMaskList.append(end <= numWords - 1)

    spanIdx = numpy.array([spanIdxList], dtype=numpy.int64)
    spanMask = numpy.array([spanMaskList], dtype=bool)

    feedObject = {
        "input_ids": inputIds,
        "attention_mask": attentionMask,
        "words_mask": wordsMask,
        "text_lengths": textLengths,
        "span_idx": spanIdx,
        "span_mask": spanMask
    }

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

    probability = 1.0 / (1.0 + numpy.exp(-logits[0]))

    candidateList = []

    for a in range(numWords):
        for b in range(glinerMaxWidth):
            end = a + b

            if end > numWords - 1:
                continue

            for c in range(len(glinerTypeAllowList)):
                score = float(probability[a][b][c])

                if score > glinerScoreMin:
                    candidateList.append({
                        "wordStart": a,
                        "wordEnd": end,
                        "label": glinerTypeAllowList[c],
                        "score": score
                    })

    candidateList.sort(key=lambda candidate: candidate["score"], reverse=True)

    takenList = []

    for a in range(len(candidateList)):
        candidate = candidateList[a]

        isOverlap = False

        for b in range(len(takenList)):
            taken = takenList[b]

            if candidate["wordStart"] <= taken["wordEnd"] and taken["wordStart"] <= candidate["wordEnd"]:
                isOverlap = True

                break

        if isOverlap == False:
            takenList.append(candidate)

            charStart = wordList[candidate["wordStart"]]["start"]
            charEnd = wordList[candidate["wordEnd"]]["end"]

            resultList.append({
                "start": charStart,
                "end": charEnd,
                "text": text[charStart:charEnd],
                "label": candidate["label"],
                "score": candidate["score"]
            })

    resultList.sort(key=lambda entity: entity["start"])

    return resultList

text = unicodedata.normalize("NFKC", "Linus Torvalds created Linux in Helsinki and later joined the Linux Foundation, while 田中太郎 works at 東京大学 in Tokyo.")

entityList = predict(text)

for a in range(len(entityList)):
    print(f"{entityList[a]['score']:.6f} | {entityList[a]['label']} | {entityList[a]['text']}")