# Copyright (c) 2024, NVIDIA CORPORATION.  All rights reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

from nemo_curator.filters.doc_filter import DocumentFilter
from nemo_curator.utils.constants import (
    bullet_list,
    common_english_words,
    ellipsis_marks,
    end_marks,
    policy_substrings,
    regex_alpha,
    regex_alphanum,
    regex_digit,
    regex_hash,
    regex_paren,
    regex_url,
    white_space_list,
)
from nemo_curator.utils.text_utils import (
    get_ngrams,
    get_paragraphs,
    get_sentences,
    get_word_splitter,
)


class NonAlphaNumericFilter(DocumentFilter):
    """
    If more than 25% of the document is non-alphanumeric, then discard.
    Intended to be applied only to English text.
    Source: Adapted from Gopher (Rae et al., 2021)
    """

    def __init__(self, max_non_alpha_numeric_to_text_ratio: float = 0.25):
        super().__init__()
        self._cutoff = max_non_alpha_numeric_to_text_ratio
        self._name = "alpha_numeric"

    def score_document(self, text):
        nchar = len(text)
        if nchar > 0:
            score = (nchar - len(regex_alphanum.findall(text))) / nchar
        else:
            # Remove the document if it is empty
            score = 1.0
        return score

    def keep_document(self, score):
        return score <= self._cutoff


class SymbolsToWordsFilter(DocumentFilter):
    """
    Remove any document with a symbol-to-word ratio greater than
    0.1 for either the hash symbol or the elipsis.
    Source: Gopher (Rae et al., 2021)

    For Chinese and Japanese text, we use external libraries to split the text
    because these languages are not separated by spaces. For all other langauges,
    such as English, we assume words are separated by spaces.
    """

    def __init__(self, max_symbol_to_word_ratio=0.1, lang="en"):
        super().__init__()
        self._cutoff = max_symbol_to_word_ratio
        self._word_splitter = get_word_splitter(lang)
        self._name = "symbol_to_word"

    def score_document(self, text):
        num_symbol_words = 0
        words = self._word_splitter(text.strip())
        for word in words:
            word = word.strip()
            # Checks if the word is an elipsis or consists mostly of symbols.
            symbol_ratio = len(regex_hash.findall(word)) / len(word)
            if word in ellipsis_marks or symbol_ratio > 0.5:
                num_symbol_words += 1
        return num_symbol_words / len(words)

    def keep_document(self, score):
        return score <= self._cutoff


class NumbersFilter(DocumentFilter):
    """
    If more than 15% of the document contains numbers, then discard.
    """

    def __init__(self, max_number_to_text_ratio=0.15):
        super().__init__()
        self._cutoff = max_number_to_text_ratio
        self._name = "numbers_ratio"

    def score_document(self, text):
        nchar = len(text)
        if nchar > 0:
            score = len(regex_digit.findall(text)) / nchar
        else:
            # Remove if the document is empty
            score = 1.0
        return score

    def keep_document(self, score):
        return score <= self._cutoff


class UrlsFilter(DocumentFilter):
    """
    If more than 20% of the document is comprised of URLs, then discard.
    """

    def __init__(self, max_url_to_text_ratio=0.2):
        super().__init__()
        self._cutoff = max_url_to_text_ratio
        self._name = "urls_ratio"

    def score_document(self, text):
        all_urls = regex_url.findall(text)
        url_chars = sum([len(url) for url in all_urls])
        nchar = len(text)
        if nchar > 0:
            score = url_chars / nchar
        else:
            # Remove if the document is empty
            score = 1.0
        return score

    def keep_document(self, score):
        return score <= self._cutoff


class BulletsFilter(DocumentFilter):
    """
    If more than 90% of the lines start with a bullet, then discard.
    Source: Gopher (Rae et al., 2021)
    """

    def __init__(self, max_bullet_lines_ratio=0.9):
        super().__init__()
        self._cutoff = max_bullet_lines_ratio
        self._sentences = None
        self._name = "bullet_ratio"

    def score_document(self, text):
        # Get sentences
        sentences = self._sentences
        if sentences is None:
            sentences = get_sentences(text)
        num_bullet_lines = 0
        for sentence in sentences:
            for bullet in bullet_list:
                if sentence.strip().startswith(bullet):
                    num_bullet_lines += 1
                    break
        return num_bullet_lines / len(sentences)

    def keep_document(self, score):
        return score <= self._cutoff


class WhiteSpaceFilter(DocumentFilter):
    """
    If the document contains a significant number
    of white space characters, then discard.
    """

    def __init__(self, max_white_space_ratio=0.25):
        super().__init__()
        self._cutoff = max_white_space_ratio
        self._name = "white_space"

    def score_document(self, text):
        # Do not strip the document since we want to
        # include leading and trailing whitepsaces as well.
        nchar = len(text)
        if nchar > 0:
            score = len([x for x in text if x in white_space_list]) / nchar
        else:
            # Remove if the document is empty
            score = 1.0
        return score

    def keep_document(self, score):
        return score <= self._cutoff


