dataset.py

import torch
from torch.utils.data import Dataset
import os
import glob
from tqdm import tqdm
import json
import shutil
import numpy as np
from .parser import LogTokenizer
from .vocab_manager import VocabManager

# --- Multiprocessing top-level functions ---
_worker_tokenizer = None


def _init_worker():
    global _worker_tokenizer
    _worker_tokenizer = LogTokenizer(VocabManager())


def _process_file(fpath, format_id, min_elo, max_elo, context_length, chunk_len):
    try:
        global _worker_tokenizer
        if format_id and not os.path.basename(fpath).startswith(f"{format_id}-"):
            return None

        # Fast metadata extraction without tokenizer
        meta = {"winner": None, "p1_rating": 1000, "p2_rating": 1000}
        p1_name = None
        p2_name = None

        with open(fpath, "r") as f:
            for line in f:
                parts = line.strip().split("|")
                if len(parts) < 2:
                    continue
                cmd = parts[1]

                if cmd == "player":
                    if len(parts) > 2:
                        pid = parts[2]
                        if len(parts) > 3:
                            name = parts[3]
                            if pid == "p1":
                                p1_name = name
                            elif pid == "p2":
                                p2_name = name

                        if len(parts) > 5:
                            try:
                                rating = float(parts[5])
                                if pid == "p1":
                                    meta["p1_rating"] = rating
                                elif pid == "p2":
                                    meta["p2_rating"] = rating
                            except:
                                pass

                elif cmd == "win":
                    if len(parts) > 2:
                        winner_name = parts[2]
                        if winner_name == p1_name:
                            meta["winner"] = "p1"
                        elif winner_name == p2_name:
                            meta["winner"] = "p2"

        p1_elo = meta.get("p1_rating", 1000) or 1000
        p2_elo = meta.get("p2_rating", 1000) or 1000
        avg_elo = (p1_elo + p2_elo) / 2

        if min_elo and avg_elo < min_elo:
            return None
        if max_elo and avg_elo > max_elo:
            return None

        winner = meta.get("winner")
        if winner is None:
            return None

        # Use global worker tokenizer
        tokens = _worker_tokenizer.tokenize_replay(fpath)
        win_label = 1.0 if winner == "p1" else 0.0
        weight = max(0.2, avg_elo / 1500.0)

        current_file_samples = []
        pad_token = _worker_tokenizer.vm.vocab[_worker_tokenizer.vm.PAD_TOKEN]

        if len(tokens) < chunk_len:
            padding = [pad_token] * (chunk_len - len(tokens))
            current_file_samples.append(tokens + padding)
        else:
            stride = context_length // 2
            for i in range(0, len(tokens) - context_length, stride):
                chunk = tokens[i : i + chunk_len]
                if len(chunk) == chunk_len:
                    current_file_samples.append(chunk)

        return {
            "samples": current_file_samples,
            "win_label": win_label,
            "weight": weight,
        }
    except Exception as e:
        return None


