# Loads an image with the tensorflow input pipeline
import glob
import os
import tensorflow as tf
import hypergan.inputs.resize_image_patch
from tensorflow.python.ops import array_ops
from hypergan.gan_component import ValidationException, GANComponent

class ImageLoader:
    """
    ImageLoader loads a set of images into a tensorflow input pipeline.
    """

    def __init__(self, batch_size):
        self.batch_size = batch_size

    def build_labels(self, dirs):
        total_labels=0
        labels = {}
        for dir in dirs:
            labels[dir.split('/')[-1]]=total_labels
            total_labels+=1
        if(len(dirs) == 1):
            labels = {}
            total_labels = 0
        return labels,total_labels

    def create(self, directory, channels=3, format='jpg', width=64, height=64, crop=False, resize=False):
        directories = glob.glob(directory+"/*")
        directories = [d for d in directories if os.path.isdir(d)]

        if(len(directories) == 0):
            directories = [directory] 
        labels,total_labels = self.build_labels(sorted(filter(os.path.isdir, directories)))

        # Create a queue that produces the filenames to read.
        if(len(directories) == 1):
            # No subdirectories, use all the images in the passed in path
            filenames = glob.glob(directory+"/*."+format)
            classes = [0 for f in filenames]
        else:
            filenames = glob.glob(directory+"/**/*."+format)
            classes = [labels[f.split('/')[-2]] for f in filenames]

        print("[loader] ImageLoader found", len(filenames), "images with", total_labels, "different class labels")
        self.file_count = len(filenames)
        if self.file_count == 0:
            raise ValidationException("No images found in '" + directory + "'")
        filenames = tf.convert_to_tensor(filenames, dtype=tf.string)
        classes = tf.convert_to_tensor(classes, dtype=tf.int32)

        input_queue = tf.train.slice_input_producer([filenames, classes])

        # Read examples from files in the filename queue.
        value = tf.read_file(input_queue[0])

        if format == 'jpg':
            img = tf.image.decode_jpeg(value, channels=channels)
        elif format == 'png':
            img = tf.image.decode_png(value, channels=channels)
        else:
            print("[loader] Failed to load format", format)
        img = tf.cast(img, tf.float32)

        label = input_queue[1]

      # Image processing for evaluation.
      # Crop the central [height, width] of the image.
        if crop:
            resized_image = hypergan.inputs.resize_image_patch.resize_image_with_crop_or_pad(img, height, width, dynamic_shape=True)
        elif resize:
            resized_image = tf.image.resize_images(img, [height, width], 1)
        else: 
            resized_image = img

        tf.Tensor.set_shape(resized_image, [height,width,channels])

        # This moves the image to a range of -1 to 1.
        float_image = resized_image / 127.5 - 1.

        # Ensure that the random shuffling has good mixing properties.
        min_fraction_of_examples_in_queue = 0.4

        # Generate a batch of images and labels by building up a queue of examples.
        x,y = self._get_data(float_image, label)

        y = tf.cast(y, tf.int64)
        y = tf.one_hot(y, total_labels, 1.0, 0.0)
        self.x = x
        self.y = y
        return x, y

    def _get_data(self, image, label):
        batch_size = self.batch_size
        num_preprocess_threads = 24
        images, label_batch = tf.train.shuffle_batch(
            [image, label],
            batch_size=batch_size,
            num_threads=num_preprocess_threads,
            capacity= batch_size*10,
            min_after_dequeue=batch_size)
        return images, tf.reshape(label_batch, [batch_size])