import re
from typing import Dict, Iterable, List, Union

import spacy
from spacy.language import Language 
from nltk import sent_tokenize
from spacy.lang.ro import Romanian

from rb.utils.downloader import download_spacy_model
from rb.core.lang import Lang

# JSON Example localhost:8081/spacy application/json
# {
#     "lang" : "en",
#     "blocks" : ["După terminarea oficială a celui de-al doilea război mondial, în conformitate cu discursul lui W. Churchill (prim ministru al Regatului Unit la acea dată), de la Fulton, s-a declanșat Războiul rece și a apărut conceptul de cortină de fier. Urmare a politicii consecvente de apărare a sistemului economic și politic (implicit a intereslor economice ale marelui capital din lumea occidentală) trupele germane, în calitate de prizonieri, aflate pe teritoriul Germaniei de Vest au fost reînarmate și au constituit baza viitorului Bundeswehr - armata regulată a R.F.G."]
# }

models = {
    Lang.EN: 'en_core_web_lg',
    Lang.NL: 'nl',
    Lang.FR: 'fr_core_news_md',
    Lang.ES: 'es',
    Lang.DE: 'de',
    Lang.IT: 'it',
    Lang.RO: 'ro_ud_ft_ner',
}

normalization = {
    'ro': [
        (re.compile("ş"), "ș"),
        (re.compile("Ş"), "Ș"),
        (re.compile("ţ"), "ț"),
        (re.compile("Ţ"), "Ț"),
        (re.compile("(\w)î(\w)"), "\g<1>â\g<2>")
    ]
}

def convertToPenn(pos: str, lang: str) -> str:
    if lang == 'fr':
        pos = pos.lower()
        if pos.startswith('noun') or pos.startswith('propn'):
            return "NN"
        if pos.startswith("verb"):
            return "VB"
        if pos.startswith("adj"):
            return "JJ"
        if pos.startswith("adv"):
            return "RB"
        if pos.startswith("adp"):
            return "IN"
        if pos.startswith("cconj"):
            return "CC"
        return ""
    if lang == 'nl':
        pos = pos.lower()
        if pos.startswith('n_') or pos.startswith('n|') or pos.startswith('propn'):
            return "NN"
        if pos.startswith("v_") or pos.startswith("v|"):
            return "VB"
        if pos.startswith("adj"):
            return "JJ"
        if pos.startswith("adv"):
            return "RB"
        if pos.startswith("adp"):
            return "IN"
        if pos.startswith("cconj") or pos.startswith("conj"):
            return "CC"
        return ""
    if lang == 'ro':
        pos = pos.lower()
        if pos.startswith("n"):
            return "NN"
        if pos.startswith("v"):
            return "VB"
        if pos.startswith("a"):
            return "JJ"
        if pos.startswith("r"):
            return "RB"
        if pos.startswith("s") or pos.startswith("cs"):
            return "IN"
        if pos.startswith("c"):
            return "CC"
        return ""
    if len(pos) > 2:
        return pos[:2]
    return pos
    

