import logging

import warnings

from typing import Dict, List

import torch

from transformers import PreTrainedTokenizer

class InstructionTuningCollator:

    def __init__(

        self,

        tokenizer: PreTrainedTokenizer,

        sys_prompt: str = "You are an AI assistant specialized in analyzing ICU patient data.",

        ignore_index: int = -100

    ) -> None:

        self.tokenizer = tokenizer

        self.sys_prompt = sys_prompt

        self.ignore_index = ignore_index

        self.response_template, self.response_token_ids = self.infer_response_template()

    def infer_response_template(self):

        logging.warning("Infer response template with v2")

        response_template, response_token_ids = self.infer_response_template_v2()

        if response_template == "":

            logging.warning("Infer response template with v1")

            response_template, response_token_ids = self.infer_response_template_v1()

        return response_template, response_token_ids

    def infer_response_template_v1(self) -> (str, List[int]):

        token = "Hi?"

        chat = [

            {"role": "user", "content": token},

        ]

        formatted_chat = self.tokenizer.apply_chat_template(

            chat,

            tokenize=False,

            add_generation_prompt=True

        )

        response_template = formatted_chat[formatted_chat.find(token) + len(token):]

        response_token_ids = self.tokenizer.encode(response_template, add_special_tokens=False)

        logging.warning(f"Inferred response template: {repr(response_template)}")

        logging.warning(f"Inferred response template token ids: {response_token_ids}")

        return response_template, response_token_ids

    def infer_response_template_v2(self) -> (str, List[int]):

        token = "Hi?"

        chat = [

            {"role": "user", "content": token},

        ]

        formatted_chat_wo_gen = self.tokenizer.apply_chat_template(

            chat,

            tokenize=False,

            add_generation_prompt=False

        )

        formatted_chat = self.tokenizer.apply_chat_template(

            chat,

            tokenize=False,

            add_generation_prompt=True

        )

        formatted_chat_wo_gen = self.tokenizer.encode(formatted_chat_wo_gen, add_special_tokens=False)

        formatted_chat = self.tokenizer.encode(formatted_chat, add_special_tokens=False)

        response_token_ids = formatted_chat[len(formatted_chat_wo_gen):]

        response_template = self.tokenizer.decode(response_token_ids)

        logging.warning(f"Inferred response template: {repr(response_template)}")

        logging.warning(f"Inferred response template token ids: {response_token_ids}")

        return response_template, response_token_ids

    def apply_chat_template(self, q_text: str, a_text: str):

        chat = [

            {"role": "system", "content": self.sys_prompt},

            {"role": "user", "content": q_text},

            {"role": "assistant", "content": a_text}

        ]

        formatted_chat = self.tokenizer.apply_chat_template(

            chat,

            tokenize=False,

            add_generation_prompt=False

        )

        return formatted_chat

    @staticmethod

    def pad_tensors(tensor_list, padding_value=0):

        max_num_events = max(tensor.shape[0] for tensor in tensor_list)

        feature_dim = tensor_list[0].shape[1]

        batch_size = len(tensor_list)

        padded_tensor = torch.full((batch_size, max_num_events, feature_dim), padding_value, dtype=torch.float)

        is_padding = torch.ones((batch_size, max_num_events), dtype=torch.bool)

        for i, tensor in enumerate(tensor_list):

            num_events = tensor.shape[0]

            padded_tensor[i, :num_events, :] = tensor

            is_padding[i, :num_events] = 0

        return padded_tensor, is_padding

    def mask_instruction(self, labels: torch.Tensor) -> torch.Tensor:

        for i in range(len(labels)):

            response_token_ids_start_idx = None

            for idx in torch.where(labels[i] == self.response_token_ids[0])[0]:

                if self.response_token_ids == labels[i][idx: idx + len(self.response_token_ids)].tolist():

                    response_token_ids_start_idx = idx

            if response_token_ids_start_idx is None:

                warnings.warn(

                    f"Could not find response key `{self.response_template}` in the "

                    f'following instance: {self.tokenizer.decode(labels[i])} '

                    f"This instance will be ignored in loss calculation. "

                    f"Note, if this happens often, consider increasing the `max_seq_length`."

                )

                labels[i, :] = self.ignore_index

            else:

                response_token_ids_end_idx = response_token_ids_start_idx + len(self.response_token_ids)

                labels[i, :response_token_ids_end_idx] = self.ignore_index

        return labels

    def __call__(self, batch: List) -> Dict[str, torch.Tensor]:

        all_text = []

        all_events = []

        for data in batch:

            text = self.apply_chat_template(

                q_text=data[0],

                a_text=data[1],

            )

            all_text.append(text)

            all_events.append(data[2])

        inputs = self.tokenizer(

            all_text,

            return_tensors="pt",

            padding=True,

            truncation=True,

            add_special_tokens=False,

        )

        input_ids = inputs["input_ids"]

        pixel_values, pixel_values_is_padding = self.pad_tensors(all_events)

        attention_mask = inputs["attention_mask"]

        labels = self.mask_instruction(input_ids.clone())

        return {

            "input_ids": input_ids,

            "pixel_values": pixel_values,

            "attention_mask": attention_mask,

            "labels": labels,

            "pixel_values_is_padding": pixel_values_is_padding,

        }

if __name__ == "__main__":

    from src.dataset.dataset import InstructionTuningDataset

    from torch.utils.data import DataLoader

    from src.model.init_llemr import init_llemr

    # llm_pretrained_model_name_or_path = "Qwen/Qwen2-0.5B-Instruct"

    llm_pretrained_model_name_or_path = "lmsys/vicuna-7b-v1.5"

    device = "cuda:0"

    llemr, tokenizer = init_llemr(llm_pretrained_model_name_or_path, hidden_size=1027)

    llemr.to(device)

    dataset = InstructionTuningDataset(split="train", source="event")

    print(len(dataset))

    collator = InstructionTuningCollator(

        tokenizer=tokenizer,

    )

    loader = DataLoader(

        dataset,

        batch_size=8,

        collate_fn=collator,

    )

    batch = next(iter(loader))

    print(batch["input_ids"].shape)

    print(batch["pixel_values"].shape)

    print(batch["attention_mask"].shape)

    print(batch["labels"].shape)

    print(batch["pixel_values_is_padding"].shape)

    for key, value in batch.items():

        batch[key] = value.to(device)

    with torch.no_grad():

        outputs = llemr(**batch)

    print(outputs.loss)

    print(outputs.logits.shape)

    llemr.train()

    for parameters in llemr.language_model.parameters():

        parameters.requires_grad = False

    outputs = llemr(**batch)

    print(outputs.loss)

    print(outputs.logits.shape)

    outputs.loss.backward()

    print("Success")