class ParenthesesFilter(DocumentFilter):
    """
    If more than 10% of the sentence is in parentheses, then discard.
    """

    def __init__(self, max_parentheses_ratio=0.1):
        super().__init__()
        self._max_parentheses_ratio = max_parentheses_ratio
        self._name = "parentheses_ratio"

    def score_document(self, text):
        nchar = len(text)
        if nchar > 0:
            score = len(regex_paren.findall(text)) / nchar
        else:
            # Remove if the document is empty
            score = 1.0
        return score

    def keep_document(self, score):
        return score <= self._max_parentheses_ratio


class LongWordFilter(DocumentFilter):
    """
    If the document contains a word longer than 1000 characters, then discard.
    NOTE: This seems to be catching things like minified `.js` files
    that don't have spaces anywhere.
    Source: C4 (Google)

    For Chinese and Japanese text, we use external libraries to split the text
    because these languages are not separated by spaces. For all other langauges,
    such as English, we assume words are separated by spaces.
    """

    def __init__(self, max_word_length=1000, lang="en"):
        super().__init__()
        self._max_word_length = max_word_length
        self._word_splitter = get_word_splitter(lang)
        self._name = "max_word_length"

    def score_document(self, text):
        return max(len(w) for w in self._word_splitter(text.strip()))

    def keep_document(self, score):
        return score <= self._max_word_length


class WordCountFilter(DocumentFilter):
    """
    If a document contains a number of words not
    within a specified range, then discard.

    For Chinese and Japanese text, we use external libraries to split the text
    because these languages are not separated by spaces. For all other langauges,
    such as English, we assume words are separated by spaces.
    """

    def __init__(self, min_words=50, max_words=100000, lang="en"):
        super().__init__()
        self._min_words = min_words
        self._max_words = max_words
        self._word_splitter = get_word_splitter(lang)
        self._name = "word_count"

    def score_document(self, text):
        return len(self._word_splitter(text.strip()))

    def keep_document(self, score):
        return self._min_words <= score <= self._max_words


class BoilerPlateStringFilter(DocumentFilter):
    """
    If more than 40% of paragraphs contain boilerplate strings, then discard.
    This includes things like "terms of use", "privacy policy", etc.
    Source: Adapted significantly from Google C4 processing.
    """

    def __init__(
        self,
        remove_if_at_top_or_bottom=True,
        max_boilerplate_string_ratio=0.4,
    ):
        super().__init__()
        self._remove_if_at_top_or_bottom = remove_if_at_top_or_bottom
        self._max_boilerplate_string_ratio = max_boilerplate_string_ratio
        self._boilerplate_paragraph_indices = []
        self._max_ratio = 1.0
        self._name = "boilerplate_string_ratio"

    def score_document(self, text):
        # Initialize variables
        boilerplate_paragraph_count = 0

        # Get the paragraphs
        paragraphs = get_paragraphs(text)

        # Check each paragraph
        for idx, paragraph in enumerate(paragraphs):
            paragraph = paragraph.strip().lower()
            if "lorem ipsum" in paragraph:
                return self._max_ratio
            if any(p in paragraph for p in policy_substrings):
                boilerplate_paragraph_count += 1

        return boilerplate_paragraph_count / len(paragraphs)

    def keep_document(self, score):
        return score <= self._max_boilerplate_string_ratio


class MeanWordLengthFilter(DocumentFilter):
    """
    If the mean word length is not in a specified range, then discard.

    For Chinese and Japanese text, we use external libraries to split the text
    because these languages are not separated by spaces. For all other langauges,
    such as English, we assume words are separated by spaces.
    """

    def __init__(
        self,
        min_mean_word_length=3,
        max_mean_word_length=10,
        lang="en",
    ):
        super().__init__()
        self._min_cutoff = min_mean_word_length
        self._max_cutoff = max_mean_word_length
        self._word_splitter = get_word_splitter(lang)
        self._name = "mean_word_length"

    def score_document(self, text):
        word_lens = [len(w) for w in self._word_splitter(text.strip()) if len(w) > 0]
        return sum(word_lens) / len(word_lens)

    def keep_document(self, score):
        return self._min_cutoff <= score <= self._max_cutoff


class RepeatedLinesFilter(DocumentFilter):
    """
    If the document shrinks by > 30% in terms of number of lines after
    removing duplicate lines, then discard.
    Source: Gopher (Rae et al., 2021)
    """

    def __init__(self, max_repeated_line_fraction=0.7):
        super().__init__()
        self._cutoff = max_repeated_line_fraction
        self._name = "repeated_lines"

    def score_document(self, text):
        sentences = self._sentences
        if sentences is None:
            sentences = get_sentences(text)
        return len(set(sentences)) / len(sentences)

    def keep_document(self, score):
        return score >= self._cutoff