class SpacyParser:

    _INSTANCE = None

    def __init__(self):
        self.ner = spacy.load('xx_ent_wiki_sm')
        # self.romanian = Romanian()
        self.pipelines = {
            lang: spacy.util.get_lang_class(lang.value)()
            for lang in models
        }
        # for pipeline in self.pipelines.values():
        #     component = pipeline.create_pipe('tagger')   # 3. create the pipeline components
        #     pipeline.add_pipe(component)
        self.loaded_models = {}
    
    @classmethod
    def get_instance(cls) -> "SpacyParser":
        if cls._INSTANCE is None:
            cls._INSTANCE = SpacyParser()
        return cls._INSTANCE
        
    def preprocess(self, text: str, lang: str) -> str:
        if lang not in normalization:
            return text
        for pattern, replacement in normalization[lang]:
            text = re.sub(pattern, replacement, text)
        return text

    def get_tokens_lemmas(self, sentences: Iterable, lang: str) -> Iterable:
        if lang not in self.pipelines:
            return None
        pipeline = self.pipelines[lang]
        # sbd = pipeline.create_pipe('sentencizer')
        # pipeline.add_pipe(sbd)
        doc = pipeline.pipe((sent[:1].lower() + sent[1:] for sent in sentences), batch_size=100000, n_threads=16)
        # print([sent.string.strip() for sent in doc.sents])
        # print(len(doc.sents))
        # print("====================")
        # for token in doc:
        #     print(token.text, token.lemma_, token.pos_, token.tag_, token.dep_,
        #   token.shape_, token.is_alpha, token.is_stop)
        # print("====================")
        return doc
        # return [(token.text, token.lemma_) for token in doc]

    def tokenize_sentences(self, block: str) -> List[str]:
        return sent_tokenize(block)
    
    def get_model(self, lang: Union[Lang, str]) -> Language:
        if isinstance(lang, str):
            lang = Lang(lang)
        if lang not in self.loaded_models:
            folder = "resources/{}/spacy/{}".format(lang.value, models[lang])
            for model in [models[lang], folder]:
                try:
                    self.loaded_models[lang] = spacy.load(model)
                    break
                except:
                    pass
            else:
                if not download_spacy_model(lang, models[lang]):
                    raise FileNotFoundError("Spacy model not found")
                self.loaded_models[lang] = spacy.load(folder)
        return self.loaded_models[lang]

    def parse(self, sentence: str, lang: Lang):
        model = self.get_model(lang)
        doc = model(sentence)
        # print([sent.string.strip() for sent in doc.sents])
        # for chunk in doc.noun_chunks:
        #     print(chunk.text, chunk.root.text, chunk.root.dep_,
        #        chunk.root.head.text)

        # print("********************")
        # for token in doc:
        #     print(token.text, token.lemma_, token.pos_, token.tag_, token.dep_,
        #   token.shape_, token.is_alpha, token.is_stop)
        # print("********************")
        # return [(token.text, token.lemma_, token.pos_, token.tag_) for token in doc]
        return doc


    def process(self, doc):
        lang = Lang(doc["lang"])
        for block in doc["blocks"]:
            sents = sent_tokenize(block["text"])
            block["sentences"] = list()

            for sent in sents:
                ne = self.ner(sent)
                tokens = self.parse(sent, lang)
                # print(ne)
                # print(pos)
                res_sent = {}
                res_sent["text"] = sent
                res_sent["words"] = []
                # get pos tags 
                for w in tokens:
                    wp = {"text" : w.text}
                    wp["index"] = w.i
                    wp["lemma"] = w.lemma_
                    wp["pos"] = convertToPenn(w.tag_, lang)
                    wp["dep"] = w.dep_
                    wp["ner"] = w.ent_type_
                    wp["head"] = w.head.i
                    res_sent["words"].append(wp)
                # get named entities 
                for ent in [token for token in ne if token.ent_type != 0]:
                    for w in res_sent["words"]:
                        # or (' ' in ent[0] and w["word"] in ent[0])
                        if w["index"] == ent.i:
                            w["ner"] = ent.ent_type_
                block["sentences"].append(res_sent)
        return doc   

if __name__ == "__main__":
    spacyInstance = SpacyParser()

    sent = """
        După terminarea oficială a celui de-al doilea război mondial, în conformitate cu discursul lui W. Churchill (prim ministru al Regatului Unit la acea dată), de la Fulton, s-a declanșat Războiul rece și a apărut conceptul de cortină de fier. Urmare a politicii consecvente de apărare a sistemului economic și politic (implicit a intereslor economice ale marelui capital din lumea occidentală) trupele germane, în calitate de "prizonieri", aflate pe teritoriul Germaniei de Vest au fost reînarmate și au constituit baza viitorului "Bundeswehr" - armata regulată a R.F.G.

        Pe fondul evenimentelor din 1948 din Cehoslovacia (expulzări ale etnicilor germani, alegeri, reconstrucție economică) apare infiltrarea agenților serviciilor speciale ale S.U.A. și Marii Britanii cu rol de "agitatori". Existând cauza, trupele sovietice nu părăsesc Europa Centrală și de Est cucerită-eliberată, staționând pe teritoriul mai multor state. Aflate pe linia de demarcație dintre cele două blocuri foste aliate, armata sovietică nu a plecat din Ungaria decât după dizolvarea Tratatului de la Varșovia.
        """

    # sent = """
    #     După terminarea oficială a celui de-al doilea război mondial, în conformitate cu discursul lui Churchill, de la Fulton, s-a declanșat Războiul rece și a apărut conceptul de cortină de fier."""

    # print(spacyInstance.get_ner(sent))
    # print(spacyInstance.get_tokens_lemmas(sent))
    # doc = spacyInstance.parse("My sister has a dog. She loves him.", 'en')
    doc = spacyInstance.parse("Pensée des enseignants, production d’écrits, ingénierie éducative, enseignement à distance, traitement automatique de la langue, outils cognitifs, feedback automatique", 'fr')
    for token in doc:
        print(convertToPenn(token.tag_, 'fr'))   
    # print(spacyInstance.preprocess("coborî", 'ro'))
