from __future__ import annotations

from abc import ABC, abstractmethod
from collections.abc import Iterable

from transformers import AutoTokenizer, PreTrainedTokenizerBase

ENGLISH_STOP_WORDS = [
    "!",
    '"',
    "''",
    "``",
    "#",
    "$",
    "%",
    "&",
    "'",
    "(",
    ")",
    "*",
    "+",
    ",",
    "-",
    ".",
    "/",
    ":",
    ";",
    "<",
    "=",
    ">",
    "?",
    "@",
    "[",
    "\\",
    "]",
    "^",
    "_",
    "`",
    "{",
    "|",
    "}",
    "~",
    "a",
    "about",
    "above",
    "across",
    "after",
    "afterwards",
    "again",
    "against",
    "ain",
    "all",
    "almost",
    "alone",
    "along",
    "already",
    "also",
    "although",
    "always",
    "am",
    "among",
    "amongst",
    "amoungst",  # codespell:ignore
    "amount",
    "an",
    "and",
    "another",
    "any",
    "anyhow",
    "anyone",
    "anything",
    "anyway",
    "anywhere",
    "are",
    "aren",
    "around",
    "as",
    "at",
    "back",
    "be",
    "became",
    "because",
    "become",
    "becomes",
    "becoming",
    "been",
    "before",
    "beforehand",
    "behind",
    "being",
    "below",
    "beside",
    "besides",
    "between",
    "beyond",
    "bill",
    "both",
    "bottom",
    "but",
    "by",
    "call",
    "can",
    "cannot",
    "cant",  # codespell:ignore
    "co",
    "con",
    "could",
    "couldn",
    "couldnt",
    "cry",
    "d",
    "de",
    "describe",
    "detail",
    "did",
    "didn",
    "do",
    "does",
    "doesn",
    "doing",
    "don",
    "done",
    "down",
    "due",
    "during",
    "each",
    "eg",
    "eight",
    "either",
    "eleven",
    "else",
    "elsewhere",
    "empty",
    "enough",
    "etc",
    "even",
    "ever",
    "every",
    "everyone",
    "everything",
    "everywhere",
    "except",
    "few",
    "fifteen",
    "fifty",
    "fill",
    "find",
    "fire",
    "first",
    "five",
    "for",
    "former",
    "formerly",
    "forty",
    "found",
    "four",
    "from",
    "front",
    "full",
    "further",
    "get",
    "give",
    "go",
    "had",
    "hadn",
    "has",
    "hasn",
    "hasnt",
    "have",
    "haven",
    "having",
    "he",
    "hence",
    "her",
    "here",
    "hereafter",
    "hereby",
    "herein",
    "hereupon",
    "hers",
    "herself",
    "him",
    "himself",
    "his",
    "how",
    "however",
    "hundred",
    "i",
    "ie",
    "if",
    "in",
    "inc",
    "indeed",
    "interest",
    "into",
    "is",
    "isn",
    "it",
    "its",
    "itself",
    "just",
    "keep",
    "last",
    "latter",
    "latterly",
    "least",
    "less",
    "ll",
    "ltd",
    "m",
    "ma",
    "made",
    "many",
    "may",
    "me",
    "meanwhile",
    "might",
    "mightn",
    "mill",
    "mine",
    "more",
    "moreover",
    "most",
    "mostly",
    "move",
    "much",
    "must",
    "mustn",
    "my",
    "myself",
    "name",
    "namely",
    "needn",
    "neither",
    "never",
    "nevertheless",
    "next",
    "nine",
    "no",
    "nobody",
    "none",
    "noone",  # codespell:ignore
    "nor",
    "not",
    "nothing",
    "now",
    "nowhere",
    "o",
    "of",
    "off",
    "often",
    "on",
    "once",
    "one",
    "only",
    "onto",
    "or",
    "other",
    "others",
    "otherwise",
    "our",
    "ours",
    "ourselves",
    "out",
    "over",
    "own",
    "part",
    "per",
    "perhaps",
    "please",
    "put",
    "rather",
    "re",
    "s",
    "same",
    "see",
    "seem",
    "seemed",
    "seeming",
    "seems",
    "serious",
    "several",
    "shan",
    "she",
    "should",
    "shouldn",
    "show",
    "side",
    "since",
    "sincere",
    "six",
    "sixty",
    "so",
    "some",
    "somehow",
    "someone",
    "something",
    "sometime",
    "sometimes",
    "somewhere",
    "still",
    "such",
    "system",
    "t",
    "take",
    "ten",
    "than",
    "that",
    "the",
    "their",
    "theirs",
    "them",
    "themselves",
    "then",
    "thence",
    "there",
    "thereafter",
    "thereby",
    "therefore",
    "therein",
    "thereupon",
    "these",
    "they",
    "thick",
    "thin",
    "third",
    "this",
    "those",
    "though",
    "three",
    "through",
    "throughout",
    "thru",
    "thus",
    "to",
    "together",
    "too",
    "top",
    "toward",
    "towards",
    "twelve",
    "twenty",
    "two",
    "un",
    "under",
    "until",
    "up",
    "upon",
    "us",
    "ve",
    "very",
    "via",
    "was",
    "wasn",  # codespell:ignore
    "we",
    "well",
    "were",
    "weren",
    "what",
    "whatever",
    "when",
    "whence",
    "whenever",
    "where",
    "whereafter",
    "whereas",
    "whereby",
    "wherein",
    "whereupon",
    "wherever",
    "whether",
    "which",
    "while",
    "whither",
    "who",
    "whoever",
    "whole",
    "whom",
    "whose",
    "why",
    "will",
    "with",
    "within",
    "without",
    "won",
    "would",
    "wouldn",
    "y",
    "yet",
    "you",
    "your",
    "yours",
    "yourself",
    "yourselves",
]


class WordTokenizer(ABC):
    @abstractmethod
    def set_vocab(self, vocab: Iterable[str]):
        pass

    @abstractmethod
    def get_vocab(self, vocab: Iterable[str]):
        pass

    @abstractmethod
    def tokenize(self, text: str, **kwargs) -> list[int]:
        pass

    @abstractmethod
    def save(self, output_path: str):
        pass

    @staticmethod
    @abstractmethod
    def load(input_path: str):
        pass


class TransformersTokenizerWrapper(WordTokenizer):
    def __init__(self, tokenizer: PreTrainedTokenizerBase):
        super().__init__()
        self.tokenizer = tokenizer

    def tokenize(self, sentence: str):
        encoded = self.tokenizer(sentence)
        return encoded["input_ids"][0]

    def set_vocab(self, vocab: Iterable[str]):
        pass

    def get_vocab(self, vocab: Iterable[str]):
        return self.tokenizer.get_vocab()

    def save(self, output_path: str):
        self.tokenizer.save_pretrained(output_path)

    @staticmethod
    def load(input_path: str):
        return TransformersTokenizerWrapper(AutoTokenizer.from_pretrained(input_path, use_fast=True))