class RepeatedParagraphsFilter(DocumentFilter):
    """
    If the document shrinks by > 30% in terms of number of lines after
    removing duplicate paragraphs, then discard.
    Source: Gopher (Rae et al., 2021)
    """

    def __init__(self, max_repeated_paragraphs_ratio=0.7):
        super().__init__()
        self._max_repeated_paragraphs_ratio = max_repeated_paragraphs_ratio
        self._name = "repeated_paragraphs"

    def score_document(self, text):
        paragraphs = self._paragraphs
        if paragraphs is None:
            paragraphs = get_paragraphs(text)
        return len(set(paragraphs)) / len(paragraphs)

    def keep_document(self, score):
        return score >= self._max_repeated_paragraphs_ratio


class RepeatedLinesByCharFilter(DocumentFilter):
    """
    If the document shrinks by > 20% in terms of number of lines
    after removing duplicate lines, then discard.
    Source: Gopher (Rae et al., 2021)
    """

    def __init__(self, max_repeated_lines_char_ratio=0.8):
        super().__init__()
        self._cutoff = max_repeated_lines_char_ratio
        self._name = "repeated_lines_char"

    def score_document(self, text):
        sentences = self._sentences
        if sentences is None:
            sentences = get_sentences(text)

        return len("".join(set(sentences))) / len("".join(sentences))

    def keep_document(self, score):
        return score >= self._cutoff


class RepeatedParagraphsByCharFilter(DocumentFilter):
    """
    If the document shrinks by > 10% in terms of number of lines after
    removing duplicate paragraphs, then discard.
    Source: Gopher (Rae et al., 2021)
    """

    def __init__(self, max_repeated_paragraphs_char_ratio=0.8):
        super().__init__()
        self._cutoff = max_repeated_paragraphs_char_ratio
        self._name = "repeated_paragraphs_char"

    def score_document(self, text):
        paragraphs = self._paragraphs
        if paragraphs is None:
            paragraphs = get_paragraphs(text)

        return len("".join(set(paragraphs))) / len("".join(paragraphs))

    def keep_document(self, score):
        return score >= self._cutoff


class RepeatingTopNGramsFilter(DocumentFilter):
    """
    If the document shrinks by > x% in terms of number of characters after
    removing the top n-grams, then discard.
    Source: Gopher (Rae et al., 2021)

    For Chinese and Japanese text, we use external libraries to split the text
    because these languages are not separated by spaces. For all other langauges,
    such as English, we assume words are separated by spaces.
    """

    def __init__(self, n=2, max_repeating_ngram_ratio=0.2, lang="en"):
        super().__init__()
        self._n = n
        self._cutoff = max_repeating_ngram_ratio
        self._max_ratio = 1.0
        self._word_splitter = get_word_splitter(lang)
        self._name = f"repeating_top_{n}grams"

    def score_document(self, text):
        ngrams = self._ngrams
        if ngrams is None:
            split_text = self._word_splitter(text.strip())
            if len(split_text) < self._n:
                return self._max_ratio
            ngrams = get_ngrams(split_text, self._n)
        unique_ngrams = set(ngrams)
        # Find the most frequent ngram in the zipped ngram list
        counts = {
            ngram: {"freq": 0, "num_chars": sum(len(word) for word in ngram)}
            for ngram in unique_ngrams
        }
        for ngram in ngrams:
            counts[ngram]["freq"] += 1
        most_frqnt_ngram = " ".join(max(counts, key=lambda x: counts[x]["freq"]))
        # Find the number of characters the most frequent ngram
        # contributes to the document
        nchar = len(text)
        len_diff = nchar - len(text.replace(most_frqnt_ngram, ""))
        if nchar > 0:
            score = len_diff / nchar
        else:
            # Remove if the document is empty
            score = 1.0
        return score

    def keep_document(self, score):
        return score <= self._cutoff


