from typing import Dict, Optional, Sequence, List
from dataclasses import dataclass, field
import json
import torch
import transformers
import tokenizers

from PIL import Image
import os
import copy
import numpy as np
import cv2

from transformers import CLIPProcessor, CLIPModel, CLIPImageProcessor
import dataset.conversation as conversation_lib
from dataset.constants import IGNORE_INDEX, EVENT_TOKEN_INDEX, DEFAULT_EVENT_TOKEN, DEFAULT_EV_START_TOKEN, DEFAULT_EV_END_TOKEN
from torch.utils.data import Dataset
from packaging import version

local_rank = None
IS_TOKENIZER_GREATER_THAN_0_14 = version.parse(tokenizers.__version__) >= version.parse('0.14')

@dataclass
class ModelArguments:
    model_name_or_path: Optional[str] = field(default="facebook/opt-125m")
    version: Optional[str] = field(default="v0")
    freeze_backbone: bool = field(default=False)
    tune_mm_mlp_adapter: bool = field(default=False)
    vision_tower: Optional[str] = field(default=None)
    mm_vision_select_layer: Optional[int] = field(default=-1)   # default to the last layer
    pretrain_mm_mlp_adapter: Optional[str] = field(default=None)
    mm_projector_type: Optional[str] = field(default='linear')
    mm_use_im_start_end: bool = field(default=False)
    mm_use_im_patch_token: bool = field(default=True)
    mm_patch_merge_type: Optional[str] = field(default='flat')
    mm_vision_select_feature: Optional[str] = field(default="patch")

@dataclass
class DataArguments:
    data_path: str = field(default=None,
                           metadata={"help": "Path to the training data."})
    lazy_preprocess: bool = False
    is_multimodal: bool = False
    image_folder: Optional[str] = field(default=None)
    image_aspect_ratio: str = 'square'
    event_folder: Optional[str] = field(default=None)

