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

from utils.utils import generate_mask
from PIL import Image

class TrainDataset(torch.utils.data.Dataset):
    def __init__(self, data_path, transform=None):
        self.data = os.listdir(os.path.join(data_path, 'color'))
        self.data_path = data_path
        self.transform = transform
        self.ToTensor = transforms.ToTensor()

        # Directorio para guardar las imágenes en blanco y negro
        self.bw_directory = os.path.join(data_path, 'bw')
        if not os.path.exists(self.bw_directory):
            os.makedirs(self.bw_directory)

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

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

        color_img = Image.open(os.path.join(self.data_path, 'color', image_name)).convert('RGB')

        bw_name = "bw_" + image_name
        dfm_name = 'dfm_' + image_name

        bw_img = np.expand_dims(np.array(Image.open(os.path.join(self.bw_directory, bw_name)).convert('L')), 2)
        dfm_img = np.expand_dims(np.array(Image.open(os.path.join(self.bw_directory, dfm_name)).convert('L')), 2)

        bw_img = np.concatenate([bw_img, dfm_img], axis=2)

        if self.transform:
            result = self.transform(image=np.array(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):
        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'))] * 100
        self.data_path = data_path
        self.transform = transform

        # Directorio para guardar las imágenes en blanco y negro
        self.bw_directory = os.path.join(data_path, 'bw')
        if not os.path.exists(self.bw_directory):
            os.makedirs(self.bw_directory)

        np.random.shuffle(self.color_data)
        self.ToTensor = transforms.ToTensor()

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

    def __getitem__(self, idx):
        color_img = Image.open(os.path.join(self.data_path, 'color', self.color_data[idx])).convert('RGB')
        bw_name = "bw_" + self.data[idx]
        dfm_name = "dfm_" + self.data[idx]

        bw_img = np.expand_dims(np.array(Image.open(os.path.join(self.bw_directory, bw_name)).convert('L')), 2)
        dfm_img = np.expand_dims(np.array(Image.open(os.path.join(self.bw_directory, dfm_name)).convert('L')), 2)

        if self.transform:
            result = self.transform(image=np.array(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