from model.request_models import InputElement
from model.response_models import EntityNerResponseModel, ResponseNerModel

from typing import List, Dict

from flair.data import Sentence
from flair.models import SequenceTagger

from pathlib import Path

from src.helpers import DownloadModel


def sentence_to_json(sentence: Sentence, tag: str = 'ner') -> List[EntityNerResponseModel]:
    result = []
    print("Values : ", sentence.to_dict(tag))
    values = sentence.to_dict(tag)['entities']
    for value in values:
        entity_response = EntityNerResponseModel(
            entity_group=value['labels'][0].value,
            word=value['text'],
            start=int(value['start_pos']),
            end=int(value['end_pos']),
            score=float(value['labels'][0].score))
        result.append(entity_response)
    return result


class ResponsesNer:
    responses: Dict[int, ResponseNerModel] = {}

    @staticmethod
    def check_id(request_id: int) -> bool:
        if ResponsesNer.responses.get(request_id):
            return True
        return False

    @staticmethod
    def get_response(request_id: int) -> ResponseNerModel:
        if ResponsesNer.responses.get(request_id):
            return ResponsesNer.responses[request_id]


class NerResponse:
    def __init__(self, name="./data/ner-model.pt", batch_size: int = 4, device: str = "gpu"):

        # load the model you trained
        if name == "./data/ner-model.pt":
            model_path = Path(name)
            if not model_path.is_file():
                DownloadModel().download_ner_model(name)
        self.model = SequenceTagger.load(name)
        self.batch_size = batch_size
        self.device = device

    def predict_json(self, input_json: List[InputElement], request_id: int) -> ResponseNerModel:
        ResponsesNer.responses[request_id] = ResponseNerModel(progression="in going")
        sentences = []
        for element in input_json:
            text = element.text
            lang = element.lang
            sentences.append(Sentence(text))
        self.model.predict(sentences, mini_batch_size=self.batch_size, verbose=True, embedding_storage_mode=self.device)
        result = []
        for sentence in sentences:
            result.append(sentence_to_json(sentence))
        ResponsesNer.responses[request_id] = ResponseNerModel(progression="done", result=result)
        print(request_id, " done ")
        return ResponseNerModel(progression="done", result=result)

