#!/usr/bin/env python3
import os
from typing import List, Tuple, Dict
from string import punctuation
from math import log


SMOOTHING_CONSTANT = 10000
OCC_PENALTY_CONSTANT_WORD = 1000
OCC_PENALTY_CONSTANT_BGR = 5000

MAX_POS = 10000

class Word:
    def __init__(self, word: str, occurences: List[Tuple[str, int]]):
        self.word = word
        self.length_modifier = 1 + log(len(word))
        self.occurence_list = occurences
        self.occurence_dict = dict(occurences)
        self.occurence_penalty: Dict[str, float] = self.__compute_occurence_penalties()


    def __compute_occurence_penalties(self) -> Dict[str, float]:
        penalty: Dict[Tuple[str, float]] = dict({(lang, 1) for lang, position in self.occurence_list})
        has_cs_and_csd: bool = "csd" in penalty.keys() and "cs" in penalty.keys()
        has_sk_and_skd: bool = "skd" in penalty.keys() and "sk" in penalty.keys()
        for lang, position in self.occurence_list:
            partial_penalty = log((MAX_POS - position) / OCC_PENALTY_CONSTANT_WORD + 1)
            for pen_lang in penalty.keys():
                if pen_lang == lang:
                    continue
                if pen_lang in ("cs", "csd") and lang in ("cs", "csd") and has_cs_and_csd:
                    continue
                if pen_lang in ("sk", "skd") and lang in ("sk", "skd") and has_sk_and_skd:
                    continue
                penalty[pen_lang] += partial_penalty
        return penalty


    def compute_score(self) -> Dict[str, float]:
        result_dict: Dict[str, float] = dict()
        for lang, order in self.occurence_list:
            result_dict[lang] = log((SMOOTHING_CONSTANT / (order + 1)) * self.length_modifier) / self.occurence_penalty.get(lang)
        return result_dict


    def __str__(self):
        return f"Word: {self.word}\n\tLength modifier: {self.length_modifier}x\n" + \
                f"\tOccurence in other languages: {sorted(self.occurence_list, key=lambda x: x[1])}\n" + \
                f"\tOccurence penalty: {self.occurence_penalty}\n" + \
                f"\tTOTAL SCORE: {sorted(self.compute_score().items(), key=lambda x: x[1], reverse=True)}"


class Bigram:
    def __init__(self, first: str, second: str, occurences: List[Tuple[str, int]]):
        self.first = first
        self.second = second
        self.length_modifier = 1 + log(len(first)) + log(len(second))
        self.occurence_list = occurences
        self.occurence_penalty = self.__compute_occurence_penalties()
    

    def __compute_occurence_penalties(self):
        has_cs_and_csd = "cs" in self.occurence_list and "csd" in self.occurence_list
        has_sk_and_skd = "sk" in self.occurence_list and "skd" in self.occurence_list
        penalty = 0
        for lang, position in self.occurence_list:
            if (lang == "csd" and has_cs_and_csd) or (lang == "skd" and has_sk_and_skd):
                continue
            penalty += log(position / OCC_PENALTY_CONSTANT_BGR + 1)
        return penalty


    def compute_score(self) -> Dict[str, float]:
        result_dict: Dict[str, float] = dict()
        for lang, order in self.occurence_list:
            #print((SMOOTHING_CONSTANT / (order + 1)) * self.length_modifier)
            result_dict[lang] = log(((SMOOTHING_CONSTANT * 5) / (order + 1)) * self.length_modifier) / (self.occurence_penalty + 1)
        return result_dict
    

    def __str__(self):
        return f"Bigram: {self.first} {self.second}\n\tLength modifier: {self.length_modifier}x\n" + \
                f"\tOccurence in other languages: {sorted(self.occurence_list, key=lambda x: x[1], reverse=True)}\n" + \
                f"Occurence penalty: {self.occurence_penalty}" + \
                f"\tTOTAL SCORE: {sorted(self.compute_score().items(), key=lambda x: x[1], reverse=True)}"


