# ─────────────────────────────────────────────────────────
# EGen Labs | ErebusTN
# AXIOM-1 — EGen-Core/AXIOM-1
# File: axiom1/training/data.py
# ─────────────────────────────────────────────────────────

"""
Dataset loader and tokenizer setup for AXIOM-1.

Uses a pretrained tokenizer compatible with vocab_size=32000:
  - meta-llama/Llama-2-7b-hf (recommended)
  - Qwen/Qwen2-0.5B (alternative)

Dataset implementation supports:
  - Loading from arrow/jsonl/text files
  - Chunking sequences to fixed length
  - Proper collation with padding
"""

from pathlib import Path
from typing import Iterator

import torch
from torch.utils.data import Dataset, DataLoader, IterableDataset

from axiom1.model.config import AXIOM1Config


def get_tokenizer(name: str = "meta-llama/Llama-2-7b-hf"):
    """
    Load a pretrained tokenizer compatible with AXIOM-1 (vocab_size=32000).

    Args:
        name: HuggingFace tokenizer name or local path

    Returns:
        Tokenizer instance with verified vocab size
    """
    from transformers import AutoTokenizer

    tokenizer = AutoTokenizer.from_pretrained(name)

    if tokenizer.pad_token is None:
        tokenizer.pad_token = tokenizer.eos_token

    assert len(tokenizer) == 32000, \
        f"Expected vocab_size=32000, got {len(tokenizer)}. " \
        f"Use meta-llama/Llama-2-7b-hf or Qwen/Qwen2-0.5B."

    return tokenizer


class TextChunkDataset(Dataset):
    """
    Pre-tokenized text dataset that serves fixed-length chunks.

    Loads a pre-tokenized tensor of token IDs and splits it into
    non-overlapping chunks of seq_len tokens for training.

    Usage:
        # Pre-tokenize your corpus to a .pt file:
        # tokens = tokenizer(text, return_tensors='pt')['input_ids'].squeeze()
        # torch.save(tokens, 'train_tokens.pt')

        dataset = TextChunkDataset('train_tokens.pt', seq_len=512)
        loader  = DataLoader(dataset, batch_size=32, shuffle=True)
    """

    def __init__(self, token_file: str | Path, seq_len: int = 512) -> None:
        """
        Args:
            token_file: Path to .pt file containing 1D tensor of token IDs
            seq_len: Sequence length for each training example
        """
        self.seq_len = seq_len
        self.tokens = torch.load(token_file, map_location='cpu', weights_only=True)

        if self.tokens.ndim > 1:
            self.tokens = self.tokens.flatten()

        # Truncate to exact multiple of seq_len
        n_chunks = len(self.tokens) // seq_len
        self.tokens = self.tokens[:n_chunks * seq_len]
        self.n_chunks = n_chunks

    def __len__(self) -> int:
        return self.n_chunks

    def __getitem__(self, idx: int) -> dict[str, torch.Tensor]:
        start = idx * self.seq_len
        end = start + self.seq_len
        return {'input_ids': self.tokens[start:end]}


class StreamingTextDataset(IterableDataset):
    """
    Memory-efficient streaming dataset for large corpora.

    Reads pre-tokenized files and yields fixed-length chunks without
    loading the entire corpus into memory.
    """

    def __init__(
        self,
        token_files: list[str | Path],
        seq_len: int = 512,
        shuffle_files: bool = True
    ) -> None:
        """
        Args:
            token_files: List of paths to .pt files containing token ID tensors
            seq_len: Sequence length for each training example
            shuffle_files: Whether to shuffle file order each epoch
        """
        self.token_files = [Path(f) for f in token_files]
        self.seq_len = seq_len
        self.shuffle_files = shuffle_files

    def __iter__(self) -> Iterator[dict[str, torch.Tensor]]:
        files = list(self.token_files)
        if self.shuffle_files:
            import random
            random.shuffle(files)

        buffer = torch.tensor([], dtype=torch.long)

        for fpath in files:
            tokens = torch.load(fpath, map_location='cpu', weights_only=True)
            if tokens.ndim > 1:
                tokens = tokens.flatten()
            buffer = torch.cat([buffer, tokens])

            while len(buffer) >= self.seq_len:
                chunk = buffer[:self.seq_len]
                buffer = buffer[self.seq_len:]
                yield {'input_ids': chunk}


def build_dataloader(
    token_file: str | Path,
    cfg: AXIOM1Config,
    seq_len: int,
    batch_size: int,
    shuffle: bool = True,
    num_workers: int = 4
) -> DataLoader:
    """
    Build a DataLoader for AXIOM-1 training.

    Args:
        token_file: Path to pre-tokenized .pt file
        cfg: AXIOM1Config instance
        seq_len: Sequence length for this training phase
        batch_size: Batch size for this training phase
        shuffle: Whether to shuffle data
        num_workers: Number of data loading workers

    Returns:
        Configured DataLoader
    """
    dataset = TextChunkDataset(token_file, seq_len=seq_len)
    return DataLoader(
        dataset,
        batch_size=batch_size,
        shuffle=shuffle,
        num_workers=num_workers,
        pin_memory=True,
        drop_last=True
    )
