"""Defines pattern search pipeline based on ahocorasik."""
import itertools
from typing import Callable, Iterable, Iterator, List, Optional, Tuple

import ahocorasick
from spacy.tokens import Doc, Span
from spacy.tokens.token import Token
import textspan


# Sometimes matched text is different from original text
# since `PatternSearcher` can match the `lemma`.
# This extension holds the matched text.
PATTERN_MATCH_AS = "pattern_match_as"
Span.set_extension(PATTERN_MATCH_AS, default=None, force=True)


class PatternSearcher:
    def __init__(
        self,
        model: ahocorasick.Automaton,
        label: str = "matched",
        lower: bool = False,
        normalizer: Optional[Callable[[Token], str]] = None,
        lemma: bool = False,
        extend_span_to_token_boundary: bool = False,
        ignore_space: bool = True,
    ):
        self.model = model
        self.lower = lower
        self.normalizer = normalizer
        self.lemma = lemma
        self.label = label
        self.extend_span_to_token_boundary = extend_span_to_token_boundary
        self.ignore_space = ignore_space

    @staticmethod
    def get_model_from_words(words: Iterable[str]) -> ahocorasick.Automaton:
        model = ahocorasick.Automaton()
        for word in words:
            model.add_word(word, word)
        model.make_automaton()
        return model

    def get_char_spans(self, text: str) -> Iterator[Tuple[int, int, str]]:
        for j, word in self.model.iter(text):
            i = j - len(word) + 1
            yield i, j + 1, word

    def _get_spans_from_matches(
        self, matches: Iterable[Tuple[int, int, str]], tokens: List[Tuple[int, int]]
    ) -> Iterator[Tuple[int, int]]:
        idx = sorted((i, j) for i, j, _ in matches)
        spans = textspan.lift_spans_index(idx, tokens)
        for (l, l_ok), (r, r_ok) in spans:
            if self.extend_span_to_token_boundary or (l_ok and r_ok):
                yield (l, r)

    def _search_by_normalizer(
        self, doc: Doc, normalizer: Callable[[Token], str], ignore_space: bool
    ) -> Iterable[Tuple[int, int]]:
        text = ""
        token_spans: List[Tuple[int, int]] = []
        left = 0
        for token in doc:
            token_text = normalizer(token)
            text += token_text
            right = left + len(token_text)
            if not ignore_space:
                text += token.whitespace_
            token_spans.append((left, right))
            if ignore_space:
                left = right
            else:
                left = right + len(token.whitespace_)
        matches = self.get_char_spans(text)
        return self._get_spans_from_matches(matches, token_spans)

    def __call__(self, doc: Doc) -> Doc:
        normalizers: List[Callable[[Token], str]] = [lambda x: x.text]
        if self.lower:
            normalizers.append(lambda x: x.text.lower())
        if self.lemma:
            normalizers.append(lambda x: x.lemma_)
        if self.normalizer is not None:
            normalizers.append(self.normalizer)

        spans: Iterable[Tuple[int, int]] = []
        for normalizer in normalizers:
            spans = itertools.chain(
                spans, self._search_by_normalizer(doc, normalizer, ignore_space=False)
            )
            if self.ignore_space:
                spans = itertools.chain(
                    spans,
                    self._search_by_normalizer(doc, normalizer, ignore_space=True),
                )

        ents = list(doc.ents)
        for i, j in spans:
            ent = Span(doc, i, j, label=self.label)
            ents.append(ent)
        selected = textspan.remove_span_overlaps_idx([(s.start, s.end) for s in ents])
        doc.ents = tuple(ents[i] for i in selected)
        return doc
