"""Pretrain dataset for HuggingFace datasets."""

from typing import Dict
import torch
from taoTrain.config import TrainingConfig
from taoTrain.data.hf_base import BaseHFDataset


class PretrainDataset(BaseHFDataset):
    """Dataset for pretraining with raw text."""
    
    def _preprocess(self):
        """Tokenize text data."""
        dataset_config = self.config.dataset
        text_column = dataset_config.text_column
        
        def tokenize_function(examples):
            # Concatenate all texts
            concatenated_examples = {
                k: sum(examples[k], []) for k in examples.keys()
            }
            
            total_length = len(concatenated_examples[text_column])
            # We'll use max_seq_length for training
            total_length = (total_length // self.config.model.max_seq_length) * self.config.model.max_seq_length
            
            # Tokenize
            tokenized = self.tokenizer(
                concatenated_examples[text_column],
                truncation=False,  # We'll chunk below
                return_special_tokens_mask=False,
            )
            
            # Chunk tokenized text
            result = {
                "input_ids": [],
                "attention_mask": [],
            }
            
            for i in range(0, total_length, self.config.model.max_seq_length):
                result["input_ids"].append(
                    tokenized["input_ids"][i:i + self.config.model.max_seq_length]
                )
                result["attention_mask"].append(
                    tokenized["attention_mask"][i:i + self.config.model.max_seq_length]
                )
            
            return result
        
        # Preprocess in batches
        self.data = self.data.map(
            tokenize_function,
            batched=True,
            batch_size=100,
            remove_columns=self.data.column_names,
            desc="Tokenizing...",
        )
    
    def __getitem__(self, idx: int) -> Dict[str, torch.Tensor]:
        """Get preprocessed sample."""
        item = self.data[idx]
        
        input_ids = torch.tensor(item["input_ids"], dtype=torch.long)
        attention_mask = torch.tensor(item["attention_mask"], dtype=torch.long)
        
        # For pretrain, labels = input_ids shifted by 1 (next token prediction)
        # Position i predicts token at position i+1
        labels = input_ids[1:].clone()
        labels = torch.cat([labels, torch.tensor([-100])], dim=0)
        
        # Mark padding tokens as -100 to ignore in loss computation
        labels[attention_mask == 0] = -100
        
        return {
            "input_ids": input_ids,
            "attention_mask": attention_mask,
            "labels": labels,
        }