import os
import json
from typing import List

import numpy as np

from augmentex.base import BaseAug


class WordAug(BaseAug):
    """Augmentation at the level of words.
    """

    def __init__(
        self,
        min_aug: int = 1,
        max_aug: int = 5,
        aug_rate: float = 0.3,
        punc_skip: float = 0.3,
    ) -> None:
        """
        Args:
            min_aug (int, optional): The minimum amount of augmentation. Defaults to 1.
            max_aug (int, optional): The maximum amount of augmentation. Defaults to 5.
            aug_rate (float, optional): Percentage of the phrase to which augmentations will be applied. Defaults to 0.3.
            punc_skip (float, optional): _description_. Defaults to 0.3.
        """
        super().__init__(min_aug=min_aug, max_aug=max_aug)
        stopwords_file = os.path.join(
            os.path.dirname(os.path.abspath(__file__)
                            ), "static_data", "stopwords_ru.json"
        )
        orfo_and_typos_file = os.path.join(
            os.path.dirname(os.path.abspath(__file__)),
            "static_data",
            "orfo_and_typos_words_ru.json",
        )
        with open(stopwords_file) as f:
            self.stopwords = json.load(f)
        with open(orfo_and_typos_file) as f:
            self.orfo_and_typos = json.load(f)
        self.aug_rate = aug_rate
        self.punc_skip = 0.3
        self.__actions = ["replace", "delete", "swap", "stopword", "reverse"]

    @property
    def actions_list(self) -> List[str]:
        """
        Returns:
            List[str]: A list of possible methods.
        """

        return self.__actions

    def _reverse_case(self, word: str) -> str:
        """Changes the case of the first letter to the reverse.

        Args:
            word (str): The initial word.

        Returns:
            str: A word with a different case of the first letter.
        """
        if len(word):
            if word[0].isupper():
                word = word.lower()
            else:
                word = word.capitalize()

        return word

    def _replace(self, word: str) -> str:
        """Replaces a word with the correct spelling with a word with spelling errors.

        Args:
            word (str): A word with the correct spelling.

        Returns:
            str: A misspelled word.
        """
        word_probas = self.orfo_and_typos.get(word, [[word], [1.0]])
        word = np.random.choice(word_probas[0], p=word_probas[1])

        return word

    def _delete(self) -> str:
        """Deletes a random word.

        Returns:
            str: Empty string.
        """

        return ""

    def _stopword(self, word: str) -> str:
        """Adds a stop word before the word.

        Args:
            word (str): Just word.

        Returns:
            str: Stopword + word.
        """
        stopword = np.random.choice(self.stopwords)

        return " ".join([stopword, word])

    def augment(self, text: str, action: str) -> str:
        aug_sent_arr = text.split()
        aug_idxs = self.aug_indexing(aug_sent_arr, self.aug_rate, clip=True)
        for idx in aug_idxs:
            if action == "delete":
                aug_sent_arr[idx] = self._delete()
            elif action == "reverse":
                aug_sent_arr[idx] = self._reverse_case(aug_sent_arr[idx])
            elif action == "swap":
                swap_idx = np.random.randint(0, len(aug_sent_arr) - 1)
                aug_sent_arr[swap_idx], aug_sent_arr[idx] = (
                    aug_sent_arr[idx],
                    aug_sent_arr[swap_idx],
                )
            elif action == "stopword":
                aug_sent_arr[idx] = self._stopword(aug_sent_arr[idx])
            elif action == "replace":
                aug_sent_arr[idx] = self._replace(aug_sent_arr[idx])
            else:
                raise NameError(
                    """These type of augmentation is not available, please check EDAAug.actions_list() to see
                available augmentations"""
                )

        return " ".join(aug_sent_arr)

    def random(self, text: str, actions: List[str] = []) -> str:
        """Applies a random method from the list to the selected text.

        Args:
            text (str): Source text.
            actions (List[str], optional): A list of methods that you want to apply.. Defaults to all actions.

        Returns:
            str: Text with augmentation.
        """
        if not actions:
            actions = self.__actions
        action = np.random.choice(actions)
        new = self.augment(text, action)

        return new

    def __call__(self, text: str, actions: List[str] = []) -> str:
        new = self.random(text, actions)

        return new
