import os
import numpy as np

from PIL import Image
from re import split, compile
from tensorflow.keras.utils import Sequence


def list_filenames(data_path, img_extension='png', filename_prefix=None):
    if filename_prefix is None:
        files_list = [file for file in os.listdir(data_path) if file.endswith(img_extension)]
    else:
        files_list = [file for file in os.listdir(data_path) if file.endswith(img_extension) and file.startswith(filename_prefix)]

    files_list.sort(key=lambda l: [int(s) if s.isdigit() else s.lower() for s in split(compile(r'(\d+)'), l)])
    files_list = [os.path.join(data_path, file) for file in files_list]
    print('Found {} files in {}'.format(len(files_list), data_path))
    return files_list


class Dataset(Sequence):
    def __init__(self, file_list, batch_size=32, crop_dim=None, resize_dim=None, shuffle=True, mode='RGB'):
        self.files_list = file_list
        self.batch_size = batch_size

        self.crop_dim = crop_dim
        self.resize_dim = resize_dim
        self.shuffle = shuffle
        self.on_epoch_end()

        self.mode=mode

    def __len__(self):
        return int(np.ceil(len(self.files_list) / self.batch_size))

    def __getitem__(self, idx):
        batch_files = self.files_list[idx * self.batch_size : (idx + 1) * self.batch_size]
        images = [self.load_images(f) for f in batch_files]
        return np.stack(images)

    def on_epoch_end(self):
        if self.shuffle:
            np.random.shuffle(self.files_list)

    @staticmethod
    def center_crop(image, crop_dim):
        h, w = image.size
        crop_h, crop_w = crop_dim

        top = max(0, (w - crop_w) // 2)
        left = max(0, (h - crop_h) // 2)
        right = min(h - 0, (h + crop_h) // 2)
        bottom = min(w - 0, (w + crop_w) // 2)

        return image.crop((left, top, right, bottom))

    def load_images(self, filepath):
        if self.mode=='RGB':
            image = Image.open(filepath).convert('RGB')
        else:
            image = Image.open(filepath)
        if self.crop_dim:
            image = self.center_crop(image, crop_dim=self.crop_dim)
        if self.resize_dim:
            image = image.resize(self.resize_dim)

        image = np.array(image).astype(np.float32)
        image = image / 255.0
        if image.ndim == 2:
            image = np.expand_dims(image, -1)
        return image