class DictLanguageIdentifier:
    def __init__(self, dict_path: str, bgr_path: str):
        self.words : Dict[str, Word] = dict()
        self.bigrams : Dict[Tuple[str, str], Bigram] = dict()
        self.supported_languages: List[str] = []
        self.__load_dictionaries(dict_path, bgr_path)

    def __load_dictionaries(self, dict_path: str, bgr_path: str):
        dicts = os.listdir(dict_path)
        wordlist: Dict[str, List[Tuple[str, int]]] = dict()
        bgrlist: Dict[Tuple[str, str], List[Tuple[str, int]]] = dict()
        for d in dicts:
            lang_id = d[:-4]
            self.supported_languages.append(lang_id)
            with open(os.path.join(dict_path, d), "r") as dict_f:
                for idx, line in enumerate(dict_f.readlines()):
                    if line.strip() not in wordlist:
                        wordlist[line.strip()] = []
                    wordlist[line.strip()].append((lang_id, idx))
            with open(os.path.join(bgr_path, d), "r") as bgr_f:
                for idx, line in enumerate(bgr_f.readlines()):
                    first, second = line.strip().split("\t")
                    if (first, second) not in bgrlist:
                        bgrlist[(first, second)] = []
                    bgrlist[(first, second)].append((lang_id, idx))
        for word, langs in wordlist.items():
            self.words[word] = Word(word, langs)
        for bgr, langs in bgrlist.items():
            self.bigrams[(bgr[0], bgr[1])] = Bigram(bgr[0], bgr[1], langs)


    def __preprocess(self, sentence: str) -> List[str]:
        sentence = sentence.strip().split()
        result = []
        for word in sentence:
            word = word.lower().strip().strip(punctuation)
            if not word.isalpha():
                continue
            result.append(word)
        return result


    def __find_bigrams(self, tokens: List[str], score_dict, debug=False):
        prev = None
        for token in tokens:
            if prev is None:
                prev = token
                continue
            if (prev, token) in self.bigrams:
                scores = self.bigrams[(prev, token)].compute_score()
                if debug:
                    print(self.bigrams[(prev, token)])
                for lang, score in scores.items():
                    score_dict[lang] += score
            prev = token



    def __assign_score(self, word: str, score_dict, debug=False) -> bool:
        if word not in self.words:
            if debug:
                print(f"Word {word} not in dictionaries.")
            return False
        scores = self.words[word].compute_score()
        if debug:
            print(self.words[word])
        for lang, score in scores.items():
            score_dict[lang] += score
        return True


    def __normalize(self, result_list):
        conf_sum = sum(map(lambda x: x[1], result_list))
        try:
            return list(map(lambda x: (x[0], x[1] / conf_sum), result_list))
        except ZeroDivisionError:
            return [("???", 1.0)]


    def predict(self, sentence: str, debug=False):
        score_dict = dict({(lang_id, 0) for lang_id in self.supported_languages})
        tokens = self.__preprocess(sentence)
        unknowns = 1
        for token in tokens:
            is_known = self.__assign_score(token, score_dict, debug)
            if not is_known:
                unknowns += 1
        self.__find_bigrams(tokens, score_dict, debug)
        results = list(map(lambda x: (x[0], x[1] / unknowns), sorted(list(score_dict.items()), key=lambda x: x[1], reverse=True)))
        norm = results #self.__normalize(results)
        if debug:
            print(results)
        if all(map(lambda conf: conf[1] == 0.0, results)):
            return ("???", 0.0)
        return norm[0]

if __name__ == "__main__":
    test_ident = DictLanguageIdentifier("dictionaries", "bigram_dicts")
    while True:
        try:
            print(test_ident.predict(input(">>> "), debug=True))
        except (EOFError, KeyboardInterrupt) as err:
            break
