import torch
import os
import torchvision.transforms as transforms
import matplotlib.pyplot as plt
import numpy as np

from utils.utils import generate_mask


class TrainDataset(torch.utils.data.Dataset):
    def __init__(self, data_path, transform = None, mults_amount = 1):
        self.data = os.listdir(os.path.join(data_path, 'color'))
        self.data_path = data_path
        self.transform = transform
        self.mults_amount = mults_amount
        
        self.ToTensor = transforms.ToTensor()
    def __len__(self):
        return len(self.data)
    
    def __getitem__(self, idx):
        image_name = self.data[idx]
        
        color_img = plt.imread(os.path.join(self.data_path, 'color', image_name))
        

        if self.mults_amount > 1:
            mult_number = np.random.choice(range(self.mults_amount))
            
            bw_name = image_name[:image_name.rfind('.')] + '_' + str(mult_number) + '.png'
            dfm_name = image_name[:image_name.rfind('.')] + '_' + str(mult_number) + '_dfm.png'
        else:
            bw_name = self.data[idx]
            dfm_name =  os.path.splitext(self.data[idx])[0] + '0_dfm.png'
            
            
        bw_img =  np.expand_dims(plt.imread(os.path.join(self.data_path, 'bw', bw_name)), 2)
        dfm_img =  np.expand_dims(plt.imread(os.path.join(self.data_path, 'bw', dfm_name)), 2)
        
        bw_img = np.concatenate([bw_img, dfm_img], axis = 2)
        
        if self.transform:
            result = self.transform(image = color_img, mask = bw_img)
            color_img = result['image']
            bw_img = result['mask']
          
        dfm_img = bw_img[:, :, 1]
        bw_img = bw_img[:, :, 0]
        
        color_img = self.ToTensor(color_img)
        bw_img = self.ToTensor(bw_img)
        
        dfm_img = self.ToTensor(dfm_img)
        
        color_img = (color_img - 0.5) / 0.5
        
        mask = generate_mask(bw_img.shape[1], bw_img.shape[2])
        hint = torch.cat((color_img * mask, mask), 0)
        
        return bw_img, color_img, hint, dfm_img
    
class FineTuningDataset(torch.utils.data.Dataset):
    def __init__(self, data_path, transform = None, mult_amount = 1):
        self.data = [x for x in os.listdir(os.path.join(data_path, 'real_manga')) if x.find('_dfm') == -1] 
        self.color_data = [x for x in os.listdir(os.path.join(data_path, 'color'))]
        self.data_path = data_path
        self.transform = transform
        self.mults_amount = mult_amount
        
        np.random.shuffle(self.color_data)
        
        self.ToTensor = transforms.ToTensor()
    def __len__(self):
        return len(self.data)
    
    def __getitem__(self, idx):
        color_img = plt.imread(os.path.join(self.data_path, 'color', self.color_data[idx]))
        
        image_name = self.data[idx]
        if self.mults_amount > 1:
            mult_number = np.random.choice(range(self.mults_amount))
            
            bw_name = image_name[:image_name.rfind('.')] + '_' + str(self.mults_amount) + '.png'
            dfm_name = image_name[:image_name.rfind('.')] + '_' + str(self.mults_amount) + '_dfm.png'
        else:
            bw_name = self.data[idx]
            dfm_name =  os.path.splitext(self.data[idx])[0] + '_dfm.png'
        
        
        bw_img =  np.expand_dims(plt.imread(os.path.join(self.data_path, 'real_manga', image_name)), 2)
        dfm_img =  np.expand_dims(plt.imread(os.path.join(self.data_path, 'real_manga', dfm_name)), 2)
        
        if self.transform:
            result = self.transform(image = color_img)
            color_img = result['image']
            
            result = self.transform(image = bw_img,  mask = dfm_img)
            bw_img = result['image']
            dfm_img = result['mask']
        
        color_img = self.ToTensor(color_img)
        bw_img = self.ToTensor(bw_img)
        dfm_img = self.ToTensor(dfm_img)
        
        color_img = (color_img - 0.5) / 0.5
        
        return bw_img, dfm_img, color_img