import re
from collections import Counter

from tqdm import tqdm


class BPETokenizer:
    def __init__(self):
        self.merges: list[tuple[str, str]] = []
        self.vocab: dict[str, int] = {}
        self.inverse_vocab: dict[int, str] = {}
        self.base_tokens: set[str] = set()
        self.unk_token: str = "<unk>"

    def _create_initial_tokens(self, text: str) -> list[str]:
        words = re.findall(r"\w+|[^\w\s]+", text)
        initial_tokens = []
        for word in words:
            for char in word:
                initial_tokens.append(char)
                self.base_tokens.add(char)
            initial_tokens.append("</w>")
        self.base_tokens.add("</w>")
        return initial_tokens

    def train(self, text: str, turns: int = 10, verbose: bool = False) -> None:
        initial_tokens = self._create_initial_tokens(text)
        if turns < 1:
            return

        initial_length = len(initial_tokens)
        if verbose:
            print(f"Original length of vocab = {initial_length}")
        instances_merged = 0

        for t in tqdm(range(turns), desc="Training tokenizer..."):
            if verbose:
                print(f"\nTurn {t + 1}")

            freq = Counter()
            merge_rule_tracker = {}
            i = 0

            while i + 1 < len(initial_tokens):
                if len(initial_tokens[i]) >= 4 and initial_tokens[i][-4:] == "</w>":
                    i += 1
                    continue
                to_check = initial_tokens[i] + initial_tokens[i + 1]
                freq[to_check] += 1
                if to_check not in merge_rule_tracker:
                    merge_rule_tracker[to_check] = (
                        initial_tokens[i],
                        initial_tokens[i + 1],
                    )
                i += 1

            max_freq = max(freq.values(), default=0)
            if max_freq <= 1:
                break

            candidates = [pair for pair in freq if freq[pair] == max_freq]
            to_merge = min(candidates)

            if verbose:
                print(f"\tmax_freq: {max_freq}")
                print(f"\tto_merge: {to_merge}")

            instances_merged += max_freq

            temp_tokens = []
            i = 0
            while i < len(initial_tokens):
                if (
                    i + 1 < len(initial_tokens)
                    and initial_tokens[i] + initial_tokens[i + 1] == to_merge
                ):
                    temp_tokens.append(to_merge)
                    i += 2
                else:
                    temp_tokens.append(initial_tokens[i])
                    i += 1

            self.merges.append(merge_rule_tracker[to_merge])

            if verbose:
                print(f"\tVocab Size at the beginning of turn: {len(initial_tokens)}")
                print(f"\tInstances merged this turn:          {max_freq}")
                print(f"\tVocab Size at the end of turn:       {len(temp_tokens)}")

            initial_tokens = temp_tokens

        if verbose:
            print(
                f"\nNew size of vocab should be = {initial_length - instances_merged}"
            )
            print(f"Actual vocab length =         {len(initial_tokens)}")

        all_tokens = self.base_tokens.copy()
        for part1, part2 in self.merges:
            all_tokens.add(part1 + part2)
        all_tokens.add(self.unk_token)
        sorted_tokens = sorted(all_tokens)
        self.vocab = {token: idx for idx, token in enumerate(sorted_tokens)}
        self.inverse_vocab = {idx: token for token, idx in self.vocab.items()}

    def encode(self, text: str) -> list[int]:
        tokens = []
        words = re.findall(r"\w+|[^\w\s]+", text)
        for word in words:
            for char in word:
                if char in self.base_tokens:
                    tokens.append(char)
                else:
                    tokens.append(self.unk_token)
            tokens.append("</w>")

        for part1, part2 in self.merges:
            merged = part1 + part2
            temp = []
            i = 0
            while i < len(tokens):
                if (
                    i + 1 < len(tokens)
                    and tokens[i] == part1
                    and tokens[i + 1] == part2
                ):
                    temp.append(merged)
                    i += 2
                else:
                    temp.append(tokens[i])
                    i += 1
            tokens = temp

        ids = [self.vocab.get(token, self.vocab[self.unk_token]) for token in tokens]
        return ids

    def decode(self, ids: list[int]) -> str:
        tokens = [self.inverse_vocab.get(id_, self.unk_token) for id_ in ids]
        text = "".join(tokens).replace("</w>", " ")
        return text.strip()
