from kithara.dataset.utils import HFtokenize
import ray
import numpy as np
from kithara.dataset.utils import HFtokenize
from typing import Dict, Any, Optional
from kithara.model import ModelImplementationType
from kithara.dataset.text_completion import TextCompletionDataset
from transformers import AutoTokenizer


class SFTDataset(TextCompletionDataset):
    """A dataset class for Supervised Fine-Tuning (SFT) tasks.

    Args:
        source (ray.data.Dataset): The source Ray dataset containing the training data.
        tokenizer (Optional[AutoTokenizer]): HuggingFace tokenizer instance.
        tokenizer_handle (Optional[str]): Handle/name of the tokenizer to load if not provided.
        column_mapping (Optional[Dict[str, str]]): Mapping of source column names to expected
            column names ("prompt" and "answer").
        model_type (Optional[ModelImplementationType]): Type of model implementation to use.
            Please specify model_type or set MODEL_IMPLEMENTATION in global state. Global
            state is automatically set upon model initialization. Supported types:
            ModelImplementationType.KERASHUB, ModelImplementationType.MAXTEXT
        max_seq_len (int): Maximum sequence length for tokenization (default: 1024). Sequences
            will be padded to this length.
    """

    def __init__(
        self,
        source: ray.data.Dataset,
        tokenizer: Optional[AutoTokenizer] = None,
        tokenizer_handle: Optional[str] = None,
        column_mapping: Optional[Dict[str, str]] = None,
        model_type: Optional["ModelImplementationType"] = None,
        max_seq_len: int = 1024,
    ):
        super().__init__(
            source=source,
            tokenizer=tokenizer,
            tokenizer_handle=tokenizer_handle,
            model_type=model_type,
            max_seq_len=max_seq_len,
        )
        self.column_mapping = {"prompt": "prompt", "answer": "answer"}
        if column_mapping:
            self.column_mapping = {**self.column_mapping, **column_mapping}

    def task_transform(self, sample: Dict[str, str]) -> Dict[str, str]:
        """Transform the raw sample into a standardized prompt-answer format.

        Args:
            sample (Dict[str, str]): Raw sample containing prompt and answer.

        Returns:
            Dict[str, str]: Transformed sample with standardized keys.

        """
        return {
            "prompt": sample[self.column_mapping["prompt"]],
            "answer": sample[self.column_mapping["answer"]],
        }

    def model_transform(self, sample: Dict[str, str]) -> Dict[str, np.ndarray]:
        """Transform the prompt-answer pair into model inputs.

        Args:
            sample (Dict[str, str]): Sample containing prompt and answer.

        Returns:
            Tuple[np.ndarray, np.ndarray, np.ndarray]: Tuple containing input_ids,
            attention_mask, and label_ids.

        """
        prompt, answer = sample["prompt"], sample["answer"]
        full_seq = HFtokenize(
            f"<bos>{prompt}{answer}<eos>", self.tokenizer, seq_len=self.max_seq_len
        )
        prompt_seq = HFtokenize(
            f"<bos>{prompt}",
            self.tokenizer,
            seq_len=self.max_seq_len,
            padding="do_not_pad",
        )
        num_prompt_tokens = len(prompt_seq["input_ids"][0])

        input_ids = full_seq["input_ids"]  # [1, S]
        attention_mask = full_seq["attention_mask"]

        label_ids = np.roll(input_ids, -1)
        label_ids[:, -1] = self.tokenizer.pad_token_id
        label_ids[:, : num_prompt_tokens - 1] = self.tokenizer.pad_token_id
        return input_ids, attention_mask, label_ids
