import os
import pickle
from typing import BinaryIO

import regex as re

PAT = re.compile(
    r"""'(?:[sdmt]|ll|ve|re)| ?\p{L}+| ?\p{N}+| ?[^\s\p{L}\p{N}]+|\s+(?!\S)|\s+"""  
)



def find_chunk_boundaries(
    file: BinaryIO,
    desired_num_chunks: int,
    split_special_token: bytes,
) -> list[int]:
    """
    Chunk the file into parts that can be counted independently.
    May return fewer chunks if the boundaries end up overlapping.
    """
    assert isinstance(split_special_token, bytes), "Must represent special token as a bytestring"

    # Get total file size in bytes
    file.seek(0, os.SEEK_END)
    file_size = file.tell()
    file.seek(0)

    chunk_size = file_size // desired_num_chunks

    # Initial guesses for chunk boundary locations, uniformly spaced
    # Chunks start on previous index, don't include last index
    chunk_boundaries = [i * chunk_size for i in range(desired_num_chunks + 1)]
    chunk_boundaries[-1] = file_size

    mini_chunk_size = 4096  # Read ahead by 4k bytes at a time

    for bi in range(1, len(chunk_boundaries) - 1):
        initial_position = chunk_boundaries[bi]
        file.seek(initial_position)  # Start at boundary guess
        while True:
            mini_chunk = file.read(mini_chunk_size)  # Read a mini chunk

            # If EOF, this boundary should be at the end of the file
            if mini_chunk == b"":
                chunk_boundaries[bi] = file_size
                break

            # Find the special token in the mini chunk
            found_at = mini_chunk.find(split_special_token)
            if found_at != -1:
                chunk_boundaries[bi] = initial_position + found_at
                break
            initial_position += mini_chunk_size

    # Make sure all boundaries are unique, but might be fewer than desired_num_chunks
    return sorted(set(chunk_boundaries))

# Usage
def run_train_bpe(
        path: str | os.PathLike, vocab_size: int, special_tokens: list[str]
        ) -> tuple[dict[int, bytes], list[tuple[bytes, bytes]]]:
    # initialize vocab
    vocab = {}
    for i in range(0,256):
        vocab[i] = bytes([i]) 
    for i in range (0,len(special_tokens)):
        special_token = special_tokens[i]
        vocab[256 + i] = special_token.encode("utf-8",errors = "ignore")


    with open(path, "rb") as f:
        num_processes = 4
        boundaries = find_chunk_boundaries(f, num_processes, b"<|endoftext|>")

        counts: dict[tuple[int, ...], int] = {}
        merges: list[tuple[bytes,bytes]] = []

     # The following is a serial implementation, but you can parallelize this
        # by sending each start/end pair to a set of processes.
        for start, end in zip(boundaries[:-1], boundaries[1:]):
            f.seek(start)
            chunk = f.read(end - start).decode("utf-8", errors="ignore")
            chunk_part = [chunk]
            for special_token in special_tokens:
                new_part = []
                for part in chunk_part:
                    new_part.extend(part.split(special_token))
                chunk_part = new_part

            # Run pre-tokenization on your chunk and store the counts for each pre-token
            for chunk in chunk_part:
                for match in PAT.finditer(chunk):
                    raw = match.group().encode("utf-8")
                    pieces = tuple(raw)
                    counts[pieces] = counts.get(pieces,0) + 1    

        # we use the extend-update algo:
        # 1. calculate pair_count and 
        pair_count: dict[tuple[int,int], int] = {}
        pair_affected: dict[tuple[int,int], set] = {}
        for token_seq, num in counts.items():
            for i in range (0, len(token_seq) - 1):
                pair = (token_seq[i], token_seq[i + 1])
                pair_count[pair] = pair_count.get(pair,0) + num
                if pair not in pair_affected:
                    pair_affected[pair] = set()
                pair_affected[pair].add(token_seq) 

        # merge step 
        while len(vocab) < vocab_size:    
            best_pair, _ = max(pair_count.items(), key=lambda item: ( item[1],(vocab[item[0][0]], vocab[item[0][1]])))
            
            affected_token_seq = pair_affected[best_pair]

            merged_bytes = vocab[best_pair[0]] + vocab[best_pair[1]]
            
            new_token_id = len(vocab)
            vocab[new_token_id] = merged_bytes
            merges.append((vocab[best_pair[0]],vocab[best_pair[1]]))

            for token_seq in tuple(affected_token_seq):
                if token_seq not in counts:
                    continue 

                i = 0
                parts = []
                while i < len(token_seq):
                    if i != len(token_seq) - 1 and \
                        token_seq[i] == best_pair[0] and token_seq[i + 1] == best_pair[1]:
                        parts.append(new_token_id)
                        i += 2
                    else:
                        parts.append(token_seq[i])
                        i += 1

                old_sequence_num = counts.pop(token_seq)
                
                # remove every old_pair in this sequence
                # Note: pair_affected's is set
                old_pair_dict = {}
                for i in range (0,len(token_seq)-1):
                    old_pair = (token_seq[i],token_seq[i + 1])
                    old_pair_dict[old_pair] = old_pair_dict.get(old_pair,0) + 1
                for old_pair, num in old_pair_dict.items():
                    pair_affected[old_pair].remove(token_seq)
                    pair_count[old_pair] -= old_sequence_num * num
                    # check if the old_pair should be deleted 
                    if not pair_affected[old_pair]:
                        del pair_affected[old_pair]
                    assert pair_count[old_pair] >= 0,(
                        f"pair_count became negative for {old_pair}: {pair_count[old_pair]}"
                    )
                    if pair_count[old_pair] == 0:
                        del pair_count[old_pair]


                # merge and update affected items.
                new_token_seq = tuple(parts)

                # recalculate pairs
                new_pair_dict = {}
                for i in range (0, len(new_token_seq)-1):
                    new_pair = (new_token_seq[i],new_token_seq[i + 1])
                    new_pair_dict[new_pair] = new_pair_dict.get(new_pair,0) + 1
                for new_pair, num in new_pair_dict.items():     
                    if new_pair not in pair_affected:
                        pair_affected[new_pair] = set()
                    pair_affected[new_pair].add(new_token_seq)
                    pair_count[new_pair] = pair_count.get(new_pair,0) + old_sequence_num * num
                # update counts table. The new_sequence may already in counts, 
                # even though they are not derivated from the same old_token_seq. 
                counts[new_token_seq] = counts.get(new_token_seq,0) + old_sequence_num

            

    return (vocab, merges)


def save_bpe(filepath: str,vocab: dict[int, bytes], merges: list[tuple[bytes,bytes]]):
    with open(filepath, "wb") as f:
        pickle.dump((vocab, merges),f)