class RepeatingDuplicateNGramsFilter(DocumentFilter):
    """
    If the document shrinks by > x% in terms of number of characters
    after removing all duplicate n-grams, then discard.
    Source: Gopher (Rae et al., 2021)

    For Chinese and Japanese text, we use external libraries to split the text
    because these languages are not separated by spaces. For all other langauges,
    such as English, we assume words are separated by spaces.
    """

    def __init__(self, n=2, max_repeating_duplicate_ngram_ratio=0.2, lang="en"):
        super().__init__()
        self._n = n
        self._cutoff = max_repeating_duplicate_ngram_ratio
        self._max_ratio = 1.0
        self._word_splitter = get_word_splitter(lang)
        self._name = f"repeating_dup_{n}gram"

    def score_document(self, text):
        ngrams = self._ngrams
        if ngrams is None:
            split_text = self._word_splitter(text.strip())
            if len(split_text) < self._n:
                return self._max_ratio
            ngrams = get_ngrams(split_text, self._n)

        counts = {}
        duplicated_nchar = 0
        overlapping_ngrams = 0
        for ngram in ngrams:
            counts[ngram] = counts.get(ngram, 0) + 1
            if counts[ngram] > 1:
                # Count the number of characters in this ngram that haven't been counted already
                duplicated_ngrams = sum(
                    len(gram) for gram in ngram[overlapping_ngrams:]
                )
                # Count the spaces between the ngrams
                nspaces = min(self._n - overlapping_ngrams, self._n - 1)
                duplicated_nchar += duplicated_ngrams + nspaces
                overlapping_ngrams = self._n
            overlapping_ngrams = max(overlapping_ngrams - 1, 0)

        nchar = len(text)
        if nchar > 0:
            score = duplicated_nchar / nchar
        else:
            # Remove if the document is empty
            score = 1.0
        return score

    def keep_document(self, score):
        return score <= self._cutoff


class PunctuationFilter(DocumentFilter):
    """
    If more than 85% of the sentences do not end with a
    punctuation mark, then discard.
    Source: Google C4 processing
    """

    def __init__(self, max_num_sentences_without_endmark_ratio=0.85):
        super().__init__()
        self._cutoff = max_num_sentences_without_endmark_ratio
        self._name = "punctuation"

    def score_document(self, text):
        sentences = self._sentences
        if sentences is None:
            sentences = get_sentences(text)
        num_sentence_without_endmarks = len(
            [s for s in sentences if not s.strip().endswith(end_marks)]
        )
        return num_sentence_without_endmarks / len(sentences)

    def keep_document(self, score):
        return score <= self._cutoff


class EllipsisFilter(DocumentFilter):
    """
    If more than 30% of the sentences end with an elipsis, then discard.
    Source: Google C4 processing
    """

    def __init__(self, max_num_lines_ending_with_ellipsis_ratio=0.3):
        super().__init__()
        self._cutoff = max_num_lines_ending_with_ellipsis_ratio
        self._name = "ellipsis"

    def score_document(self, text):
        sentences = self._sentences
        if sentences is None:
            sentences = get_sentences(text)
        num_lines_ending_with_ellipsis = 0
        for sentence in sentences:
            for ellipsis in ellipsis_marks:
                if sentence.strip().lower().endswith(ellipsis):
                    num_lines_ending_with_ellipsis += 1
                    break
        return num_lines_ending_with_ellipsis / len(sentences)

    def keep_document(self, score):
        return score <= self._cutoff


class CommonEnglishWordsFilter(DocumentFilter):
    """
    If the sentence contains at least 2 common English words, then keep it.
    NOTE: We purposefully check for the lowercase versions of those common words
    to remove documents with over-capitalization.

    For Chinese and Japanese text, we use external libraries to split the text
    because these languages are not separated by spaces. For all other langauges,
    such as English, we assume words are separated by spaces.
    """

    def __init__(self, min_num_common_words=2, stop_at_false=True):
        super().__init__()
        self._cutoff = min_num_common_words
        self._stop_at_false = stop_at_false
        self._word_splitter = get_word_splitter("en")
        self._name = "common_english_words"

    def score_document(self, text):
        common_word_counter = 0
        for word in self._word_splitter(text.strip()):
            if word in common_english_words:
                common_word_counter += 1
            if self._stop_at_false and common_word_counter >= self._cutoff:
                return common_word_counter

        return common_word_counter

    def keep_document(self, score):
        return score >= self._cutoff


class WordsWithoutAlphabetsFilter(DocumentFilter):
    """
    80% of words in a document must contain at least one alphabetic character.
    Source: Gopher (Rae et al., 2021)

    For Chinese and Japanese text, we use external libraries to split the text
    because these languages are not separated by spaces. For all other langauges,
    such as English, we assume words are separated by spaces.
    """

    def __init__(self, min_words_with_alphabets=0.8, lang="en"):
        super().__init__()
        self._cutoff = min_words_with_alphabets
        self._word_splitter = get_word_splitter(lang)
        self._name = "words_without_alphabets"

    def score_document(self, text):
        num_english_alpha = 0
        words = self._word_splitter(text.strip())
        for word in words:
            if regex_alpha.search(word):
                num_english_alpha += 1

        return num_english_alpha / len(words)

    def keep_document(self, score):
        return score >= self._cutoff


class PornographicUrlsFilter(DocumentFilter):
    """
    Check if any of the URLs within the document point to pornography.
    """

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

    def score_document(self, text):
        all_urls = regex_url.findall(text)
        for url in all_urls:
            if "porn" in url:
                return 1

        return 0

    def keep_document(self, score):
        return score != 1
