from spacy.language import Language
from spacy.tokens import Span, Doc
from transformers import (
    AutoConfig,
    AutoTokenizer,
    AutoModelForTokenClassification,
)
from transformers import pipeline
from transformers.utils import logging


# REQUIRES sentencizer
# https://huggingface.co/m3hrdadfi/typo-detector-distilbert-en?text=The+review+thoroughla+assessed+all+aspects+of+JLENS+SuR+and+CPG+esign+maturit+and+confidence+.
@Language.factory("typo_component")
class TypoFactory:
    def __init__(self, nlp: Language, name: str):
        self.nlp = nlp
        logging.disable_progress_bar()

        model_name_or_path = "m3hrdadfi/typo-detector-distilbert-en"
        tokenizer = AutoTokenizer.from_pretrained(model_name_or_path)
        self.model = AutoModelForTokenClassification.from_pretrained(
            model_name_or_path,
            config=AutoConfig.from_pretrained(model_name_or_path),
        )
        self.transformer_nlp = pipeline(
            "token-classification",
            model=self.model,
            tokenizer=tokenizer,
            aggregation_strategy="average",
        )

        if not Span.has_extension("typos"):
            Span.set_extension("typos", default=None)

        if not Doc.has_extension("doc_typos"):
            Doc.set_extension("doc_typos", default=None)

    def __call__(self, doc):
        all_typos = []
        for sentence in doc.sents:
            txt = sentence.text
            typos = [
                txt[r["start"] : r["end"]] for r in self.transformer_nlp(txt)
            ]
            all_typos.extend(typos)
            sentence._.set("typos", {"count": len(typos), "words": typos})
        doc._.set(
            "doc_typos",
            {
                "counter": [{x: all_typos.count(x)} for x in set(all_typos)],
                "count": len(all_typos),
            },
        )
        return doc
