from torch.utils.data import Dataset
from PIL import Image
import os
import io
import json
import random
import torch
import numpy as np
from einops import rearrange
try:
    from aoss_client.client import Client
except:
    try:
        from petrel_client.client import Client
    except:
        Client = None
from glob import glob
from xtuner.registry import BUILDER
from xtuner.dataset.utils import expand2square
from src.datasets.utils import crop2square, encode_fn
from xtuner.utils import DEFAULT_IMAGE_TOKEN, IMAGE_TOKEN_INDEX
from src.datasets.understanding.caption_prompts import dense_prompts, short_prompts

from datasets import load_from_disk
from divergent_memories.data.constants import SYNTHETIC_SHAPES, SYNTHETIC_COLORS, SYNTHETIC_PATTERNS, SYNTHETIC_POSITIONS


class CaptionDataset(Dataset):
    def __init__(self,
                 data_path,
                #  local_folder,
                 image_size,
                #  ceph_folder=None,
                #  ceph_config=None,
                 tokenizer=None,
                 template_map_fn=None,
                 max_length=2048,
                 min_image_size=80,
                 image_length=256,
                 pad_image=True,
                 brief=False,
                #  cap_folder=None,
                #  cap_prompt_folder=None,
                #  cap_source='caption',
                 ):
        super().__init__()
        self.data_path = data_path
        # self._load_data(data_path)
        # self.local_folder = local_folder
        # self.cap_folder = local_folder if cap_folder is None else cap_folder
        # self.cap_source = cap_source

        # self.cap_prompt_folder = cap_prompt_folder

        self.image_size = image_size

        self.tokenizer = BUILDER.build(tokenizer)
        self.prompt_template = template_map_fn['template']
        self.template_map_fn = BUILDER.build(template_map_fn)
        self.max_length = max_length
        self.image_length = image_length
        self.pad_image = pad_image
        self.min_image_size = min_image_size

        # self.FILE_CLIENT = None
        # self.ceph_folder = ceph_folder
        # self.ceph_config = ceph_config
        # self.use_ceph = ((Client is not None) and (ceph_folder is not None)
        #                  and (ceph_config is not None) and os.path.exists(ceph_config))

        self.brief = brief
        self.caption_prompts = short_prompts if self.brief else dense_prompts

        def map_1():


        def map_2(data_sample):
            pixel_values = self._process_image(data_sample["image"])["pixel_values"]
            return {
                'input_ids': input_ids,
                'attention_mask': attention_mask,
                'labels': labels,
                'pixel_values': torch.stack(pixel_values_list) if len(pixel_values_list) > 0 else None
        }

        self.images_data = load_from_disk(data_path)['train']
        self.data_list = self.data_list.map(map_1, batched=True, batch_size=1, remove_columns=["synthetic_color", "synthetic_pattern", "synthetic_position", "synthetic_shape"])

    # def _load_data(self, data_path: str):      # image path and annotation path are saved in a json file
    #     if data_path.endswith('.json'):
    #         with open(data_path, 'r') as f:
    #             self.data_list = json.load(f)
    #     else:
    #         json_files = glob(f"{data_path}/*.json")
    #         data_list = []
    #         for json_file in json_files:
    #             with open(json_file, 'r') as f:
    #                 data_list += json.load(f)

    #         self.data_list = data_list

    #     print(f"Load {len(self.data_list)} data samples from {data_path}", flush=True)

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

    # def _read_ceph(self, ceph_path):
    #     if self.FILE_CLIENT is None:
    #         self.FILE_CLIENT = Client(self.ceph_config)
    #     data_bytes = self.FILE_CLIENT.get(ceph_path)

    #     return io.BytesIO(data_bytes)

    # def _read_image(self, image_file):
    #     if self.use_ceph:
    #         image = Image.open(
    #             self._read_ceph(
    #                 os.path.join(self.ceph_folder, image_file)
    #             )
    #         )
    #     else:
    #         image = Image.open(
    #             os.path.join(self.local_folder, image_file)
    #         )
    #     assert image.width > self.min_image_size and image.height > self.min_image_size, f"Image: {image.size}"
    #     assert image.width / image.height > 0.1, f"Image: {image.size}"
    #     assert image.width / image.height < 10, f"Image: {image.size}"
    #     return image.convert('RGB')

    # def _read_json(self, annotation_file):
    #     if self.use_ceph:
    #         annotation = json.load(
    #             self._read_ceph(
    #                 os.path.join(self.ceph_folder, annotation_file)
    #             )
    #         )
    #     else:
    #         with open(os.path.join(self.local_folder, annotation_file), 'r') as f:
    #             annotation = json.load(f)

    #     return annotation

    def _process_image(self, image):
        data = dict()
        if self.pad_image:
            image = expand2square(image, (127, 127, 127))
        else:
            image = crop2square(image)

        image = image.resize(size=(self.image_size, self.image_size))
        pixel_values = torch.from_numpy(np.array(image)).float()
        pixel_values = pixel_values / 255
        pixel_values = 2 * pixel_values - 1
        pixel_values = rearrange(pixel_values, 'h w c -> c h w')

        data.update(pixel_values=pixel_values)
        return data

    # def _process_text(self, text):
    #     assert DEFAULT_IMAGE_TOKEN not in text, text
    #     data_dict = dict(conversation=[{'input': f"{DEFAULT_IMAGE_TOKEN}\n{random.choice(self.caption_prompts)}",
    #                                     'output': text.strip()}])
    #     data_dict.update(self.template_map_fn(data_dict))
    #     data_dict.update(encode_fn(data_dict, self.tokenizer, self.max_length,
    #                                self.image_length, True, True))

    #     assert (torch.tensor(data_dict['input_ids']).long() == IMAGE_TOKEN_INDEX).sum() == self.image_length, \
    #         "Error in image format"

    #     data_dict['type'] = 'image2text'
    #     return data_dict

    def _process_text(self, caption_prompt, text):
        assert DEFAULT_IMAGE_TOKEN not in text, text
        data_dict = dict(conversation=[{'input': f"{DEFAULT_IMAGE_TOKEN}\n{caption_prompt}",
                                         'output': text.strip()}])
        data_dict.update(self.template_map_fn(data_dict))
        data_dict.update(encode_fn(data_dict, self.tokenizer, None,
                                    self.image_length, True, True))

        assert (torch.tensor(data_dict['input_ids']).long() == IMAGE_TOKEN_INDEX).sum() == self.image_length, \
            "Error in image format"

        data_dict['type'] = 'image2text'
        return data_dict

    def _retry(self):
        return self.__getitem__(random.choice(range(self.__len__())))

    def __getitem__(self, idx):
        try:
            data_sample = self.data_list[idx]
            image = self._read_image(data_sample['image']).convert('RGB')
            data = self._process_image(image)
            del image

            if 'caption_prompt' in data_sample:
                with open(f"{self.cap_prompt_folder}/{data_sample['caption_prompt']}", 'r') as f:
                    caption_prompt = json.load(f)["caption_prompt"]
            else:
                caption_prompt = random.choice(self.caption_prompts)

            with open(f"{self.cap_folder}/{data_sample['annotation']}", 'r') as f:
                caption = json.load(f)[self.cap_source]
            data.update(self._process_text(caption_prompt, caption))

            data.update(image_dir=self.local_folder, image_file=data_sample['image'])

            return data

        except Exception as e:
            print(f"Error when reading {self.data_path}:{data_sample['image']}: {e}", flush=True)
            return self._retry()
