from accelerate.logging import get_logger
logger = get_logger(__name__, log_level="INFO")


import torch
class UniversalPrompting():
    def __init__(self, text_tokenizer,
                 max_prompt_len=8000, max_gen_length=377, ignore_id=-100, system_prompt=None):
        """
        :param text_tokenizer: original text tokenizer
        """
        self.text_tokenizer = text_tokenizer
        self.max_gen_length = max_gen_length
        self.max_prompt_len = max_prompt_len

    def lm_prompt(self, text_ids_pairs):
        prompts_list, responses_list = text_ids_pairs
        pad_id = self.text_tokenizer.pad_token_id

        if responses_list.shape[1] < self.max_gen_length:
            max_seq_len = prompts_list.shape[1] + responses_list.shape[1]
        else:
            max_seq_len = prompts_list.shape[1] + self.max_gen_length

        sequence_ids = []
        attention_masks = []
        label_ids = []

        for prompt_ids, resp_ids in zip(prompts_list, responses_list):
            prompt_ids = prompt_ids.tolist()
            resp_ids   = resp_ids.tolist()

            temp_ids = prompt_ids + resp_ids
            temp_masks = [1] * len(temp_ids)
            temp_labels = temp_ids.copy()

            if len(temp_ids) < max_seq_len:
                pad_len = max_seq_len - len(temp_ids)
                temp_ids.extend([pad_id] * pad_len)
                temp_labels.extend([pad_id] * pad_len)
                temp_masks.extend([0] * pad_len)
            else:
                temp_ids = temp_ids[:max_seq_len]
                temp_labels = temp_labels[:max_seq_len]
                temp_masks = temp_masks[:max_seq_len]

            sequence_ids.append(torch.tensor(temp_ids).unsqueeze(0))
            attention_masks.append(torch.tensor(temp_masks).unsqueeze(0))
            label_ids.append(torch.tensor(temp_labels).unsqueeze(0))

        input_ids = torch.cat(sequence_ids, dim=0)
        attention_masks = torch.cat(attention_masks, dim=0)
        label_ids = torch.cat(label_ids, dim=0)
        

        return input_ids, label_ids, prompts_list.shape[1]

        
    

    def mask_prompt(self):
        pass

    def __call__(self, input):
        prompts, responses = input

        enc = self.text_tokenizer(
            prompts,
            padding=False,
            truncation=False,
            return_length=True
        )
        lengths = enc["length"]
        keep_indices = [i for i, L in enumerate(lengths) if L <= self.max_prompt_len]
        drop_num = len(prompts) - len(keep_indices)
        
        prompts = [self.text_tokenizer.apply_chat_template(
            [{"role": "user", "content": prompts[i]}],
            tokenize=False,
            add_generation_prompt=True
        ) for i in keep_indices]
        responses = [responses[i] for i in keep_indices]

        prompt_ids = self.text_tokenizer(
            prompts,
            padding=True,
            return_tensors="pt",
            padding_side = "left"
        )['input_ids']
        response_ids = self.text_tokenizer(
            responses,
            padding=True,
            return_tensors="pt",
            padding_side = "right"
        )['input_ids']

        input_ids_lm, labels_lm, start_pos = self.lm_prompt((prompt_ids, response_ids))
        return input_ids_lm, labels_lm, start_pos, drop_num


if __name__ == '__main__':
    pass