#! /usr/bin/env python3

from dict_lang_ident import DictLanguageIdentifier #baseline
from language_identifier import LanguageIdentifier #officebot model
import fasttext #fasttext public model
from math import log
from typing import Tuple


FASTTEXT_PATH = "/nlp/projekty/officebot/trac-git/officebot_dev/models/lid.176.bin"

FASTTEXT_TO_OUR = { "__label__nl" : '__label__dutch',
           '__label__sw' : "__label__swedish",
           '__label__sr' : "__label__serbian",
           '__label__fr' : '__label__french',
           '__label__ru' : '__label__russian',
           '__label__sk' : '__label__slovak',
           "__label__cs" : '__label__czech',
           "__label__pt" : '__label__portuguese',
           "__label__pl" : '__label__polish',
           "__label__hu" : '__label__hungarian',
           "__label__en" : '__label__english',
           '__label__es' : '__label__spanish',
           "__label__fi" : '__label__finnish',
           "__label__de" : '__label__german',
           "__label__no" : '__label__norwegian',
           '__label__ja' : "__label__japanese",
           "__label__uk" : "__label__ukrainian",
           "__label__it" : "__label__italian"
           }

LONG_TO_SHORT = dict([(value, key[9:]) for key, value in FASTTEXT_TO_OUR.items()])
print(LONG_TO_SHORT)

class LangIdentEnsemble(object):
    def __init__(self, dict_paths: Tuple[str, str] = ("dictionaries/", "bigram_dicts/")):
        self.models = {"officebot" : LanguageIdentifier(),
                        "fasttext" : fasttext.load_model(FASTTEXT_PATH),
                        "baseline" : DictLanguageIdentifier(dict_paths[0], dict_paths[1]) }
        while not self.models["officebot"].ready():
            continue


    @staticmethod
    def __decide(text, fasttext, our, baseline):
        council = [fasttext, our, baseline]
        council_votes = dict()
        for vote, _ in council:
            if vote == "???":
                continue
            if vote not in council_votes:
                council_votes[vote] = 0
            council_votes[vote] += 1
        if any(map(lambda x: x > 1, council_votes.values())):
            return max(council_votes.items(), key=lambda x: x[1])
        if our[1] > 0.9:
            return our
        if fasttext[1] > 0.6:
            return fasttext
        if baseline[1] > 0.15*log(len(text.strip().split())):
            return baseline
        if (our[1] < 0.5 and fasttext[1] < 0.2):
            return baseline
        return our


    def identify_language(self, sentence, debug=False, ret_all=False):
        fasttext_prediction, fasttext_confidence = self.models["fasttext"].predict(sentence.replace("\n", ""))
        fasttext_prediction = fasttext_prediction[0][9:]
        dict_prediction, dict_confidence = self.models["baseline"].predict(sentence)
        our_prediction, our_confidence = self.models["officebot"].classify_single(sentence)
        if debug:
            print("Pretrained Fasttext model:", fasttext_prediction, fasttext_confidence)
            print("Dictionary identifier:", dict_prediction, dict_confidence)
            print("OfficeBot FastText model:", our_prediction, our_confidence)
        prediction_pack = ((fasttext_prediction, fasttext_confidence), 
                (our_prediction, our_confidence), 
                (dict_prediction, dict_confidence))
        ret = self.__decide(sentence, *prediction_pack)
        ret = (ret[0], sum(map(lambda x: x[1], filter(lambda x: x[0] == ret[0], prediction_pack))))
        if not ret_all:
            return ret
        return ret, {"fasttext": (fasttext_prediction, fasttext_confidence),
                     "baseline": (dict_prediction, dict_confidence),
                     "officebot": (our_prediction, our_confidence)}


if __name__ == "__main__":
    ensemble = LangIdentEnsemble()
    while True:
        try:
            sentence = input(">>>")
            print("Voting result:", ensemble.identify_language(sentence, debug=True))
        except (KeyboardInterrupt, EOFError):
            break

