import torch
from torch.utils.data import Dataset
from torchvision import transforms as T

import os
import json
import numpy as np
from PIL import Image
import h5py

import sys
BASE_DIR = os.path.dirname(os.path.abspath(__file__))
sys.path.append(os.path.join(BASE_DIR, '../../'))
from utils.direct_field.df_cardia import direct_field
from libs.datasets import augment as standard_aug
from utils.direct_field.utils_df import class2dist


class AcdcDataset(Dataset):
    def __init__(self, data_list, df_used=False, joint_augment=None, augment=None, target_augment=None, df_norm=True, boundary=False):
        self.joint_augment = joint_augment
        self.augment = augment
        self.target_augment = target_augment
        self.data_list = data_list
        self.df_used = df_used
        self.df_norm = df_norm
        self.boundary = boundary
        
        with open(data_list, 'r') as f:
            self.data_infos = json.load(f)

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

    def __getitem__(self,index):
        img = h5py.File(self.data_infos[index],'r')['image']
        gt = h5py.File(self.data_infos[index],'r')['label']
        # print(np.unique(gt))
        img = np.array(img)[:,:,None].astype(np.float32)
        gt = np.array(gt)[:,:,None].astype(np.float32)
        # print(np.unique(gt))

        if self.joint_augment is not None:
            img, gt = self.joint_augment(img, gt)
        if self.augment is not None:
            img = self.augment(img)
        if self.target_augment is not None:
            gt = self.target_augment(gt)

        if self.df_used:
            gt_df = direct_field(gt.numpy()[0], norm=self.df_norm)
            gt_df = torch.from_numpy(gt_df)
        else:
            gt_df = None
        
        if self.boundary:
            dist_map = torch.from_numpy(class2dist(gt.numpy()[0], C=4))
        else:
            dist_map = None

        return img, gt, gt_df, dist_map