#!/usr/bin/env python
# -*- coding:utf-8 -*-
# @Time    : 2023/12/5 20:39
# @Author  : yebulk
import os
from torch.utils.data import Dataset
import os
import json
import numpy as np
import torch
from utils_prompt import *

img_shape = {
    "resnet": (512, 2048),
    "clip": (49, 2048),
    "detr": (100, 256),
    "vit": (145, 1024),
}

def split_chinese_string(string):
    split_result = string.split("：", 1)
    if len(split_result) > 1:
        return split_result[1]
    else:
        return string

def load_data_std(args):
    problems = json.load(open(os.path.join(args.data_root, '../Data/output.json')))
    pid_splits = json.load(open(os.path.join(args.data_root, '../Data/split.json')))
    # captions = json.load(open(args.caption_file))["captions"]

    # for qid in problems:
    #     problems[qid]['caption'] = captions[qid] if qid in captions else ""

    train_qids = pid_splits['%s' % (args.train_split)]
    val_qids = pid_splits['%s' % (args.val_split)]
    test_qids = pid_splits['%s' % (args.test_split)]
    print(f"number of train problems: {len(train_qids)}\n")
    print(f"number of val problems: {len(val_qids)}\n")
    print(f"number of test problems: {len(test_qids)}\n")

    qids = {'train': train_qids, 'val': val_qids, 'test': test_qids}
    return problems, qids,


def load_data_img(args):
    problems = json.load(open(os.path.join(args.data_root, 'output.json')))
    pid_splits = json.load(open(os.path.join(args.data_root, 'split.json')))
    # captions = json.load(open(args.caption_file))["captions"]

    name_maps = json.load(open('/mnt/sdb1/jsy/yjw/prompt/myCode/vision_features/name_map.json'))

    # check
    if args.img_type == "resnet":
        image_features = np.load('../vision_features_pics_Meme/resnet.npy')
        image_features = np.expand_dims(image_features, axis=1)
        image_features = image_features.repeat(512, axis=1)
    elif args.img_type == "clip":
        image_features = np.load('../vision_features_pics_Meme/clip.npy')
    elif args.img_type == "detr":
        image_features = torch.load('/mnt/sdb1/jsy/yjw/prompt/myCode/vision_features/detr.pth').to('cpu')
    elif args.img_type == "vit":
        image_features = torch.load("/mnt/sdb1/jsy/yjw/prompt/myCode/vision_features/vit.pth").to('cpu')
    else:
        image_features = np.load('/mnt/sdb1/jsy/yjw/prompt/myCode/vision_features/vit.pth')
    print("img_features size: ", image_features.shape)

    # for qid in problems:
    #     problems[qid]['caption'] = captions[qid] if qid in captions else ""

    train_qids = pid_splits['%s' % (args.train_split)]
    val_qids = pid_splits['%s' % (args.val_split)]
    test_qids = pid_splits['%s' % (args.test_split)]
    print(f"number of train problems: {len(train_qids)}\n")
    print(f"number of val problems: {len(val_qids)}\n")
    print(f"number of test problems: {len(test_qids)}\n")

    qids = {'train': train_qids, 'val':val_qids,'test':test_qids}
    return problems, qids, name_maps, image_features


class MetaphorDatasetStd(Dataset):
    """
    Creating a custom dataset for reading the dataset and
    loading it into the dataloader to pass it to the
    neural network for finetuning the model

    """

    def __init__(
            self, problems, qids, tokenizer, answer_len, question_len, args, test_le=None
    ):
        self.tokenizer = tokenizer
        self.data = {qid: problems[qid] for qid in qids}
        self.answer_len = answer_len
        self.question_len = question_len
        self.questions = []
        self.answers = []
        if test_le is not None:
            test_le_data = json.load(open(test_le))["preds"]
        else:
            test_le_data = None
        idx = 0
        for qid in self.data:
            if test_le_data is not None:
                curr_le_data = test_le_data[idx]
                idx += 1
            else:
                curr_le_data = None
            question, answer = build_train_pair(problems, qid, args, curr_le_data)
            self.questions.append(question)
            self.answers.append(answer)

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

    def __getitem__(self, index):
        question_text = str(self.questions[index])
        answer_text = str(self.answers[index])
        


        question = self.tokenizer.batch_encode_plus(
            [question_text],
            max_length=self.question_len,
            pad_to_max_length=True,
            truncation=True,
            padding="max_length",
            return_tensors="pt",
        )
        answer = self.tokenizer.batch_encode_plus(
            [answer_text],
            max_length=self.answer_len,
            pad_to_max_length=True,
            truncation=True,
            padding="max_length",
            return_tensors="pt",
        )
        question_ids = question["input_ids"].squeeze()
        question_mask = question["attention_mask"].squeeze()
        answer_ids = answer["input_ids"].squeeze().tolist()

        return {
            "input_ids": question_ids,
            "attention_mask": question_mask,
            "labels": answer_ids,
        }

class MeatphorDatasetImg(Dataset):
    """
        Creating a custom dataset for reading the dataset and
        loading it into the dataloader to pass it to the
        neural network for finetuning the model

        """

    def __init__(
            self, problems, qids, name_maps, tokenizer, question_len, answer_len, args, image_features, test_le=None
    ):
        """
        Initializes a Dataset class

        Args:
            dataframe (pandas.DataFrame): Input dataframe
            tokenizer (transformers.tokenizer): Transformers tokenizer
            source_len (int): Max length of source text
            target_len (int): Max length of target text
            source_text (str): column name of source text
            target_text (str): column name of target text
        """
        self.tokenizer = tokenizer
        self.data = {qid: problems[qid] for qid in qids}
        self.answer_len = answer_len
        self.question_len = question_len
        self.questions = []
        self.answers = []
        self.image_ids = []
        if test_le is not None:
            test_le_data = json.load(open(test_le))["preds"]
        else:
            test_le_data = None
        idx = 0
        for qid in self.data:
            if test_le_data is not None:
                curr_le_data = test_le_data[idx]
                idx += 1
            else:
                curr_le_data = None
            question, answer = build_train_pair(problems, qid, args, curr_le_data)
            self.questions.append(question)
            self.answers.append(answer)
            pic_name = str(problems[qid]['Pic_id'])
            if str(pic_name) in name_maps:
                i_vectors = image_features[int(name_maps[str(pic_name)])]
                self.image_ids.append(i_vectors)
            else:
                shape = img_shape[args.img_type]
                self.image_ids.append(np.zeros(shape))

    def __len__(self):
        """returns the length of dataframe"""

        return len(self.answers)

    def __getitem__(self, index):
        """return the input ids, attention masks and target ids"""

        image_ids = self.image_ids[index]

        question_text = str(self.questions[index])
        answer_text = str(self.answers[index])

        question = self.tokenizer.batch_encode_plus(
            [question_text],
            max_length=self.question_len,
            pad_to_max_length=True,
            truncation=True,
            padding="max_length",
            return_tensors="pt",
        )
        answer = self.tokenizer.batch_encode_plus(
            [answer_text],
            max_length=self.answer_len,
            pad_to_max_length=True,
            truncation=True,
            padding="max_length",
            return_tensors="pt",
        )
        question_ids = question["input_ids"].squeeze()
        question_mask = question["attention_mask"].squeeze()
        answer_ids = answer["input_ids"].squeeze().tolist()

        image_ids = torch.tensor(image_ids).squeeze()

        return {
            "input_ids": question_ids,
            "attention_mask": question_mask,
            "labels": answer_ids,
            "image_ids": image_ids
        }