import os
import kenlm
import sentencepiece as spm
from tokenizers import normalizers


class KenlmModel:
    def __init__(
        self,
        vocabulary_size: str,
        ngram: str,
        pruning: str,
        normalize_nfd: bool = True,
        normalize_numbers: bool = True,
        normalize_puctuation: bool = True,
    ):
        self.model = kenlm.Model(os.path.join("files", f"jomleh-sp-{vocabulary_size}-o{ngram}-prune{pruning}.probing"))
        self.tokenizer = spm.SentencePieceProcessor(os.path.join("files", f"jomleh-sp-{vocabulary_size}.model"))

        norm_list = []
        if normalize_numbers:
            norm_list += [normalizers.Replace("۱", "۰"),
                          normalizers.Replace("۲", "۰"),
                          normalizers.Replace("۳", "۰"),
                          normalizers.Replace("۴", "۰"),
                          normalizers.Replace("۵", "۰"),
                          normalizers.Replace("۶", "۰"),
                          normalizers.Replace("۷", "۰"),
                          normalizers.Replace("۸", "۰"),
                          normalizers.Replace("۹", "۰"),
                          normalizers.Replace(".", "")]
        if normalize_puctuation:
            norm_list += [normalizers.Replace(".", ""),
                          normalizers.Replace("!", ""),
                          normalizers.Replace("؛", ""),
                          normalizers.Replace("،", ""),
                          normalizers.Replace("؟", "")]
        if normalize_nfd:
            norm_list += [normalizers.NFD()]
        norm_list += [normalizers.Strip()]

        self.normalizer = normalizers.Sequence(norm_list)

    @classmethod
    def from_pretrained(
        cls,
        vocabulary_size: str,
        ngram: str,
        pruning: str,
    ):
        return cls(vocabulary_size, ngram, pruning)

    def score(self, doc: str):
        doc = self.normalizer.normalize_str(doc)
        doc = ' '.join(self.tokenizer.encode(doc, out_type=str))
        return self.model.score(doc)

    def perplexity(self, doc: str):
        doc = self.normalizer.normalize_str(doc)
        doc = ' '.join(self.tokenizer.encode(doc, out_type=str))
        log_score = self.model.score(doc)
        length = len(doc.split()) + 1
        return round(10.0 ** (-log_score / length), 1)