from eole.utils.logging import logger
from eole.transforms import register_transform
from .transform import Transform, TransformConfig
from pydantic import Field
import re


class TerminologyConfig(TransformConfig):
    termbase_path: str | None = Field(default=None, description="Path to a dictionary file with terms.")
    src_spacy_language_model: str | None = Field(
        default=None,
        description="Name of the spaCy language model for the source corpus.",
    )
    tgt_spacy_language_model: str | None = Field(
        default=None,
        description="Name of the spaCy language model for the target corpus.",
    )
    term_corpus_ratio: float | None = Field(default=0.3, description="Ratio of corpus to augment with terms.")
    term_example_ratio: float | None = Field(default=0.2, description="Maximum terms allowed in an example.")
    src_term_stoken: str | None = Field(default="｟src_term_start｠", description="The source term start token.")
    tgt_term_stoken: str | None = Field(default="｟tgt_term_start｠", description="The target term start token.")
    tgt_term_etoken: str | None = Field(default="｟tgt_term_end｠", description="The target term end token.")
    term_source_delimiter: str | None = Field(
        default="｟fuzzy｠",
        description="Any special token used for augmented source sentences. "
        "The default is the fuzzy token used in the "
        "FuzzyMatch transform.",
    )


class TermMatcher(object):
    def __init__(
        self,
        termbase_path,
        src_spacy_language_model,
        tgt_spacy_language_model,
        term_example_ratio,
        src_term_stoken,
        tgt_term_stoken,
        tgt_term_etoken,
        delimiter,
        term_corpus_ratio=0.2,
    ):
        import spacy

        self.term_example_ratio = term_example_ratio
        self.src_nlp = spacy.load(src_spacy_language_model, disable=["parser", "ner"])
        self.tgt_nlp = spacy.load(tgt_spacy_language_model, disable=["parser", "ner"])

        # We exclude tokenization for contractions in
        # order to avoid inconsistencies with pyeoleok's tokenization.
        # (e.g. "I ca n't" with spacy, "I can ' t" with pyonmttok)
        self.src_nlp.tokenizer.rules = {
            key: value
            for key, value in self.src_nlp.tokenizer.rules.items()
            if "'" not in key and "’" not in key and "‘" not in key
        }
        self.tgt_nlp.tokenizer.rules = {
            key: value
            for key, value in self.tgt_nlp.tokenizer.rules.items()
            if "'" not in key and "’" not in key and "‘" not in key
        }
        self.internal_termbase = self._create_internal_termbase(termbase_path)
        self.automaton = self._create_automaton()
        self.term_corpus_ratio = term_corpus_ratio
        self.src_term_stoken = src_term_stoken
        self.tgt_term_stoken = tgt_term_stoken
        self.tgt_term_etoken = tgt_term_etoken
        self.delimiter = delimiter

    def _create_internal_termbase(self, termbase_path):
        logger.debug("Creating termbase with lemmas for Terminology transform")

        # Use Spacy's stopwords to get rid of junk entries
        src_stopwords = self.src_nlp.Defaults.stop_words
        tgt_stopwords = self.tgt_nlp.Defaults.stop_words
        termbase = list()
        with open(termbase_path, mode="r", encoding="utf-8") as file:
            pairs = file.readlines()
            for pair in pairs:
                src_term, tgt_term = map(str, pair.split("\t"))
                src_lemma = " ".join("∥".join(tok.lemma_.split(" ")) for tok in self.src_nlp(src_term)).strip()
                tgt_lemma = " ".join(tok.lemma_ for tok in self.tgt_nlp(tgt_term)).strip()
                if src_lemma.lower() not in src_stopwords and tgt_lemma.lower() not in tgt_stopwords:
                    termbase.append((src_lemma, tgt_lemma))
        logger.debug(f"Created termbase with {len(termbase)} lemmas " f"for Terminology transform")
        return termbase

    def _create_automaton(self):
        import ahocorasick

        automaton = ahocorasick.Automaton()
        for term in self.internal_termbase:
            automaton.add_word(term[0], (term[0], term[1]))
        automaton.make_automaton()
        return automaton

    def _src_sentence_with_terms(self, source_string, target_string) -> tuple:

        maybe_augmented = source_string.split(self.delimiter)
        source_only = maybe_augmented[0].strip()
        augmented_part = maybe_augmented[1].strip() if len(maybe_augmented) > 1 else None

        doc_src = self.src_nlp(source_only)
        doc_tgt = self.tgt_nlp(target_string)

        # Perform tokenization with spacy for consistency.
        tokenized_source = [tok.text for tok in doc_src]
        lemmatized_source = ["∥".join(tok.lemma_.lower().split(" ")) for tok in doc_src]
        lemmatized_target = [tok.lemma_.lower() for tok in doc_tgt]

        lemmatized_source_string = " ".join(lemmatized_source)

        offset = 0
        source_with_terms = list()
        term_counter = 0

        max_terms_allowed = int(len(tokenized_source) * self.term_example_ratio)
        is_match = False
        for match_end, (src_entry, tgt_entry) in self.automaton.iter_long(lemmatized_source_string):

            if term_counter == max_terms_allowed:
                break

            match_start = match_end - len(src_entry) + 1

            # We ensure that the target lemma is present in the lemmatized
            # target string, that the match is an exact match (there is
            # whitespace before or after the term)
            # and we perform some bound checking.
            if (
                (tgt_entry.lower() not in " ".join(lemmatized_target).lower())
                or (
                    len(lemmatized_source_string) != match_end + 1
                    and not (lemmatized_source_string[match_end + 1].isspace())
                )
                or (not lemmatized_source_string[match_start - 1].isspace() and match_start != 0)
            ):
                continue
            else:
                term_counter += 1

                # Map the lemmatized string match index to
                # the lemmatized list index
                lemma_list_index = 0
                for i, w in enumerate(lemmatized_source):
                    if lemma_list_index == match_start:
                        lemma_list_index = i
                        break
                    else:
                        lemma_list_index += len(w) + 1

                # We need to know if the term is multiword
                num_words_in_src_term = len(src_entry.split(" "))
                src_term = " ".join(
                    tokenized_source[lemma_list_index : lemma_list_index + num_words_in_src_term]
                ).strip()

                # Join multiword target lemmas with a unique separator so
                # we can treat them as single word and not change the indices.
                tgt_term = tgt_entry.replace(" ", "∥").rstrip().lower()
                source_with_terms.append(
                    f"{lemmatized_source_string[offset: match_start]}"
                    f"{self.src_term_stoken}∥{src_term}∥{self.tgt_term_stoken}∥"
                    f"{tgt_term}∥{self.tgt_term_etoken}"
                )

                offset = match_end + 1
                is_match = True

        if is_match:
            source_with_terms.append(lemmatized_source_string[offset:])
            tokenized_source_with_terms = "".join(source_with_terms).split(" ")

            if not (len(tokenized_source) == len(lemmatized_source) == len(tokenized_source_with_terms)):
                final_string = " ".join(tokenized_source)
                fixed_punct = re.sub(r" ([^\w\s｟\-\–])", r"\1", final_string)
                return fixed_punct.split(" "), not is_match

            # Construct the final source from the lemmatized list
            # that contains the terms. We compare the tokens in the
            # term-augmented lemma list with the tokens in the original
            # lemma list. If the lemma is the same, then we replace with
            # the token from the original tokenized source list. If they
            # are not the same, it means the lemma has been augemented
            # with a term, so we inject this in the final list.
            completed_tokenized_source = list()
            for idx in range(len(tokenized_source_with_terms)):
                # Restore the spaces in multi-word terms
                src_lemma = tokenized_source_with_terms[idx].replace("∥", " ")
                if lemmatized_source[idx].replace("∥", " ") == src_lemma:
                    completed_tokenized_source.append(tokenized_source[idx])
                else:
                    completed_tokenized_source.append(src_lemma)

            if augmented_part is not None:
                final_string = " ".join(completed_tokenized_source + [self.delimiter] + augmented_part.split(" "))
            else:
                final_string = " ".join(completed_tokenized_source)

            fixed_punct = re.sub(r" ([^\w\s｟\-\–])", r"\1", final_string)
            return fixed_punct.split(" "), is_match
        else:
            final_string = " ".join(tokenized_source)
            fixed_punct = re.sub(r" ([^\w\s｟\-\–])", r"\1", final_string)
            return fixed_punct.split(" "), not is_match