def split_event_by_time(event_npy, time_interval=50000):
    # Extract data from the event_npy dictionary
    p = event_npy['p']
    t = event_npy['t']
    x = event_npy['x']
    y = event_npy['y']

    # Calculate the time bin for each timestamp based on the given time interval
    time_bins = (t // time_interval) * time_interval

    # Get the unique time bins (intervals)
    unique_bins = np.unique(time_bins)

    # Split the data according to the time bins
    split_data = [
        {
            'p': p[time_bins == bin],
            't': t[time_bins == bin],
            'x': x[time_bins == bin],
            'y': y[time_bins == bin]
        }
        for bin in unique_bins
    ]

    return split_data

def preprocess_multimodal(
    sources: Sequence[str],
    data_args: DataArguments
) -> Dict:
    is_multimodal = data_args.is_multimodal
    if not is_multimodal:
        return sources

    for source in sources:
        for sentence in source:
            if DEFAULT_EVENT_TOKEN in sentence['value']:
                sentence['value'] = sentence['value'].replace(DEFAULT_EVENT_TOKEN, '').strip()
                sentence['value'] = DEFAULT_EVENT_TOKEN + '\n' + sentence['value']
                sentence['value'] = sentence['value'].strip()
                if "mmtag" in conversation_lib.default_conversation.version:
                    sentence['value'] = sentence['value'].replace(DEFAULT_EVENT_TOKEN, '<Event>' + DEFAULT_EVENT_TOKEN + '</Event>')
            replace_token = DEFAULT_EVENT_TOKEN
            if data_args.mm_use_im_start_end: 
                replace_token = DEFAULT_EV_START_TOKEN + replace_token + DEFAULT_EV_END_TOKEN
            sentence["value"] = sentence["value"].replace(DEFAULT_EVENT_TOKEN, replace_token)

    return sources


@dataclass
class TrainingArguments(transformers.TrainingArguments):
    cache_dir: Optional[str] = field(default=None)
    optim: str = field(default="adamw_torch")
    remove_unused_columns: bool = field(default=False)
    freeze_mm_mlp_adapter: bool = field(default=False)
    mpt_attn_impl: Optional[str] = field(default="triton")
    model_max_length: int = field(
        default=512,
        metadata={
            "help":
            "Maximum sequence length. Sequences will be right padded (and possibly truncated)."
        },
    )
    double_quant: bool = field(
        default=True,
        metadata={"help": "Compress the quantization statistics through double quantization."}
    )
    quant_type: str = field(
        default="nf4",
        metadata={"help": "Quantization data type to use. Should be one of `fp4` or `nf4`."}
    )
    bits: int = field(
        default=16,
        metadata={"help": "How many bits to use."}
    )
    lora_enable: bool = False
    lora_r: int = 64
    lora_alpha: int = 16
    lora_dropout: float = 0.05
    lora_weight_path: str = ""
    lora_bias: str = "none"
    mm_projector_lr: Optional[float] = None
    group_by_modality_length: bool = field(default=False)

def rank0_print(*args):
    if local_rank == 0:
        print(*args)


def preprocess_plain(
    sources: Sequence[str],
    tokenizer: transformers.PreTrainedTokenizer,
) -> Dict:
    # add end signal and concatenate together
    conversations = []
    for source in sources:
        assert len(source) == 2
        assert DEFAULT_EVENT_TOKEN in source[0]['value']
        source[0]['value'] = DEFAULT_EVENT_TOKEN
        conversation = source[0]['value'] + source[1]['value'] + conversation_lib.default_conversation.sep
        conversations.append(conversation)
    # tokenize conversations
    input_ids = [tokenizer_event_token(prompt, tokenizer, return_tensors='pt') for prompt in conversations]
    targets = copy.deepcopy(input_ids)
    for target, source in zip(targets, sources):
        tokenized_len = len(tokenizer_event_token(source[0]['value'], tokenizer))
        target[:tokenized_len] = IGNORE_INDEX

    return dict(input_ids=input_ids, labels=targets)

def _add_speaker_and_signal(header, source, get_conversation=True):
    """Add speaker and start/end signal on each round."""
    BEGIN_SIGNAL = "### "
    END_SIGNAL = "\n"
    conversation = header
    for sentence in source:
        from_str = sentence["from"]
        if from_str.lower() == "human":
            from_str = conversation_lib.default_conversation.roles[0]
        elif from_str.lower() == "gpt":
            from_str = conversation_lib.default_conversation.roles[1]
        else:
            from_str = 'unknown'
        sentence["value"] = (BEGIN_SIGNAL + from_str + ": " +
                             sentence["value"] + END_SIGNAL)
        if get_conversation:
            conversation += sentence["value"]
    conversation += BEGIN_SIGNAL
    return conversation

def preprocess_qwen(sources, tokenizer: transformers.PreTrainedTokenizer, has_image: bool = False, max_len=2048, system_message: str = "You are a helpful assistant.") -> Dict:
    # roles = {"human": "<|im_start|>user", "gpt": "<|im_start|>assistant"}
    roles = {"human": "user", "gpt": "assistant"}

    # Add image tokens to tokenizer as a special tokens
    # Use a deepcopy of tokenizer so that we don't modify on the tokenizer
    tokenizer = copy.deepcopy(tokenizer)
    # When there is actually an image, we add the image tokens as a special token
    if has_image:
        tokenizer.add_tokens(["<event>"], special_tokens=True)

    image_token_index = tokenizer.convert_tokens_to_ids("<event>")
    im_start, im_end = tokenizer.additional_special_tokens_ids[0], tokenizer.additional_special_tokens_ids[1]
    # unmask_tokens = ["<|im_start|>", "<|im_start|>", "\n"]
    unmask_tokens_idx =  [198, im_start, im_end]
    nl_tokens = tokenizer("\n").input_ids

    # Reset Qwen chat templates so that it won't include system message every time we apply
    chat_template = "{% for message in messages %}{{'<|im_start|>' + message['role'] + '\n' + message['content'] + '<|im_end|>' + '\n'}}{% endfor %}{% if add_generation_prompt %}{{ '<|im_start|>assistant\n' }}{% endif %}"
    tokenizer.chat_template = chat_template

    # _system = tokenizer("system").input_ids + nl_tokens
    # _user = tokenizer("user").input_ids + nl_tokens
    # _assistant = tokenizer("assistant").input_ids + nl_tokens

    # Apply prompt templates
    input_ids, targets = [], []
    for i, source in enumerate(sources):
        if roles[source[0]["from"]] != roles["human"]:
            source = source[1:]

        input_id, target = [], []

        # New version, use apply chat template
        # Build system message for each sentence
        input_id += tokenizer.apply_chat_template([{"role" : "system", "content" : system_message}])
        target += [IGNORE_INDEX] * len(input_id)

        for conv in source:
            # Make sure llava data can load
            try:
                role = conv["role"]
                content = conv["content"]
            except:
                role = conv["from"]
                content = conv["value"]

            role =  roles.get(role, role)
            
            conv = [{"role" : role, "content" : content}]
            encode_id = tokenizer.apply_chat_template(conv)
            input_id += encode_id
            if role in ["user", "system"]:
                target += [IGNORE_INDEX] * len(encode_id)
            else:
                target += encode_id
                         
        assert len(input_id) == len(target), f"{len(input_id)} != {len(target)}"
        for idx, encode_id in enumerate(input_id):
            if encode_id in unmask_tokens_idx:
                target[idx] = encode_id
            if encode_id == image_token_index:
                input_id[idx] = EVENT_TOKEN_INDEX
        input_ids.append(input_id)
        targets.append(target)
    input_ids = torch.tensor(input_ids, dtype=torch.long)
    targets = torch.tensor(targets, dtype=torch.long)

    return dict(
        input_ids=input_ids,  # tensor(bs x seq_len)
        labels=targets,  # tensor(bs x seq_len)
    )

def preprocess_v1(
    sources,
    tokenizer: transformers.PreTrainedTokenizer,
    has_image: bool = False
) -> Dict:
    conv = conversation_lib.default_conversation.copy()
    roles = {"human": conv.roles[0], "gpt": conv.roles[1]}

    # Apply prompt templates
    conversations = []
    for i, source in enumerate(sources):
        if roles[source[0]["from"]] != conv.roles[0]:
            # Skip the first one if it is not from human
            source = source[1:]

        conv.messages = []
        for j, sentence in enumerate(source):
            role = roles[sentence["from"]]
            assert role == conv.roles[j % 2], f"{i}"
            conv.append_message(role, sentence["value"])
        conversations.append(conv.get_prompt())

    # Tokenize conversations

    if has_image:
        input_ids = torch.stack([tokenizer_event_token(prompt, tokenizer, return_tensors='pt') for prompt in conversations], dim=0)
    else:
        input_ids = tokenizer(
            conversations,
            return_tensors="pt",
            padding="longest",
            max_length=tokenizer.model_max_length,
            truncation=True,
        ).input_ids

    targets = input_ids.clone()

    assert conv.sep_style == conversation_lib.SeparatorStyle.TWO

    # Mask targets
    sep = conv.sep + conv.roles[1] + ": "
    for conversation, target in zip(conversations, targets):
        total_len = int(target.ne(tokenizer.pad_token_id).sum())

        rounds = conversation.split(conv.sep2)
        cur_len = 1
        target[:cur_len] = IGNORE_INDEX
        for i, rou in enumerate(rounds):
            if rou == "":
                break

            parts = rou.split(sep)
            if len(parts) != 2:
                break
            parts[0] += sep

            if has_image:
                round_len = len(tokenizer_event_token(rou, tokenizer))
                instruction_len = len(tokenizer_event_token(parts[0], tokenizer)) - 2
            else:
                round_len = len(tokenizer(rou).input_ids)
                instruction_len = len(tokenizer(parts[0]).input_ids) - 2

            if i != 0 and not tokenizer.legacy and IS_TOKENIZER_GREATER_THAN_0_14:
                round_len -= 1
                instruction_len -= 1

            target[cur_len : cur_len + instruction_len] = IGNORE_INDEX

            cur_len += round_len
        target[cur_len:] = IGNORE_INDEX

        if cur_len < tokenizer.model_max_length:
            if cur_len != total_len:
                target[:] = IGNORE_INDEX
                print(
                    f"WARNING: tokenization mismatch: {cur_len} vs. {total_len}."
                    f" (ignored)"
                )

    return dict(
        input_ids=input_ids,
        labels=targets,
    )


def tokenizer_event_token(prompt, tokenizer, event_token_index=EVENT_TOKEN_INDEX, return_tensors=None):
    prompt_chunks = [tokenizer(chunk).input_ids for chunk in prompt.split('<event>')]

    def insert_separator(X, sep):
        return [ele for sublist in zip(X, [sep]*len(X)) for ele in sublist][:-1]

    input_ids = []
    offset = 0
    if len(prompt_chunks) > 0 and len(prompt_chunks[0]) > 0 and prompt_chunks[0][0] == tokenizer.bos_token_id:
        offset = 1
        input_ids.append(prompt_chunks[0][0])

    for x in insert_separator(prompt_chunks, [event_token_index] * (offset + 1)):
        input_ids.extend(x[offset:])

    if return_tensors is not None:
        if return_tensors == 'pt':
            return torch.tensor(input_ids, dtype=torch.long)
        raise ValueError(f'Unsupported tensor type: {return_tensors}')
    return input_ids

def _tokenize_fn(strings: Sequence[str],
                 tokenizer: transformers.PreTrainedTokenizer) -> Dict:
    """Tokenize a list of strings."""
    tokenized_list = [
        tokenizer(
            text,
            return_tensors="pt",
            padding="longest",
            max_length=tokenizer.model_max_length,
            truncation=True,
        ) for text in strings
    ]
    input_ids = labels = [
        tokenized.input_ids[0] for tokenized in tokenized_list
    ]
    input_ids_lens = labels_lens = [
        tokenized.input_ids.ne(tokenizer.pad_token_id).sum().item()
        for tokenized in tokenized_list
    ]
    return dict(
        input_ids=input_ids,
        labels=labels,
        input_ids_lens=input_ids_lens,
        labels_lens=labels_lens,
    )

def _mask_targets(target, tokenized_lens, speakers):
    # cur_idx = 0
    cur_idx = tokenized_lens[0]
    tokenized_lens = tokenized_lens[1:]
    target[:cur_idx] = IGNORE_INDEX
    for tokenized_len, speaker in zip(tokenized_lens, speakers):
        if speaker == "human":
            target[cur_idx+2:cur_idx + tokenized_len] = IGNORE_INDEX
        cur_idx += tokenized_len

def preprocess(
    sources: Sequence[str],
    tokenizer: transformers.PreTrainedTokenizer,
    has_image: bool = False
) -> Dict:
    """
    Given a list of sources, each is a conversation list. This transform:
    1. Add signal '### ' at the beginning each sentence, with end signal '\n';
    2. Concatenate conversations together;
    3. Tokenize the concatenated conversation;
    4. Make a deepcopy as the target. Mask human words with IGNORE_INDEX.
    """
    if conversation_lib.default_conversation.sep_style == conversation_lib.SeparatorStyle.PLAIN:
        return preprocess_plain(sources, tokenizer)

    if conversation_lib.default_conversation.version == "qwen":
        return preprocess_qwen(sources, tokenizer, has_image=has_image)

    if conversation_lib.default_conversation.version.startswith("v1"):
        return preprocess_v1(sources, tokenizer, has_image=has_image)

    conversations = []
    for source in sources:
        header = f"{conversation_lib.default_conversation.system}\n\n"
        conversation = _add_speaker_and_signal(header, source)
        conversations.append(conversation)
    # tokenize conversations
    def get_tokenize_len(prompts):
        return [len(tokenizer_event_token(prompt, tokenizer)) for prompt in prompts]

    if has_image:
        input_ids = [tokenizer_event_token(prompt, tokenizer, return_tensors='pt') for prompt in conversations]
    else:
        conversations_tokenized = _tokenize_fn(conversations, tokenizer)
        input_ids = conversations_tokenized["input_ids"]

    targets = copy.deepcopy(input_ids)
    for target, source in zip(targets, sources):
        if has_image:
            tokenized_lens = get_tokenize_len([header] + [s["value"] for s in source])
        else:
            tokenized_lens = _tokenize_fn([header] + [s["value"] for s in source], tokenizer)["input_ids_lens"]
        speakers = [sentence["from"] for sentence in source]
        _mask_targets(target, tokenized_lens, speakers)

    return dict(input_ids=input_ids, labels=targets)

def generate_event_image(x, y, p):  
    height, width = y.max() + 1, x.max() + 1 
    event_image = np.ones((height, width, 3), dtype=np.uint8) * 255

    blue_mask = p == 0
    red_mask = p == 1

    event_image[y[blue_mask], x[blue_mask]] = np.array([0, 0, 255])  # blue
    event_image[y[red_mask], x[red_mask]] = np.array([255, 0, 0])  # red

    return event_image

def preprocess_batch(event_images, processor):
    pixel_values = processor.preprocess(event_images, return_tensors='pt')['pixel_values']
    return pixel_values

class EventChatDataset(Dataset):
    """Dataset for supervised fine-tuning."""

    def __init__(self, data_path: str,
                 tokenizer: transformers.PreTrainedTokenizer,
                 data_args: DataArguments):
        super(EventChatDataset, self).__init__()

        list_data_dict = json.load(open(data_path, "r"))

        rank0_print("Formatting inputs...Skip in lazy mode")
        self.tokenizer = tokenizer
        self.list_data_dict = list_data_dict
        self.data_args = data_args

    def __len__(self):
        return len(self.list_data_dict)
    
    def get_event_images_list(self, event_npy, n):
        x, y, p, t = event_npy['x'], event_npy['y'], event_npy['p'], event_npy['t']

        total_events = len(t)
        events_per_image = total_events // n  

        event_image_list = [] 

        for i in range(n):
            start_idx = i * events_per_image
            end_idx = (i + 1) * events_per_image if i < n - 1 else total_events  

            x_part = x[start_idx:end_idx]
            y_part = y[start_idx:end_idx]
            p_part = p[start_idx:end_idx]

            event_img = generate_event_image(x_part, y_part, p_part)
            
            event_image_list.append(event_img) 

        return event_image_list  


    def split_event_by_time(self, event_npy, time_interval=50000):
        """
        Split event data into time intervals (default 50ms).
        
        :param event_npy: Dictionary containing event data with keys 'p', 't', 'x', 'y'.
        :param time_interval: Time interval for splitting in microseconds, default is 50ms (50,000 microseconds).
        :return: A list of dictionaries where each dictionary corresponds to a time interval with the split data.
        """
        # Extract data from the event_npy dictionary
        p = event_npy['p']
        t = event_npy['t']
        x = event_npy['x']
        y = event_npy['y']

        # Calculate the time bin for each timestamp based on the given time interval
        time_bins = (t // time_interval) * time_interval

        # Get the unique time bins (intervals)
        unique_bins = np.unique(time_bins)

        # Split the data according to the time bins
        split_data = [
            {
                'p': p[time_bins == bin],
                't': t[time_bins == bin],
                'x': x[time_bins == bin],
                'y': y[time_bins == bin]
            }
            for bin in unique_bins
        ]

        return split_data


    # @property
    # def lengths(self):
    #     length_list = []
    #     for sample in self.list_data_dict:
    #         img_tokens = 128 if 'image' in sample else 0
    #         length_list.append(sum(len(conv['value'].split()) for conv in sample['conversations']) + img_tokens)
    #     return length_list

    # @property
    # def modality_lengths(self):
    #     length_list = []
    #     for sample in self.list_data_dict:
    #         cur_len = sum(len(conv['value'].split()) for conv in sample['conversations'])
    #         cur_len = cur_len if 'image' in sample else -cur_len
    #         length_list.append(cur_len)
    #     return length_list

    def __getitem__(self, i) -> Dict[str, torch.Tensor]:
        sources = self.list_data_dict[i]
        if isinstance(i, int):
            sources = [sources]
        assert len(sources) == 1, "Don't know why it is wrapped to a list"  # FIXME
        if 'event_data' in sources[0]:
            image_file = self.list_data_dict[i]['event_data']
            image_folder = self.data_args.image_folder
            if self.data_args.event_folder:
                event_file = os.path.splitext(image_file)[0] + ".npy"
                event_folder = self.data_args.event_folder
            processor = self.data_args.event_processor
            if not hasattr(self.data_args, 'event_folder') or not self.data_args.event_folder:
                try:
                    image = Image.open(os.path.join(image_folder, image_file))
                    if image.mode != 'RGB':
                        image = image.convert('RGB')
                    image_np = np.array(image)
                    if image_np.ndim == 2 or image_np.shape[2] != 3:
                        image_np = np.stack([image_np]*3, axis=-1)
                        image = Image.fromarray(image_np.astype('uint8'))
                except OSError:
                    print(f"Error loading image: {image_file}. Returning default image...")
                    image = Image.new('RGB', (640, 480), (255, 255, 255))
            if self.data_args.image_aspect_ratio == 'pad':
                def expand2square(pil_img, background_color):
                    width, height = pil_img.size
                    if width == height:
                        return pil_img
                    elif width > height:
                        result = Image.new(pil_img.mode, (width, width), background_color)
                        result.paste(pil_img, (0, (width - height) // 2))
                        return result
                    else:
                        result = Image.new(pil_img.mode, (height, height), background_color)
                        result.paste(pil_img, ((height - width) // 2, 0))
                        return result
                image = expand2square(image, tuple(int(x*255) for x in processor.image_mean))
                image = processor.preprocess(image, return_tensors='pt')['pixel_values'][0]
            else:
                if self.data_args.event_folder:
                    if self.data_args.spatial_temporal_encoder:
                        event_npy = np.load(os.path.join(event_folder, event_file), allow_pickle=True)                   
                        # if event_npy['t'].max() - event_npy['t'].min() >= 100000:
                        if event_npy.ndim == 4:
                            event_tensor = torch.tensor(event_npy)
                            event_list = list(event_tensor)
                            event_list = event_list[:5]
                        else:
                            event_npy = np.array(event_npy).item()      
                            event_data = event_npy['event_data']
                            # event_data = np.array(event_data).item()  
                            event_feature = event_npy['event_features']

                            # event_img_list = self.get_event_images_list(event_npy, self.data_args.n_event_images)
                            # event_list = []
                            # for event in event_img_list:
                            #     event = processor.preprocess(event, return_tensors='pt')['pixel_values'][0]
                            #     event_list.append(event)
                    elif self.data_args.QFormer:
                        event_npy = np.load(os.path.join(event_folder, event_file), allow_pickle=True)
                        event_npy = np.array(event_npy).item()
                        event_npy_list = self.split_event_by_time(event_npy, 50000)
                        sampled_event_list = event_npy_list[::10]

                        event_images_batch = [generate_event_image(event['x'], event['y'], event['p']) for event in sampled_event_list]
                        event_img_tensor_batch = preprocess_batch(event_images_batch, processor)
                        event_list = [event_img_tensor_batch[i] for i in range(len(event_img_tensor_batch))]                     
                    else:
                        event_npy = np.load(os.path.join(event_folder, event_file), allow_pickle=True)
                        event_npy = np.array(event_npy).item()
                        event_img = generate_event_image(event_npy['x'], event_npy['y'], event_npy['p'])
                        event_img = processor.preprocess(event_img, return_tensors='pt')['pixel_values'][0]
                else:
                    try:
                        image = processor.preprocess(image, return_tensors='pt')['pixel_values'][0]
                    except (OSError, ValueError) as e:
                        print(f"Error processing image: {image_file}. Error: {e}. Returning default image...")
                        raise e             

            sources = preprocess_multimodal(
                copy.deepcopy([e["conversations"] for e in sources]),
                self.data_args)
        else:
            sources = copy.deepcopy([e["conversations"] for e in sources])
        data_dict = preprocess(
            sources,
            self.tokenizer,
            has_image=('event_data' in self.list_data_dict[i]))
        if isinstance(i, int):
            data_dict = dict(input_ids=data_dict["input_ids"][0],
                             labels=data_dict["labels"][0])

        # image exist in the data
        # if 'event_data' in self.list_data_dict[i] and not getattr(self.data_args, 'event_folder', None) and not self.data_args.spatial_temporal_encoder:
        #     data_dict['event_data'] = image
        #     return data_dict
    
        # if self.data_args.spatial_temporal_encoder:
        #     data_dict['events_list'] = event_list
        # elif self.data_args.QFormer:
        #     data_dict['events_list'] = event_list
        # else:
        #     if self.data_args.event_folder:
        #         data_dict['events'] = event_img
        #     else:
        #         raise("input error, please check input")
        
        data_dict['event_data'] = event_data
        data_dict['event_feature'] = event_feature

        return data_dict
    
@dataclass
class DataCollatorForEventChatDataset(object):
    """Collate examples for supervised fine-tuning."""

    tokenizer: transformers.PreTrainedTokenizer

    def pad_sequence(self, input_ids, batch_first, padding_value):
        if self.tokenizer.padding_side == "left":
            input_ids = [torch.flip(_input_ids, [0]) for _input_ids in input_ids]
        input_ids = torch.nn.utils.rnn.pad_sequence(input_ids, batch_first=batch_first, padding_value=padding_value)
        if self.tokenizer.padding_side == "left":
            input_ids = torch.flip(input_ids, [1])
        return input_ids

    def __call__(self, instances: Sequence[Dict]) -> Dict[str, torch.Tensor]:
        input_ids, labels = tuple([instance[key] for instance in instances] for key in ("input_ids", "labels"))
        input_ids = [_input_ids[: self.tokenizer.model_max_length] for _input_ids in input_ids]
        labels = [_labels[: self.tokenizer.model_max_length] for _labels in labels]
        if self.tokenizer.pad_token_id is None:
            self.tokenizer.pad_token_id = 0 
        input_ids = self.pad_sequence(input_ids, batch_first=True, padding_value=self.tokenizer.pad_token_id)
        labels = self.pad_sequence(labels, batch_first=True, padding_value=IGNORE_INDEX)
        batch = dict(input_ids=input_ids, labels=labels.long() if labels.dtype == torch.int32 else labels, attention_mask=input_ids.ne(self.tokenizer.pad_token_id))

        # if 'event_data' in instances[0]:
        #     event_tensors = [instance['event_data'] for instance in instances]
        #     if all(x is not None and x.shape == event_tensors[0].shape for x in event_tensors):
        #         batch['event_tensors'] = torch.stack(event_tensors)
        #     else:
        #         batch['event_tensors'] = event_tensors
        # elif 'events' in instances[0]:
        #     events = [instance['events'] for instance in instances]
        #     batch['event_tensors'] = torch.stack(events)

        # elif 'events_list' in instances[0]:
        #     events_list = [instance['events_list'] for instance in instances]
        #     batch['event_tensors'] = events_list
        
        # else:
        #     raise("input error, please check input")
        
        event_data = [instance['event_data'] for instance in instances]
        event_feature = [instance['event_feature'] for instance in instances]
        
        batch['event_data'] = event_data       
        event_feature = [torch.tensor(ev_f) for ev_f in event_feature]
        event_feature = self.pad_sequence(event_feature, batch_first=True, padding_value=self.tokenizer.pad_token_id)
        batch['event_feature'] = event_feature

        return batch
    

def make_supervised_data_module(tokenizer: transformers.PreTrainedTokenizer,
                                data_args) -> Dict:
    train_dataset = EventChatDataset(tokenizer=tokenizer,
                                data_path=data_args.data_path,
                                data_args=data_args)
    data_collator = DataCollatorForEventChatDataset(tokenizer=tokenizer)
    return dict(train_dataset=train_dataset,
                eval_dataset=None,
                data_collator=data_collator)


    