# import os.path
# import torchvision.transforms as transforms
# from data.base_dataset import BaseDataset, get_transform
from data.base_dataset import BaseDataset
# from data.image_folder import make_dataset
# from PIL import Image
# import PIL
import h5py
import random
import torch
import numpy

class TestDataset(BaseDataset):
    def initialize(self, opt):
        '''
        self.opt = opt
        self.root = opt.dataroot
        self.dir_A = os.path.join(opt.dataroot, opt.phase + 'A')
        self.dir_B = os.path.join(opt.dataroot, opt.phase + 'B')

        self.A_paths = make_dataset(self.dir_A)
        self.B_paths = make_dataset(self.dir_B)

        self.A_paths = sorted(self.A_paths)
        self.B_paths = sorted(self.B_paths)
        self.A_size = len(self.A_paths)
        self.B_size = len(self.B_paths)
        self.transform = get_transform(opt)
        '''

        self.opt = opt
        self.opt.A_paths_file = '/home/zhenghan/MRF/B'
        with open(self.opt.A_paths_file, 'r') as f:
            self.A_paths = f.read()
        self.A_paths = self.A_paths.replace('\n',' ').split()
        
        
        self.A_imgs = []
        self.B_imgs = []
        self.masks = []
        #self.opt.patchSize=0
        
        for A_path in self.A_paths:
            A_path1 = '/raid/zhenghan/data/MRF/data/'+A_path+'.mat'
            f = h5py.File(A_path1)
            A_img = numpy.transpose(f['imMRF']).astype('float32')
            #B_img = numpy.transpose(f['goal']).astype('float32')
            mask = numpy.transpose(f['mask']).astype('float32')
            f.close()
            
            if opt.goal_type == 'T2':
                A_path2='/raid/zhenghan/data/MRF/dataT2/'+A_path+'_T2.mat'
                #print(A_path2)
                f = h5py.File(A_path2)
                B_img = numpy.transpose(f['t2big']).astype('float32')
                maskT2 = numpy.transpose(f['maskT2']).astype('float32')
                
                f.close()
                mask = mask*maskT2
                
            if opt.goal_type == 'T1':
                A_path2='/raid/zhenghan/data/MRF/dataT1/'+A_path+'_T1.mat'
                
                f = h5py.File(A_path2)
                B_img = numpy.transpose(f['t1big']).astype('float32')
                maskT1 = numpy.transpose(f['maskT1']).astype('float32')
                
                f.close()
                mask = mask*maskT1
            
            
            A_img = A_img[53:201,58:229,:]
            B_img = B_img[53:201,58:229]
            mask = mask[53:201,58:229]
            
            
            
            A_img = numpy.concatenate((A_img[:,:,0:int(opt.input_nc/2)],A_img[:,:,2304:2304+int(opt.input_nc/2)]), axis=2)
            
            mask = mask[:,:,numpy.newaxis]
            if B_img.ndim==2:
                B_img = B_img[:,:,numpy.newaxis]
            assert A_img.ndim==3 and B_img.ndim==3, "# of dim is not 3 for training image"
            
            A_img = torch.from_numpy(A_img)
            B_img = torch.from_numpy(B_img)
            mask = torch.from_numpy(mask)
            
            if opt.data_GPU:
                A_img = A_img.cuda()
                B_img = B_img.cuda()
                mask = mask.cuda()
            

            A_img = A_img.permute(2,0,1)
            B_img = B_img.permute(2,0,1)
            mask = mask.permute(2,0,1)
            
            
            
            self.A_imgs.append(A_img)
            self.B_imgs.append(B_img)
            self.masks.append(mask)
            self.A_size = len(self.A_imgs)
            '''
            elif self.opt.patchSize==1:
                #print(self.input_A_t.shape)
                for k in range(mask.shape[1]):
                    for j in range(mask.shape[2]):
                        if mask[0,k,j]==1:
                            self.A_imgs.append(A_img[:,k:k+1,j:j+1])
                            self.B_imgs.append(B_img[:,k:k+1,j:j+1])
                            self.masks.append(mask[:,k:k+1,j:j+1])
                self.A_size = len(self.A_imgs)
                            
            else:
                self.A_imgs.append(A_img)
                self.B_imgs.append(B_img)
                self.masks.append(mask)
                self.A_size = int(len(self.A_imgs)*24990/(opt.patchSize**2))
                
            '''
            print("loaded image: %s" % A_path)
            self.num_imgs = len(self.A_paths)
        



    def __getitem__(self, index):
        '''
        A_path = self.A_paths[index % self.A_size]
        index_A = index % self.A_size
        index_B = random.randint(0, self.B_size - 1)
        B_path = self.B_paths[index_B]
        # print('(A, B) = (%d, %d)' % (index_A, index_B))
        A_img = Image.open(A_path).convert('RGB')
        B_img = Image.open(B_path).convert('RGB')

        A = self.transform(A_img)
        B = self.transform(B_img)
        if self.opt.which_direction == 'BtoA':
            input_nc = self.opt.output_nc
            output_nc = self.opt.input_nc
        else:
            input_nc = self.opt.input_nc
            output_nc = self.opt.output_nc

        if input_nc == 1:  # RGB to gray
            tmp = A[0, ...] * 0.299 + A[1, ...] * 0.587 + A[2, ...] * 0.114
            A = tmp.unsqueeze(0)

        if output_nc == 1:  # RGB to gray
            tmp = B[0, ...] * 0.299 + B[1, ...] * 0.587 + B[2, ...] * 0.114
            B = tmp.unsqueeze(0)
        return {'A': A, 'B': B,
                'A_paths': A_path, 'B_paths': B_path}
        '''

        
        index_A = index % self.A_size
        A_path = ''
        # print('(A, B) = (%d, %d)' % (index_A, index_B))
        A_img = self.A_imgs[index_A]
        B_img = self.B_imgs[index_A]
        mask = self.masks[index_A]
        '''
        elif self.opt.patchSize==1:
            index_A = index % self.A_size
            A_path = self.A_paths[0]
            # print('(A, B) = (%d, %d)' % (index_A, index_B))
            A_img = self.A_imgs[index_A]
            B_img = self.B_imgs[index_A]
            mask = self.masks[index_A]
            
        else:
            index_A = index % self.num_imgs
            A_path = self.A_paths[0]
            # print('(A, B) = (%d, %d)' % (index_A, index_B))
            A_img = self.A_imgs[index_A]
            B_img = self.B_imgs[index_A]
            mask = self.masks[index_A]
            
            # random crop
            
            patch_size = self.opt.patchSize
            
            A_position0, A_position1 = random.randint(0,A_img.shape[1]-patch_size), random.randint(0,A_img.shape[2]-patch_size)
            
            A_img = A_img[:, A_position0:A_position0+patch_size, A_position1:A_position1+patch_size]
            B_img = B_img[:, A_position0:A_position0+patch_size, A_position1:A_position1+patch_size]
            mask = mask[:, A_position0:A_position0+patch_size, A_position1:A_position1+patch_size]
            
            
            index_A = index // (256*256)
            p0, p1 = (index % (256*256)) // 256, (index % (256*256)) % 256
            
            mask = self.masks[index_A]
            mask = mask[:, p0:p0+1, p1:p1+1]
            
            
            A_path = self.A_paths[index_A]
            # print('(A, B) = (%d, %d)' % (index_A, index_B))
            A_img = self.A_imgs[index_A]
            B_img = self.B_imgs[index_A]
            
            
            A_img = A_img[:, p0:p0+1, p1:p1+1]
            B_img = B_img[:, p0:p0+1, p1:p1+1]
            
            '''
        
        '''
        f = h5py.File(A_path)

        A_img = numpy.transpose(f['imMRF']).astype('float32')
        B_img = numpy.transpose(f['t1big']).astype('float32')
        mask = numpy.transpose(f['mask']).astype('float32')
        
        mask = mask[:,:,numpy.newaxis]
        if B_img.ndim==2:
            B_img = B_img[:,:,numpy.newaxis]
        assert A_img.ndim==3 and B_img.ndim==3, "# of dim is not 3 for training image"
        '''
        '''
        A_img = A_img/255*2-1
        B_img = B_img/255*2-1

        # random crop
        if self.opt.patchSize0 <= 0 or self.opt.patchSize1 <= 0:
            patch_size0 = int((A_img.shape[0]*0.8))
            patch_size1 = int((A_img.shape[1]*0.8))
        else:
            patch_size0 = self.opt.patchSize0
            patch_size1 = self.opt.patchSize1
        A_position0, A_position1 = random.randint(0,A_img.shape[0]-patch_size0), random.randint(0,A_img.shape[1]-patch_size1)
        B_position0, B_position1 = random.randint(0,B_img.shape[0]-patch_size0), random.randint(0,B_img.shape[1]-patch_size1)
        A_img = A_img[A_position0:A_position0+patch_size0, A_position1:A_position1+patch_size1,:]
        B_img = B_img[B_position0:B_position0+patch_size0, B_position1:B_position1+patch_size1,:]
        '''
        '''
        A_img = torch.from_numpy(A_img)
        B_img = torch.from_numpy(B_img)
        mask = torch.from_numpy(mask)

        A_img = A_img.permute(2,0,1)
        B_img = B_img.permute(2,0,1)
        mask = mask.permute(2,0,1)
        '''

        return {'A': A_img, 'B': B_img, 'mask': mask,
                'A_paths': A_path}

    def __len__(self):
        
        return self.A_size
        

    def name(self):
        return 'TestDataset'
