from collections import defaultdict
import copy
import os
from dataclasses import dataclass, field
import random
import json
import logging
import pathlib
from typing import Dict, Optional, Sequence
import glob

from datasets import get_dataset
from fastchat.model.diffusion_llm_caption import DiffusionLlmForCausalLM

import torch
import numpy as np

import transformers
from torch.utils.data import Dataset
from transformers import Trainer, AddedToken

from fastchat.conversation import conv_one_shot as default_conversation

from transformers import BitsAndBytesConfig

# TODO: import and use code from ../data/dataset.py

IGNORE_INDEX = -100
DEFAULT_PAD_TOKEN = "[PAD]"
DEFAULT_EOS_TOKEN = "</s>"
DEFAULT_BOS_TOKEN = "</s>"
DEFAULT_UNK_TOKEN = "</s>"


@dataclass
class ModelArguments:
    model_name_or_path: Optional[str] = field(default="facebook/opt-125m")


@dataclass
class DataArguments:
    data_path: str = field(default=None,
                           metadata={"help": "Path to the training data."})
    lazy_preprocess: bool = False
    num_data: int = -1
    preprocessed_path: str = field(default=None,
                                   metadata={"help": "Path to the preprocessed training data."})


@dataclass
class TrainingArguments(transformers.TrainingArguments):
    cache_dir: Optional[str] = field(default=None)
    optim: str = field(default="adamw_torch")
    model_max_length: int = field(
        default=2048,
        metadata={
            "help":
                "Maximum sequence length. Sequences will be right padded (and possibly truncated)."
        },
    )


def safe_save_model_for_hf_trainer(trainer: transformers.Trainer,
                                   output_dir: str):
    """Collects the state dict and dump to disk."""
    state_dict = trainer.model.state_dict()
    if trainer.args.should_save:
        cpu_state_dict = {
            key: value
            for key, value in state_dict.items()
        }
        # potential bug for T5 model
        del state_dict
        trainer._save(output_dir, state_dict=cpu_state_dict)  # noqa


def get_feature_dir_info(root):
    files = glob.glob(os.path.join(root, '*_id.npy'))
    files_caption = glob.glob(os.path.join(root, '*_*.npy'))
    num_data = len(files)
    n_captions = [0] * num_data
    for f in files_caption:
        name = os.path.split(f)[-1]
        k1, k2 = os.path.splitext(name)[0].split('_')
        if k2 != 'tmp' and k2 != 'id':
            n_captions[int(k1)] += 1
    for i in range(len(n_captions)):
        if i > 0:
            n_captions[i] = n_captions[i] + n_captions[i - 1]
    return n_captions[-1], n_captions


def get_feature_test(root):
    files = glob.glob(os.path.join(root, '*_tmp.npy'))
    num_data = len(files)
    return num_data


def get_rand_des():
    text = ['Describe the image concisely.',
            'Provide a brief description of the given image.',
            'Offer a succinct explanation of the picture presented.',
            'Can you describe this image?',
            'Summarize the visual content of the image.',
            'Give a short and clear explanation of the subsequent image.',
            'Share a concise interpretation of the image provided.',
            'Present a compact description of the photo’s key features.',
            'Relay a brief, clear account of the picture shown.',
            'Render a clear and concise summary of the photo.',
            'Write a terse but informative summary of the picture.',
            'Create a compact narrative representing the image presented.']

    return text[random.randint(0, 11)]


class MSCOCOFeatureDataset:
    def __init__(self, root):
        self.root = root
        self.num_data, self.n_captions = get_feature_dir_info(root)
        tokenizer = transformers.T5Tokenizer.from_pretrained(
            pretrained_model_name_or_path='lmsys/fastchat-t5-3b-v1.0',
            model_max_length=2048,
            padding_side="right",
            use_fast=False,
        )
        head = 0
        tail = 0
        q_list = []
        a_list = []
        o_list = []
        llm_list = []
        from tqdm import tqdm
        for i in tqdm(range(8000)):
            if self.n_captions[head] > i:
                if head > 0:
                    tail = i - self.n_captions[head - 1]
                else:
                    tail = i
            else:
                head = head + 1
                tail = i - self.n_captions[head - 1]
            # c = np.load(os.path.join(self.root, f'{head}_{tail}.npy'))
            # v = np.load(os.path.join(self.root, f'{head}_{tail}_llm.npy'))
            text = np.load(os.path.join(self.root, f'{head}_{tail}_text.npy'))
            llm_text = np.load(os.path.join(self.root, f'{head}_{tail}_llm.npy'))
            original_answer = np.load(os.path.join(self.root, f'{head}_{tail}.npy'))
            llm_list.append(llm_text)
            query = get_rand_des()
            text = str(text)
            conversation = ('### Human: <Img> ' + text.strip().replace('\n', '')
                            + ' </Img> ' + query + '\n' + '### Assistant: ')
            tokenized_conversation = tokenizer(conversation, max_length=None)["input_ids"]
            q_list.append(tokenized_conversation)
            answer = tokenizer(text.strip().replace('\n', '') + '\n')["input_ids"]
            a_list.append(answer)
            o_list.append(original_answer)
        result = dict(input_ids=q_list, labels=a_list, original_labels=o_list, encode_text=llm_list)
        # json_data_dict = json.dumps(result)
        np.save('/data3/xiangyu/Data/preprocessed_data/processed_train.npy', result)
        # open file for writing, "w"
        # f = open('/data3/xiangyu/Data/preprocessed_data/processed_train.json', "w")
        # write json object to file
        # f.write(json_data_dict)


