"""
BPETrainer: learn merge rules from raw text.

Note: this recounts all pairs each iteration (clear but 0(merges x tokens)).
Fine for today's small demo. For tomorrow's full TinyStories run we'll add
incremental pair-count updates so it scales.
"""

import os
import json
from collections import Counter, defaultdict

from .regex_split import split
from .tokenizer import merge_seq
from .special_tokens import SPECIAL_TOKENS, special_token_ids


def _pairs(ids):
    """Counter of adjacent id pairs within one word."""
    c = Counter()
    for a, b in zip(ids, ids[1:]):
        c[(a, b)] += 1
    return c

class BPETrainer:
    def __init__(self, vocab_size=16384, verbose=False):
        self.vocab_size = vocab_size
        self.verbose = verbose
        self.merges = []
        self.vocab = {}

    # ---- entry points ----
    def train(self, text):
        """Train from an in-memory string *used by the test)."""
        return self._fit(Counter(split(text)))
    
    def train_from_file(self, path):
        """Stream a (possibly multi-GB) text file line by line so we never
        hold the whole corpus in memory."""
        word_freqs = Counter()
        with open(path, "r", encoding="utf-8") as f:
            for i, line in enumerate(f):
                if "<|endoftext|>" in line:
                    line = line.replace("<|endoftext|>", " ")
                word_freqs.update(split(line))
                if self.verbose and i and i % 200000 == 0:
                    print(f"  read {i:,} lines, {len(word_freqs):,} unique chunks")
        if self.verbose:
            print(f"  done reading: {len(word_freqs):,} unique chunks")
        return self._fit(word_freqs)
    
    # ---- core training ----
    def _fit(self, word_freqs):
        num_merges = self.vocab_size - 256 - len(SPECIAL_TOKENS)
        if num_merges < 0:
            raise ValueError("vocab_size too small for 256 + specials")
        
        # each unique word -> list of byte ids
        splits = {w: list(w.encode("utf-8")) for w in word_freqs}

        # running pair counts (weighted by word freq) + which words hold each pair
        pair_counts = Counter()
        where = defaultdict(set)
        for w, freq in word_freqs.items():
            for p, c in _pairs(splits[w]).items():
                pair_counts[p] += c * freq
                where[p].add(w)

        merges = []
        vocab = {i: bytes([i]) for i in range(256)}

        for m in range(num_merges):
            if not pair_counts:
                break
            best = max(pair_counts, key=pair_counts.get)
            if pair_counts[best] <= 0:
                break

            new_id = 256 + m
            merges.append(best)
            vocab[new_id] = vocab[best[0]] + vocab[best[1]]

            # only the words containing 'best' can change
            for w in list(where[best]):
                freq = word_freqs[w]
                old = splits[w]
                before = _pairs(old)
                new = merge_seq(old, best, new_id)
                splits[w] = new
                after = _pairs(new)

                for p in set(before) | set(after):
                    delta = (after[p] - before[p]) * freq
                    if delta:
                        pair_counts[p] += delta
                    if after[p] > 0:
                        where[p].add(w)
                    else:
                        where[p].discard(w)
                    if pair_counts.get(p, 0) <= 0:
                        pair_counts.pop(p, None)
                        where.pop(p, None)
            
            if self.verbose and (m < 20 or m% 500 == 0):
                piece = vocab[new_id].decode("utf-8", errors="replace")
                print(f"merge {m:5d}: {best} -> id {new_id}  {piece!r}")
        
        self.merges = merges
        self.vocab = vocab
        return merges, vocab
    
    # ---- save ----
    def save(self, out_dir):
        os.makedirs(out_dir, exist_ok=True)
        with open(os.path.join(out_dir, "merges.txt"), "w", encoding="utf-8") as f:
            for a, b in self.merges:
                f.write(f"{a} {b}\n")

        # human-readable view (not used for loading, just for inspection)
        view = {i: b.decode("utf-8", errors="replace")
                for i, b in self.vocab.items()}
        for tok, tid in special_token_ids(self.vocab_size).items():
            view[tid] = tok
        with open(os.path.join(out_dir, "vocab.json"), "w", encoding="utf-8") as f:
            json.dump(view, f, ensure_ascii=False, indent=0)

        with open(os.path.join(out_dir, "config.json"), "w", encoding="utf-8") as f:
            json.dump({"vocab_size": self.vocab_size}, f)