@register_transform(name="terminology")
class TerminologyTransform(Transform):

    config_model = TerminologyConfig

    def __init__(self, config):
        super().__init__(config)

    def _parse_config(self):
        self.termbase_path = self.config.termbase_path
        self.src_spacy_language_model = self.config.src_spacy_language_model
        self.tgt_spacy_language_model = self.config.tgt_spacy_language_model
        self.term_corpus_ratio = self.config.term_corpus_ratio
        self.term_example_ratio = self.config.term_example_ratio
        self.term_source_delimiter = self.config.term_source_delimiter
        self.src_term_stoken = self.config.src_term_stoken
        self.tgt_term_stoken = self.config.tgt_term_stoken
        self.tgt_term_etoken = self.config.tgt_term_etoken

    @classmethod
    def get_specials(cls, config):
        """Add the term tokens to the src vocab."""
        src_specials = list()
        src_specials.extend([config.src_term_stoken, config.tgt_term_stoken, config.tgt_term_etoken])
        return (src_specials, list())

    def warm_up(self, vocabs=None):
        """Create the terminology matcher."""

        super().warm_up(None)
        self.termmatcher = TermMatcher(
            self.termbase_path,
            self.src_spacy_language_model,
            self.tgt_spacy_language_model,
            self.term_example_ratio,
            self.src_term_stoken,
            self.tgt_term_stoken,
            self.tgt_term_etoken,
            self.term_source_delimiter,
            self.term_corpus_ratio,
        )

    def batch_apply(self, batch, is_train=False, stats=None, **kwargs):
        bucket_size = len(batch)
        examples_with_terms = 0

        for i, (ex, _, _) in enumerate(batch):
            # Skip half examples to improve performance. This means we set
            # a hard limit for the `term_corpus_ratio` to 0.5, which is actually
            # quite high. TODO: We can add this (skipping examples) as an option
            if i % 2 == 0:
                original_src = ex["src"]
                augmented_example, is_match = self.apply(ex, is_train, stats, **kwargs)
                if is_match and (examples_with_terms < bucket_size * self.term_corpus_ratio):
                    examples_with_terms += 1
                    ex["src"] = augmented_example["src"]
                else:
                    ex["src"] = original_src

        logger.debug(f"Added terms to {examples_with_terms}/{bucket_size} examples")
        return batch

    def apply(self, example, is_train=False, stats=None, **kwargs) -> tuple:
        """Add terms to source examples."""

        example["src"], is_match = self.termmatcher._src_sentence_with_terms(
            " ".join(example["src"]), " ".join(example["tgt"])
        )
        return example, is_match
