import os
import zipfile

import h5py
import numpy
from PIL import Image

from fuel.converters.base import check_exists, progress_bar
from fuel.datasets.hdf5 import H5PYDataset

TRAIN = 'dogs_vs_cats.train.zip'
TEST = 'dogs_vs_cats.test1.zip'


@check_exists(required_files=[TRAIN, TEST])
def convert_dogs_vs_cats(directory, output_directory,
                         output_filename='dogs_vs_cats.hdf5'):
    """Converts the Dogs vs. Cats dataset to HDF5.

    Converts the Dogs vs. Cats dataset to an HDF5 dataset compatible with
    :class:`fuel.datasets.dogs_vs_cats`. The converted dataset is saved as
    'dogs_vs_cats.hdf5'.

    It assumes the existence of the following files:

    * `dogs_vs_cats.train.zip`
    * `dogs_vs_cats.test1.zip`

    Parameters
    ----------
    directory : str
        Directory in which input files reside.
    output_directory : str
        Directory in which to save the converted dataset.
    output_filename : str, optional
        Name of the saved dataset. Defaults to 'dogs_vs_cats.hdf5'.

    Returns
    -------
    output_paths : tuple of str
        Single-element tuple containing the path to the converted dataset.

    """
    # Prepare output file
    output_path = os.path.join(output_directory, output_filename)
    h5file = h5py.File(output_path, mode='w')
    dtype = h5py.special_dtype(vlen=numpy.dtype('uint8'))
    hdf_features = h5file.create_dataset('image_features', (37500,),
                                         dtype=dtype)
    hdf_shapes = h5file.create_dataset('image_features_shapes', (37500, 3),
                                       dtype='int32')
    hdf_labels = h5file.create_dataset('targets', (37500, 1), dtype='uint8')

    # Attach shape annotations and scales
    hdf_features.dims.create_scale(hdf_shapes, 'shapes')
    hdf_features.dims[0].attach_scale(hdf_shapes)

    hdf_shapes_labels = h5file.create_dataset('image_features_shapes_labels',
                                              (3,), dtype='S7')
    hdf_shapes_labels[...] = ['channel'.encode('utf8'),
                              'height'.encode('utf8'),
                              'width'.encode('utf8')]
    hdf_features.dims.create_scale(hdf_shapes_labels, 'shape_labels')
    hdf_features.dims[0].attach_scale(hdf_shapes_labels)

    # Add axis annotations
    hdf_features.dims[0].label = 'batch'
    hdf_labels.dims[0].label = 'batch'
    hdf_labels.dims[1].label = 'index'

    # Convert
    i = 0
    for split, split_size in zip([TRAIN, TEST], [25000, 12500]):
        # Open the ZIP file
        filename = os.path.join(directory, split)
        zip_file = zipfile.ZipFile(filename, 'r')
        image_names = zip_file.namelist()[1:]  # Discard the directory name

        # Shuffle the examples
        rng = numpy.random.RandomState(123522)
        rng.shuffle(image_names)

        # Convert from JPEG to NumPy arrays
        with progress_bar(filename, split_size) as bar:
            for image_name in image_names:
                # Save image
                image = numpy.array(Image.open(zip_file.open(image_name)))
                image = image.transpose(2, 0, 1)
                hdf_features[i] = image.flatten()
                hdf_shapes[i] = image.shape

                # Cats are 0, Dogs are 1
                hdf_labels[i] = 0 if 'cat' in image_name else 1

                # Update progress
                i += 1
                bar.update(i if split == TRAIN else i - 25000)

    # Add the labels
    split_dict = {}
    sources = ['image_features', 'targets']
    for name, slice_ in zip(['train', 'test'],
                            [(0, 25000), (25000, 37500)]):
        split_dict[name] = dict(zip(sources, [slice_] * len(sources)))
    h5file.attrs['split'] = H5PYDataset.create_split_array(split_dict)

    h5file.flush()
    h5file.close()

    return (output_path,)


def fill_subparser(subparser):
    """Sets up a subparser to convert the dogs_vs_cats dataset files.

    Parameters
    ----------
    subparser : :class:`argparse.ArgumentParser`
        Subparser handling the `dogs_vs_cats` command.

    """
    return convert_dogs_vs_cats
