import json
import cv2
import numpy as np
import os

from torch.utils.data import Dataset
import pycocotools.mask as maskUtils
from torchvision import transforms
import utils.transforms as custom_transforms
from PIL import Image
class SAMDataset(Dataset):
    def __init__(self, data_path='../data/files', txt_path='../data/data_85616.txt'):
        self.data = []
        with open(txt_path, 'rt') as f:
            for line in f:
                self.data.append(eval(line))
        self.data_path = data_path
        randomresizedcrop = custom_transforms.RandomResizedCrop(
            512,
            scale=(0.9, 1),
        )
        self.transform = custom_transforms.Compose([
            randomresizedcrop,
            custom_transforms.RandomHorizontalFlip(p=0.5),
            custom_transforms.ToTensor(),
            transforms.Normalize(mean=0.5, std=0.5)
        ])



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

    def load_rle_annotations_from_json(self, json_file_path, return_pil=True):
        with open(json_file_path, 'r', encoding='utf-8') as f:
            anno_data = json.load(f)
        annotations = anno_data['annotations']
        height = int(anno_data['image']['height'])
        width = int(anno_data['image']['width'])

        map = np.zeros((height,width), dtype=np.uint16)
        for i in range(len(annotations)):
            ann = annotations[i]
            mask = maskUtils.decode(ann['segmentation'])
            map[mask != 0] = i + 1
        if return_pil:
            res = np.zeros((map.shape[0], map.shape[1], 3))
            res[:, :, 0] = map % 256
            res[:, :, 1] = map // 256
            res = Image.fromarray(res.astype(np.uint8))
            return res
        return map

    def __getitem__(self, idx):
        item = self.data[idx]

        source_filename = item['source']
        target_filename = item['target']
        prompt = item['prompt']


        source = self.load_rle_annotations_from_json(os.path.join(self.data_path, source_filename))
        target = Image.open(os.path.join(self.data_path, target_filename))


        target, source = self.transform(target, source)

        print(source.max(), source.min())
        target = target.permute(1,2,0)

        return dict(jpg=target, txt=prompt, hint=source)