"""

Filters.

@kylel, @soldni

"""

import regex
import uniseg.wordbreak
from tokenizers import Regex, pre_tokenizers

from ..core.data_types import DocResult, Document, Span
from ..core.registry import TaggerRegistry
from ..core.taggers import BaseTagger
from ..core.utils import split_paragraphs


@TaggerRegistry.add("char_length_v1")
class CharLengthV1(BaseTagger):
    def predict(self, doc: Document) -> DocResult:
        score = len(doc.text)
        return DocResult(doc=doc, spans=[Span(start=0, end=len(doc.text), type="length", score=score)])


@TaggerRegistry.add("char_length_with_paragraphs_v1")
class CharLengthWithParagraphsV1(BaseTagger):
    def predict(self, doc: Document) -> DocResult:
        spans = [
            Span(start=p.start, end=p.end, type="paragraph", score=len(p.text)) for p in split_paragraphs(doc.text)
        ]
        spans.append(Span(start=0, end=len(doc.text), type="document", score=len(doc.text)))
        return DocResult(doc=doc, spans=spans)


@TaggerRegistry.add("whitespace_tokenizer_v1")
class WhitespaceLengthV1(BaseTagger):
    WHITESPACE_REGEX = regex.compile(r"\w+|[^\w\s]+")

    def predict(self, doc: Document) -> DocResult:
        score = len(self.WHITESPACE_REGEX.split(doc.text))
        return DocResult(doc=doc, spans=[Span(start=0, end=len(doc.text), type="length", score=score)])


@TaggerRegistry.add("whitespace_tokenizer_with_paragraphs_v1")
class WhitespaceLengthParagraphsV1(WhitespaceLengthV1):
    def predict(self, doc: Document) -> DocResult:
        spans = [
            Span(start=p.start, end=p.end, type="paragraph", score=len(self.WHITESPACE_REGEX.split(p.text)))
            for p in split_paragraphs(doc.text)
        ]
        spans.append(Span(start=0, end=len(doc.text), type="document", score=sum(s.score for s in spans)))
        return DocResult(doc=doc, spans=spans)


@TaggerRegistry.add("uniseg_length_paragraphs_v1")
class UnisegParagraphsV1(BaseTagger):
    def predict(self, doc: Document) -> DocResult:
        spans = []
        for para in split_paragraphs(doc.text):
            # we ignore whitespace-only tokens when counting words
            para_length = sum(1 for w in uniseg.wordbreak.words(para.text.strip()) if w.strip())
            spans.append(Span(start=para.start, end=para.end, type="paragraph", score=para_length))

            # we have to record negative length because mixer can only work on greater than operations,
            # and we might want to drop paragraphs that are shorter than a certain length n, so we need
            # to filter on >= n
            spans.append(Span(start=para.start, end=para.end, type="negative_paragraph", score=-para_length))
        return DocResult(doc=doc, spans=spans)


@TaggerRegistry.add("uniseg_length_paragraphs_with_doc_length_v1")
class UnisegParagraphsWithDocLengthV1(UnisegParagraphsV1):
    def predict(self, doc: Document) -> DocResult:
        doc_results = super().predict(doc)
        pos_len = sum(s.score for s in doc_results.spans if s.type == "paragraph")
        neg_len = sum(s.score for s in doc_results.spans if s.type == "negative_paragraph")
        doc_results.spans.append(Span(start=0, end=len(doc.text), type="document", score=pos_len))
        doc_results.spans.append(Span(start=0, end=len(doc.text), type="negative_document", score=neg_len))
        return doc_results


@TaggerRegistry.add("olmo_pretokenizer_v1")
class OlmoPreTokenizerV1(BaseTagger):
    def __init__(self) -> None:
        self.pre_tokenizer = pre_tokenizers.Sequence(
            [
                # Split on all punctuation.
                pre_tokenizers.Split(
                    pattern=Regex(" ?[[:punct:]]"),
                    behavior="isolated",
                    invert=False,
                ),
                # Split up digits.
                pre_tokenizers.Split(
                    pattern=Regex(" ?\\d"),
                    behavior="isolated",
                    invert=False,
                ),
                pre_tokenizers.ByteLevel(add_prefix_space=False, use_regex=True),
            ]
        )

    def predict(self, doc: Document) -> DocResult:
        score = len(self.pre_tokenizer.pre_tokenize_str(doc.text))
        return DocResult(doc=doc, spans=[Span(start=0, end=len(doc.text), type="length", score=score)])


@TaggerRegistry.add("olmo_pretokenizer_with_paragraphs_v1")
class OlmoPreTokenizerParagraphsV1(OlmoPreTokenizerV1):
    def predict(self, doc: Document) -> DocResult:
        spans = [
            Span(
                start=p.start, end=p.end, type="paragraph", score=len(self.pre_tokenizer.pre_tokenize_str(p.text))
            )
            for p in split_paragraphs(doc.text)
        ]
        spans.append(Span(start=0, end=len(doc.text), type="document", score=sum(s.score for s in spans)))
        return DocResult(doc=doc, spans=spans)
