from typing import List, Dict, Union

from flair.data import Sentence
from flair.models import SequenceTagger
from langdetect import detect_langs
from model.response_models import ResponseSummaryModel
from transformers import AutoTokenizer, AutoModelForSeq2SeqLM

from piraye import NormalizerBuilder, NltkTokenizer

from model.request_models import InputSummaryType
import re
from pathlib import Path

from src.helpers import DownloadModel

normalizer = NormalizerBuilder(
    tokenization=True).space_normal().punctuation_fa().remove_extra_spaces().build()
my_tokenizer = NltkTokenizer()


class ResponsesSummary:
    responses: Dict[int, ResponseSummaryModel] = {}

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

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


class NormalizerTexts:
    def __init__(self, name: str = "./data/pos-model.pt"):
        # load the model you trained
        if name == "./data/pos-model.pt":
            model_path = Path(name)
            if not model_path.is_file():
                DownloadModel().download_pos_model(name)
        self.model = SequenceTagger.load(name)

    def get_normal_text(self, text) -> List[str]:
        text = normalizer.normalize(text)
        text = str(re.sub(r"\[[۰-۹]*]", '', text))
        text = str(re.sub(r"\[[0-9]*]", '', text))
        sents = my_tokenizer.sentence_tokenize(text)
        final_text = []
        for sent in sents:
            if len(sent.strip()) < 10:
                continue
            lang = detect_langs(sent)
            if str(lang[0]).split(":")[0] == 'fa' and float(str(lang[0]).split(":")[1]) > 0.95:
                final_text.append(sent)
        return final_text

    @staticmethod
    def get_normal_summary(text) -> str:
        WHITESPACE_HANDLER = lambda k: re.sub('\s+', ' ', re.sub('\n+', ' ', k.strip()))
        text = str(re.sub(r"\[[۰-۹]*]", '', text))
        text = str(re.sub(r"\[[0-9]*]", '', text))
        text = text.replace("<n>", "")
        return WHITESPACE_HANDLER(text)

    def predict_sentences(self, not_cleaned_sentences: List[str]) -> str:
        sentences = []
        for sen in not_cleaned_sentences:
            sentences.append(Sentence(sen))
        self.model.predict(sentences, mini_batch_size=4, embedding_storage_mode="gpu")
        result = ""
        for i in range(len(sentences)):
            sentence = sentences[i]
            values = sentence.to_dict("pos")['entities']
            tag = values[-2]['labels'][0].value.split("_")[0]
            if tag == "V":
                result += not_cleaned_sentences[i]
        return result

    def get_normalized(self, input_json: InputSummaryType) -> List[str]:
        docs = []
        for text in input_json:
            docs.append(self.get_normal_text(text))
        final_docs = []
        for sentences in docs:
            final_docs.append(self.predict_sentences(sentences))
        return final_docs


class ShortResponse:
    def __init__(self, normalizer_input: Union[NormalizerTexts, type(None)], device: str = "gpu"):
        self.maxInputLength = None
        self.model = None
        self.tokenizer = None
        self.normalizer_input = normalizer_input
        self.device = device

    def load_model(self, name="./data/mT5_multilingual_XLSum"):
        print("loading " + name + "...")
        self.tokenizer = AutoTokenizer.from_pretrained(name)
        if self.device == "gpu":
            self.model = AutoModelForSeq2SeqLM.from_pretrained(name).to('cuda:0')
        else:
            self.model = AutoModelForSeq2SeqLM.from_pretrained(name)
        self.maxInputLength = 5000
        print(name + " loaded.")
        return

    def get_summary(self, article_text) -> str:
        input_ids = self.tokenizer(
            [article_text],
            return_tensors="pt",
            truncation=True,
            max_length=self.maxInputLength
        )["input_ids"]
        output_ids = self.model.generate(
            input_ids=input_ids,
            max_length=256,
            num_beams=5,
            no_repeat_ngram_size=2,
            early_stopping=True,
        )[0]
        summary = self.tokenizer.decode(
            output_ids,
            skip_special_tokens=True,
            clean_up_tokenization_spaces=True
        )
        return NormalizerTexts.get_normal_summary(summary)

    def get_result(self, input_json: InputSummaryType, request_id: int) -> ResponseSummaryModel:
        ResponsesSummary.responses[request_id] = ResponseSummaryModel(progression="started")
        if self.normalizer_input is None:
            normalized_input = input_json
        else:
            normalized_input = self.normalizer_input.get_normalized(input_json)
        result = []
        for sentence in normalized_input:
            result.append(self.get_summary(sentence))
        ResponsesSummary.responses[request_id] = ResponseSummaryModel(progression="done", result=result)
        print(request_id, " done ")
        return ResponseSummaryModel(progression="done", result=result)


class LongResponse(ShortResponse):
    def __init__(self, normalizer_input: Union[NormalizerTexts, type(None)], device: str = "gpu"):
        super().__init__(normalizer_input, device)

    def load_model(self, name="./data/ARMAN-MSR-persian-base-PN-summary"):
        print("loading " + name + "...")
        self.tokenizer = AutoTokenizer.from_pretrained(name)
        if self.device == "gpu":
            self.model = AutoModelForSeq2SeqLM.from_pretrained(name).to('cuda:0')
        else:
            self.model = AutoModelForSeq2SeqLM.from_pretrained(name)
        self.maxInputLength = 511
        print(name + " loaded.")
        return