class PokemonShowdownDataset(Dataset):
    def __init__(
        self,
        data_dir: str,
        context_length: int = 512,
        min_elo: int = None,
        max_elo: int = None,
        format_id: str = None,
        cached: bool = False,
        file_names: set[str] = None,
        max_files: int = None,
    ):
        self.data_dir = data_dir
        self.context_length = context_length
        self.min_elo = min_elo
        self.max_elo = max_elo
        self.format_id = format_id
        self.file_names = file_names
        self.max_files = max_files

        self.vm = VocabManager()
        if not os.path.exists(os.path.join(self.vm.cache_dir, "vocab.json")):
            self.vm.build_vocab()
        else:
            self.vm.load_vocab()

        self.tokenizer = LogTokenizer(self.vm)
        self.files = glob.glob(os.path.join(data_dir, "*.log"))

        if self.file_names:
            self.files = [
                f for f in self.files if os.path.basename(f) in self.file_names
            ]

        if self.max_files and len(self.files) > self.max_files:
            import random

            # Fix the random seed for reproducible sampling of the dataset
            random.seed(42)
            self.files = random.sample(self.files, self.max_files)

        # --- CACHE LOGIC: MMAP DIRECTORY ---
        cache_name = f"mmap_ctx{context_length}_min{min_elo}_max{max_elo}_{format_id}_maxf{max_files}"
        self.cache_path = os.path.join(self.vm.cache_dir, cache_name)

        self.use_mmap = False

        # Check if cache is valid (folder exists and has metadata)
        if (
            cached
            and os.path.exists(self.cache_path)
            and os.path.exists(os.path.join(self.cache_path, "metadata.json"))
        ):
            print(f"Loading cached dataset (mmap) from {self.cache_path}...")
            try:
                self._load_mmap()
                self.use_mmap = True
                print("Cache loaded successfully (Memory Mapped).")
            except Exception as e:
                print(f"Failed to load mmap cache: {e}. Rebuilding...")
                shutil.rmtree(self.cache_path, ignore_errors=True)

        if not self.use_mmap:
            if not self.files:
                print(
                    f"Warning: No .log files found in {data_dir}. Use the downloader first."
                )

            # Prepare data and write to disk incrementally
            self._build_mmap_incrementally()
            self.use_mmap = True

    def _build_mmap_incrementally(self):
        print(f"Building dataset cache at {self.cache_path}...")
        os.makedirs(self.cache_path, exist_ok=True)

        temp_chunks_path = os.path.join(self.cache_path, "temp_chunks.bin")
        temp_wins_path = os.path.join(self.cache_path, "temp_wins.bin")
        temp_weights_path = os.path.join(self.cache_path, "temp_weights.bin")

        # Open temp files for appending
        f_chunks = open(temp_chunks_path, "wb")
        f_wins = open(temp_wins_path, "wb")
        f_weights = open(temp_weights_path, "wb")

        total_samples = 0
        chunk_len = self.context_length + 1

        # Buffer to reduce I/O calls
        buffer_chunks = []
        buffer_wins = []
        buffer_weights = []
        BUFFER_SIZE = 10000

        print(f"Processing {len(self.files)} replay files...")

        from concurrent.futures import ProcessPoolExecutor
        import multiprocessing
        from functools import partial

        process_func = partial(
            _process_file,
            format_id=self.format_id,
            min_elo=self.min_elo,
            max_elo=self.max_elo,
            context_length=self.context_length,
            chunk_len=chunk_len,
        )

        with ProcessPoolExecutor(
            max_workers=multiprocessing.cpu_count(), initializer=_init_worker
        ) as executor:
            # Use map with chunksize instead of submit to avoid building millions of futures in RAM at once
            # A large chunksize (5000) creates massive IPC payloads that cause OOM crashes with millions of files.
            # We use a much smaller chunksize (100) and limit the prefetch to keep memory usage low.
            results = executor.map(process_func, self.files, chunksize=100)

            for res in tqdm(results, total=len(self.files), desc="Processing"):
                if res and res["samples"]:
                    for s in res["samples"]:
                        buffer_chunks.append(s)
                        buffer_wins.append(res["win_label"])
                        buffer_weights.append(res["weight"])
                        total_samples += 1

                    if len(buffer_chunks) >= BUFFER_SIZE:
                        np.array(buffer_chunks, dtype=np.int32).tofile(f_chunks)
                        np.array(buffer_wins, dtype=np.float32).tofile(f_wins)
                        np.array(buffer_weights, dtype=np.float32).tofile(f_weights)

                        buffer_chunks = []
                        buffer_wins = []
                        buffer_weights = []

        # Flush remaining buffer
        if buffer_chunks:
            np.array(buffer_chunks, dtype=np.int32).tofile(f_chunks)
            np.array(buffer_wins, dtype=np.float32).tofile(f_wins)
            np.array(buffer_weights, dtype=np.float32).tofile(f_weights)

        f_chunks.close()
        f_wins.close()
        f_weights.close()

        print(f"Total samples collected: {total_samples}")

        if total_samples == 0:
            # Handle empty dataset case
            with open(os.path.join(self.cache_path, "metadata.json"), "w") as f:
                json.dump({"num_samples": 0, "chunk_len": chunk_len}, f)
            return

        # --- Finalize: Create proper memory maps ---
        # We rename the temp raw binary files to final .bin files.
        # Since we used `tofile` (C-order raw write), the format matches `memmap` exactly.

        os.rename(temp_chunks_path, os.path.join(self.cache_path, "chunks.bin"))
        os.rename(temp_wins_path, os.path.join(self.cache_path, "wins.bin"))
        os.rename(temp_weights_path, os.path.join(self.cache_path, "weights.bin"))

        # Save metadata
        self.num_samples = total_samples
        with open(os.path.join(self.cache_path, "metadata.json"), "w") as f:
            json.dump({"num_samples": total_samples, "chunk_len": chunk_len}, f)

        # Load them up
        self._load_mmap()

    def _load_mmap(self):
        with open(os.path.join(self.cache_path, "metadata.json"), "r") as f:
            meta = json.load(f)

        self.num_samples = meta["num_samples"]
        chunk_len = meta["chunk_len"]

        if self.num_samples > 0:
            # Open in read-only mode (copy-on-write)
            self.data_chunks = np.memmap(
                os.path.join(self.cache_path, "chunks.bin"),
                dtype="int32",
                mode="r",
                shape=(self.num_samples, chunk_len),
            )
            self.data_wins = np.memmap(
                os.path.join(self.cache_path, "wins.bin"),
                dtype="float32",
                mode="r",
                shape=(self.num_samples,),
            )
            self.data_weights = np.memmap(
                os.path.join(self.cache_path, "weights.bin"),
                dtype="float32",
                mode="r",
                shape=(self.num_samples,),
            )
        else:
            self.data_chunks = []

    def __len__(self):
        if not self.use_mmap:
            return 0
        return self.num_samples

    def get_vocab_size(self):
        return len(self.vm)

    def __getitem__(self, idx):
        if not self.use_mmap:
            raise RuntimeError("Dataset not initialized correctly (empty or no mmap).")

        # Read from numpy mmap (this does disk IO if not in page cache, but doesn't fill VRAM)
        chunk_np = self.data_chunks[idx]
        win_val = self.data_wins[idx]
        weight_val = self.data_weights[idx]

        # Convert to Tensor (this copies data into standard RAM, small memory footprint per batch)
        # We perform copy() to ensure we have a writable, standard array before tensor conversion
        # This prevents "UserWarning: The given NumPy array is not writeable"
        chunk = torch.from_numpy(chunk_np.copy()).long()
        win_label = torch.tensor(win_val, dtype=torch.float)
        weight = torch.tensor(weight_val, dtype=torch.float)

        # X: 0..N-1
        # Y: 1..N
        x = chunk[:-1]
        y = chunk[1:]

        return x, y, win_label, weight