import os
import time
import json
import argparse
import multiprocessing
from collections import defaultdict
from tqdm import tqdm

def get_stats_chunk(ids):
    counts = defaultdict(int)
    for pair in zip(ids, ids[1:]):
        counts[pair] += 1
    return counts


def merge_chunk(args):
    ids, pair, idx = args
    new_ids = []
    i = 0
    while i < len(ids):
        if i < len(ids) - 1 and ids[i] == pair[0] and ids[i + 1] == pair[1]:
            new_ids.append(idx)
            i += 2
        else:
            new_ids.append(ids[i])
            i += 1
    return new_ids


class ParallelBPETokenizer:
    def __init__(self):
        self.merges = {}
        self.vocab = {i: bytes([i]) for i in range(256)}

    def train(self, text, vocab_size, pct_bpe=1.0, workers=None, verbose=True):
        assert vocab_size >= 256
        num_merges = vocab_size - 256

        if verbose:
            print("Pre-processing text...")
        text_subset = text[: int(len(text) * max(0.01, min(1.0, pct_bpe)))]
        ids = list(text_subset.encode("utf-8"))

        num_procs = workers if workers is not None else max(1, (os.cpu_count() or 4) - 1)
        chunk_len = max(1, len(ids) // num_procs)
        chunks = [ids[i : i + chunk_len] for i in range(0, len(ids), chunk_len)]

        if verbose:
            print(f"Using {num_procs} workers for {len(ids)} bytes...")

        with multiprocessing.Pool(num_procs) as pool:
            for i in tqdm(range(num_merges), desc="Training BPE"):
                chunk_stats = pool.map(get_stats_chunk, chunks)

                totals = defaultdict(int)
                for stat in chunk_stats:
                    for pair, count in stat.items():
                        totals[pair] += count

                if not totals:
                    break

                best_pair = max(totals, key=totals.get)
                new_idx = 256 + i

                merge_args = [(chunk, best_pair, new_idx) for chunk in chunks]
                chunks = pool.map(merge_chunk, merge_args)

                self.merges[best_pair] = new_idx
                self.vocab[new_idx] = self.vocab[best_pair[0]] + self.vocab[best_pair[1]]

                if verbose and i % 20 == 0:
                    try:
                        decoded = self.vocab[new_idx].decode("utf-8")
                        tqdm.write(f"Merged {best_pair} -> {new_idx} ('{decoded}')")
                    except Exception:
                        pass

        if verbose:
            print(f"Training complete. Vocab size: {len(self.vocab)}")
        return self.merges

    def save(self, filename):
        save_merges = {f"{p[0]},{p[1]}": idx for p, idx in self.merges.items()}
        save_vocab = {idx: b.decode("latin1") for idx, b in self.vocab.items()}

        with open(filename, "w", encoding="utf-8") as f:
            json.dump({"merges": save_merges, "vocab": save_vocab}, f)
        print(f"Saved tokenizer to {filename}")


def parse_args():
    p = argparse.ArgumentParser(description="Multi-core BPE trainer")
    p.add_argument("--input", default=os.path.join("data", "jarvis_train.txt"))
    p.add_argument("--vocab-size", type=int, default=2048)
    p.add_argument("--pct-bpe", type=float, default=1.0)
    p.add_argument("--workers", type=int, default=max(1, (os.cpu_count() or 4) - 1))
    p.add_argument("--output", default="bpe_vocab.json")
    return p.parse_args()


if __name__ == "__main__":
    multiprocessing.freeze_support()
    args = parse_args()
    print("--- MULTI-CORE BPE TRAINER ---")

    if not os.path.exists(args.input):
        print(f"Error: {args.input} not found.")
        raise SystemExit(1)

    with open(args.input, "r", encoding="utf-8", errors="ignore") as f:
        text = f.read()

    tokenizer = ParallelBPETokenizer()
    start_time = time.time()
    tokenizer.train(
        text,
        vocab_size=args.vocab_size,
        pct_bpe=args.pct_bpe,
        workers=args.workers,
    )
    end_time = time.time()

    print(f"Total time: {end_time - start_time:.2f} seconds")
    tokenizer.save(args.output)