class LazySupervisedDataset2014(Dataset):
    def __init__(self, tokenizer: transformers.PreTrainedTokenizer):
        super(LazySupervisedDataset2014, self).__init__()
        self.tokenizer = tokenizer
        self.root = "/home/data2/xiangyu/Data/coco512_features/train_diff"
        self.num_data, self.n_captions = get_feature_dir_info(self.root)

    def __len__(self):
        return self.num_data

    def __getitem__(self, index):
        head = 0
        tail = 0
        for i in range(index + 1):
            if self.n_captions[i] > index:
                head = i
                if i > 0:
                    tail = index - self.n_captions[i - 1]
                else:
                    tail = index
                break
        # c = np.load(os.path.join(self.root, f'{head}_{tail}.npy'))
        # v = np.load(os.path.join(self.root, f'{head}_{tail}_llm.npy'))
        text = np.load(os.path.join(self.root, f'{head}_{tail}.npy'))
        llm_text = np.load(os.path.join(self.root, f'{head}_tmp.npy'))
        llm_text = llm_text[0]
        original_answer = np.load(os.path.join(self.root, f'{head}_{tail}.npy'))
        query = get_rand_des()
        text = str(text)
        conversation = ('### Human: <Img> ' + text.strip().replace('\n', '')
                        + ' </Img> ' + query + '\n' + '### Assistant: ')
        tokenized_conversation = self.tokenizer(conversation, max_length=None).input_ids
        answer = self.tokenizer(text.strip().replace('\n', '') + '\n').input_ids
        return dict(input_ids=answer, labels=answer,
                    original_labels=llm_text, encode_text=llm_text)


class LazySupervisedDatasetRandom(Dataset):
    def __init__(self, tokenizer: transformers.PreTrainedTokenizer):
        super(LazySupervisedDatasetRandom, self).__init__()
        self.tokenizer = tokenizer
        self.root = "/home/data2/xiangyu/Data/coco512_features/train_diff"
        self.num_data = get_feature_test(self.root)

    def __len__(self):
        return self.num_data

    def __getitem__(self, index):
        head = index
        tail = 0
        # c = np.load(os.path.join(self.root, f'{head}_{tail}.npy'))
        # v = np.load(os.path.join(self.root, f'{head}_{tail}_llm.npy'))
        text = np.load(os.path.join(self.root, f'{head}_{tail}.npy'))
        llm_text = np.load(os.path.join(self.root, f'{head}_tmp.npy'))
        llm_text = llm_text[0]
        original_answer = np.load(os.path.join(self.root, f'{head}_{tail}.npy'))
        query = get_rand_des()
        text = str(text)
        conversation = ('### Human: <Img> ' + text.strip().replace('\n', '')
                        + ' </Img> ' + query + '\n' + '### Assistant: ')
        tokenized_conversation = self.tokenizer(conversation, max_length=None).input_ids
        answer = self.tokenizer(text.strip().replace('\n', '') + '\n').input_ids
        return dict(input_ids=answer, labels=answer,
                    original_labels=llm_text, encode_text=llm_text)


