"""
 Copyright 2024 Intel Corporation

 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

      https://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.
 """

import re


def split_on_whitespace(document, new_line=False, tab=False):
    """
    This method also removes concatenated spaces.

    :param document: document to be splited
    :param new_line: whether to split document with '\\\\n'
    :param tag: whether to split document with '\\\\t'
    :return: word list obtained after splitting document
    """
    sep = [' '] + new_line * ['\n'] + tab * ['\t']
    sep = '|'.join(sep)
    split_document = re.split(sep, document)
    split_document = [word for word in split_document if word]
    return split_document


def get_words_from_document(document,
                            token_func=None,
                            new_line=True,
                            tab=True,):
    """
    Get words from a document. Useful to compute ratios, like the
    stopwords ratio.

    :param document: document that need to split words
    :param token_func: function of tokenizer, if specified, the function
        will be used for split document into different tokens.
    :param new_line: whether to use `\\\\n' to split words
    :param tab: whether to use '\\\\t' to split words
    :return: word list obtained from document
    """
    if token_func:
        words = token_func(document)
    else:
        words = split_on_whitespace(document, new_line, tab)
    return words


def strip(document, strip_characters):
    """
    Way faster than document.strip(strip_characters) since strip_characters is
    now a set instead of a str, and it contains a lot of elements (all the
    emojis).

    :param document: document to be processed
    :param strip_characters: characters uesd for stripping document
    :return: stripped document
    """
    if not document:
        return document
    beg_ind = 0
    end_ind = len(document)
    for i in range(len(document)):
        if document[i] in strip_characters:
            beg_ind += 1
        else:
            break
    for i in range(1, len(document) + 1):
        if document[-i] in strip_characters:
            end_ind -= 1
        else:
            break
    document_stripped = document[beg_ind:end_ind]
    return document_stripped


def words_augmentation(words, group_size, join_char):
    """
    Augment words, especially for Chinese (without a space between words) and
    Vietnamese (with a space between syllables).

    :param word: word list to be augmented
    :param group_size: the size of word groups that need to be merged
    :param join_char: characters to be added between word group
    :return: word list after augment
    """
    augmentation = [
        join_char.join(words[i:i + group_size])
        for i in range(len(words) - group_size + 1)
    ]
    return augmentation


def words_refinement(words,
                     lower_case=False,
                     strip_chars=None,
                     use_words_aug=False,
                     words_aug_group_sizes=[2],
                     words_aug_join_char=''):
    """
    Refine split words. Non reversible since the document is split on
    multiple characters, words are stripped of special characters and
    characters are converted to lower case.

    :param words: the word list to be augmented
    :param lower_case: whether to convert word to lowercase
    :param strip_chars: chars that need to be stripped in words
    :param use_words_aug: whether to use word augmentation
    :param words_aug_group_sizes: the size of word groups that need to
        be merged
    :param words_aug_join_char: characters to be added between word
        group
    :return: refined words or word list
    """

    if lower_case:
        words = [word.lower() for word in words]
    if strip_chars:
        words = [strip(word, strip_chars) for word in words]
        words = [word for word in words if word]
    if use_words_aug:
        augmentation = [
            words_augmentation(words, group_size, words_aug_join_char)
            for group_size in words_aug_group_sizes
        ]
        augmentation = [word for augm in augmentation for word in augm]
        words = words + augmentation
    return words