class LazySupervisedDataset(Dataset):
    def __init__(self, tokenizer: transformers.PreTrainedTokenizer):
        super(LazySupervisedDataset, self).__init__()
        self.tokenizer = tokenizer
        self.root = "/home/data2/xiangyu/Data/coco512_features/karpathy"
        self.num_data = get_feature_test(self.root)

    def __len__(self):
        return self.num_data

    def __getitem__(self, index):
        # c = np.load(os.path.join(self.root, f'{head}_{tail}.npy'))
        # v = np.load(os.path.join(self.root, f'{head}_{tail}_llm.npy'))
        text = np.load(os.path.join(self.root, f'{index}_text.npy'))
        tmp = np.load(os.path.join(self.root, f'{index}_tmp.npy'))
        tmp = tmp[0]
        diff = np.load(os.path.join(self.root, f'{index}_diff.npy'))
        image_id = np.load(os.path.join(self.root, f'{index}_id.npy'))
        query = get_rand_des()
        text = str(text)
        answer = self.tokenizer(text.strip().replace('\n', '') + '\n').input_ids
        return dict(input_ids=image_id, labels=answer,
                    original_labels=text, encode_text=tmp)


class LazySupervisedDatasetQA(Dataset):
    def __init__(self, tokenizer: transformers.PreTrainedTokenizer):
        super(LazySupervisedDatasetQA, self).__init__()
        self.tokenizer = tokenizer
        self.root = "/home/data2/xiangyu/Data/coco512_features/vqa_diff"
        self.num_data = get_feature_test(self.root)

    def __len__(self):
        return self.num_data

    def __getitem__(self, index):
        # c = np.load(os.path.join(self.root, f'{head}_{tail}.npy'))
        # v = np.load(os.path.join(self.root, f'{head}_{tail}_llm.npy'))
        text = np.load(os.path.join(self.root, f'{index}_question.npy'))
        tmp = np.load(os.path.join(self.root, f'{index}_tmp.npy'))
        tmp = tmp[0]
        diff = np.load(os.path.join(self.root, f'{index}_answer.npy'))
        answer = diff[1]
        image_id = np.load(os.path.join(self.root, f'{index}_id.npy'))
        query = get_rand_des()
        text = str(text)
        text = self.tokenizer(text.strip().replace('\n', '') + '\n').input_ids
        answer = self.tokenizer(answer.strip().replace('\n', '') + '\n').input_ids
        return dict(input_ids=image_id, labels=text,
                    original_labels=answer, encode_text=tmp)


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

    def __init__(self, data_path: str,
                 tokenizer: transformers.PreTrainedTokenizer,
                 preprocessed_path,
                 num_data=5000):
        super(SupervisedDataset, self).__init__()

        # save to file
        self.preprocessed_path = preprocessed_path
        if not os.path.exists("/data3/xiangyu/Data/preprocessed_data/"):
            os.mkdir("/data3/xiangyu/Data/preprocessed_data/")
        if os.path.exists(self.preprocessed_path):
            print("loading from preprocessed data")
            # data_dict = json.load(open(self.preprocessed_path, "r"))
            data_dict = np.load(self.preprocessed_path, allow_pickle=True).item()
            print("loading data finished")
        else:
            logging.warning("Loading data...")
            list_data_dict = json.load(open(data_path, "r"))

            logging.warning("Formatting inputs...")
            sources = []

            sources = [example["conversations"] for example in list_data_dict]

            data_dict = preprocess(sources, tokenizer)
            json_data_dict = json.dumps(data_dict)

            # open file for writing, "w"
            f = open(self.preprocessed_path, "w")

            # write json object to file
            f.write(json_data_dict)

        if num_data != -1:
            data_dict["input_ids"] = data_dict["input_ids"][:num_data]
            data_dict["labels"] = data_dict["labels"][:num_data]
            data_dict["original_labels"] = data_dict["original_labels"][:num_data]
            data_dict["encode_text"] = data_dict["encode_text"][:num_data]

        # Shuffle data to see more conversations, if only train on partial data
        temp = list(zip(data_dict["input_ids"], data_dict["labels"],
                        data_dict["original_labels"], data_dict["encode_text"]))
        random.shuffle(temp)
        res1, res2, res3, res4 = zip(*temp)
        data_dict["input_ids"], data_dict["labels"], data_dict["original_labels"], data_dict["encode_text"] = \
            list(res1), list(res2), list(res3), list(res4)

        # Dacheng: Get rid of short QA pair
        self.input_ids = copy.deepcopy(data_dict["input_ids"])
        self.labels = copy.deepcopy(data_dict["labels"])
        self.original_labels = copy.deepcopy(data_dict["original_labels"])
        self.encode_text = copy.deepcopy(data_dict["encode_text"])
        length_arr = defaultdict(int)
        for idx, (input, label, original_labels, encode_text) in \
                enumerate(zip(data_dict["input_ids"], data_dict["labels"],
                              data_dict["original_labels"], data_dict["encode_text"])):
            length_arr[str(len(label) // 100)] += 1
            if len(input) <= 5:
                del_idx = self.input_ids.index(input)
                self.input_ids.pop(del_idx)
                self.labels.pop(del_idx)
                self.original_labels.pop(del_idx)
                self.encode_text.pop(del_idx)
            if len(label) <= 5:
                del_idx = self.labels.index(label)
                self.input_ids.pop(del_idx)
                self.labels.pop(del_idx)
                self.original_labels.pop(del_idx)
                self.encode_text.pop(del_idx)
        for input, label in zip(self.input_ids, self.labels):
            assert len(input) >= 5
            assert len(label) >= 5

    def __len__(self):
        return len(self.input_ids)

    def __getitem__(self, i) -> Dict[str, torch.Tensor]:
        return dict(input_ids=self.input_ids[i], labels=self.labels[i],
                    original_labels=self.original_labels[i], encode_text=self.encode_text[i])


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

    tokenizer: transformers.PreTrainedTokenizer

    def __call__(self, instances: Sequence[Dict]) -> Dict[str, torch.Tensor]:
        input_ids, labels, original_labels, encode_text = \
            tuple([torch.as_tensor(instance[key]) for instance in instances]
                  for key in ("input_ids", "labels", "original_labels", "encode_text"))
        # original_labels = torch.stack(original_labels, 0)
        encode_text = torch.stack(encode_text, 0)
        # input_ids = torch.nn.utils.rnn.pad_sequence(
        #     input_ids,
        #     batch_first=True,
        #     padding_value=self.tokenizer.pad_token_id)
        labels = torch.nn.utils.rnn.pad_sequence(labels,
                                                 batch_first=True,
                                                 padding_value=IGNORE_INDEX)
        ret = dict(
            input_ids=input_ids,
            labels=labels,
            original_labels=original_labels,
            encode_text=encode_text,
            # attention_mask=input_ids.ne(self.tokenizer.pad_token_id),
        )
        torch.set_printoptions(profile="full")
        return ret


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

    tokenizer: transformers.PreTrainedTokenizer

    def __call__(self, instances: Sequence[Dict]) -> Dict[str, torch.Tensor]:
        input_ids, labels, original_labels, encode_text = \
            tuple([torch.as_tensor(instance[key]) for instance in instances]
                  for key in ("input_ids", "labels", "original_labels", "encode_text"))
        # original_labels = torch.stack(original_labels, 0)
        encode_text = torch.stack(encode_text, 0)
        # labels = torch.nn.utils.rnn.pad_sequence(labels,
        #                                          batch_first=True,
        #                                          padding_value=IGNORE_INDEX)
        original_labels = torch.nn.utils.rnn.pad_sequence(original_labels,
                                                          batch_first=True,
                                                          padding_value=IGNORE_INDEX)
        ret = dict(
            input_ids=input_ids,
            labels=labels,
            original_labels=original_labels,
            encode_text=encode_text,
        )
        torch.set_printoptions(profile="full")
        return ret


def smart_tokenizer_and_embedding_resize(
        special_tokens_dict: Dict,
        other_tokens,
        tokenizer: transformers.PreTrainedTokenizer,
        model: transformers.PreTrainedModel,
):
    """Resize tokenizer and embedding.

    Note: This is the unoptimized version that may make your embedding size not be divisible by 64.
    """
    num_new_tokens = tokenizer.add_special_tokens(special_tokens_dict)
    for new_token in other_tokens:
        num_new_tokens += tokenizer.add_tokens(AddedToken(new_token, normalized=False))

    model.resize_token_embeddings(len(tokenizer))

    if num_new_tokens > 0:
        input_embeddings = model.get_input_embeddings().weight.data
        output_embeddings = model.get_output_embeddings().weight.data

        input_embeddings_avg = input_embeddings[:-num_new_tokens].mean(
            dim=0, keepdim=True)
        output_embeddings_avg = output_embeddings[:-num_new_tokens].mean(
            dim=0, keepdim=True)

        input_embeddings[-num_new_tokens:] = input_embeddings_avg
        output_embeddings[-num_new_tokens:] = output_embeddings_avg


def make_supervised_data_module(tokenizer: transformers.PreTrainedTokenizer,
                                data_args) -> Dict:
    """Make dataset and collator for supervised fine-tuning."""
    # dataset_cls = SupervisedDataset
    # train_dataset = dataset_cls(tokenizer=tokenizer,
    #                             data_path=data_args.data_path,
    #                             preprocessed_path=data_args.preprocessed_path,
    #                             num_data=data_args.num_data)
    dataset_cls = LazySupervisedDatasetRandom
    train_dataset = dataset_cls(tokenizer=tokenizer)

    data_collator = DataCollatorForSupervisedDataset2014(tokenizer=tokenizer)
    return dict(train_dataset=train_dataset,
                eval_dataset=None,
                data_collator=data_collator)


def make_supervised_data_module_qa(tokenizer: transformers.PreTrainedTokenizer,
                                   data_args) -> Dict:
    """Make dataset and collator for supervised fine-tuning."""
    # dataset_cls = SupervisedDataset
    # train_dataset = dataset_cls(tokenizer=tokenizer,
    #                             data_path=data_args.data_path,
    #                             preprocessed_path=data_args.preprocessed_path,
    #                             num_data=data_args.num_data)
    dataset_cls = LazySupervisedDatasetQA
    train_dataset = dataset_cls(tokenizer=tokenizer)

    data_collator = DataCollatorForSupervisedDataset(tokenizer=tokenizer)
    return dict(train_dataset=train_dataset,
                eval_dataset=None,
                data_collator=data_collator)


# from fastchat.model.diffusion_llm import DiffusionLlmForCausalLM


def train():
    parser = transformers.HfArgumentParser(
        (ModelArguments, DataArguments, TrainingArguments))
    model_args, data_args, training_args = parser.parse_args_into_dataclasses()
    # quantization_config = BitsAndBytesConfig(load_in_8bit_fp32_cpu_offload=True)

    print("loading model")
    model = DiffusionLlmForCausalLM()
    #model.load_state_dict(torch.load('/home/data2/xiangyu/Data/flant5_full_train/pytorch_model.bin'))
    #model.load_state_dict(torch.load('/home/data2/xiangyu/Data/checkpoints_flant5_mm/pytorch_model.bin'))
    print("finishing loading model")
    # model.get_model().get_encoder().eval()
    params_grad = [n for n, p in model.named_parameters() if p.requires_grad]
    print(params_grad)
    # if len(params_no_grad) > 0:
    #     if training_args.fsdp is not None and len(training_args.fsdp) > 0:
    #         if len(params_no_grad) < 10:
    #             print('[WARNING] Attempting to use FSDP while {} parameters do not require gradients: {}'.format(
    #                 len(params_no_grad), params_no_grad))
    #         else:
    #             print(
    #                 '[WARNING] Attempting to use FSDP while {} parameters do not require gradients: {}...(omitted)'.format(
    #                     len(params_no_grad), ', '.join(params_no_grad[:10])))
    #         print("[WARNING] Attempting to use FSDP with partially frozen paramters, this is experimental.")
    #         print(
    #             "[WARNING] As of 4/30/23, this feature requires PyTorch-nightly build.  See here for details: https://github.com/haotian-liu/LLaVA#experimental-use-fsdp-to-save-memory-in-pretraining")
    #
    #         from torch.distributed.fsdp.fully_sharded_data_parallel import FullyShardedDataParallel as FSDP
    #         def patch_FSDP_use_orig_params(func):
    #             def wrap_func(*args, **kwargs):
    #                 use_orig_params = kwargs.pop('use_orig_params', True)
    #                 return func(*args, **kwargs, use_orig_params=use_orig_params)
    #
    #             return wrap_func
    #         FSDP.__init__ = patch_FSDP_use_orig_params(FSDP.__init__)

    # Dacheng: Note we can only use T5Tokenizer, otherwise it will prepend
    # a space before special tokens.
    tokenizer = transformers.T5Tokenizer.from_pretrained(
        model_args.model_name_or_path,
        cache_dir=training_args.cache_dir,
        model_max_length=training_args.model_max_length,
        padding_side="right",
        use_fast=False,
    )

    smart_tokenizer_and_embedding_resize(
        special_tokens_dict=dict(pad_token=DEFAULT_PAD_TOKEN),
        other_tokens=["<", "{", "\n", "}", "`", " ", "\\", "^", "\t"],
        tokenizer=tokenizer,
        model=model.get_model(),
    )

    data_module = make_supervised_data_module(tokenizer=tokenizer,
                                                 data_args=data_args)

    trainer = Trainer(model=model,
                      tokenizer=tokenizer,
                      args=training_args,
                      **data_module)

    # if list(pathlib.Path(training_args.output_dir).glob("checkpoint-*")):
    #     trainer.train(resume_from_checkpoint=True)
    # else:
    #     trainer.train()
    trainer.train()
    trainer.save_state()
    safe_save_model_for_hf_trainer(trainer=trainer,
                                   output_dir=training_args.output_dir)


if __name__ == "__main__":
    train()
