# ======================================================================== #
#
# Copyright (c) 2017 - 2020 scVAE authors
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#    http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
#
# ======================================================================== #

import os

import numpy
import PIL
import scipy.sparse

from scvae.analyses.figures import saving

IMAGE_EXTENSION = ".png"
DEFAULT_NUMBER_OF_RANDOM_EXAMPLES_FOR_COMBINED_IMAGES = 100


def combine_images_from_data_set(data_set, indices=None,
                                 number_of_random_examples=None, name=None):

    image_name = saving.build_figure_name("random_image_examples", name)
    random_state = numpy.random.RandomState(13)

    if indices is not None:
        n_examples = len(indices)
        if number_of_random_examples is not None:
            n_examples = min(n_examples, number_of_random_examples)
            indices = random_state.permutation(indices)[:n_examples]
    else:
        if number_of_random_examples is not None:
            n_examples = number_of_random_examples
        else:
            n_examples = DEFAULT_NUMBER_OF_RANDOM_EXAMPLES_FOR_COMBINED_IMAGES
        indices = random_state.permutation(
            data_set.number_of_examples)[:n_examples]

    if n_examples == 1:
        image_name = saving.build_figure_name("image_example", name)
    else:
        image_name = saving.build_figure_name("image_examples", name)

    width, height = data_set.feature_dimensions

    examples = data_set.values[indices]
    if scipy.sparse.issparse(examples):
        examples = examples.A
    examples = examples.reshape(n_examples, width, height)

    column = int(numpy.ceil(numpy.sqrt(n_examples)))
    row = int(numpy.ceil(n_examples / column))

    image = numpy.zeros((row * width, column * height))

    for m in range(n_examples):
        c = int(m % column)
        r = int(numpy.floor(m / column))
        rows = slice(r*width, (r+1)*width)
        columns = slice(c*height, (c+1)*height)
        image[rows, columns] = examples[m]

    return image, image_name


def save_image(image, name, directory):

    if not os.path.exists(directory):
        os.makedirs(directory)

    minimum = image.min()
    maximum = image.max()
    if 0 < minimum and minimum < 1 and 0 < maximum and maximum < 1:
        rescaled_image = 255 * image
    else:
        rescaled_image = (255 / (maximum - minimum) * (image - minimum))

    image = PIL.Image.fromarray(rescaled_image.astype(numpy.uint8))

    name += IMAGE_EXTENSION
    image_path = os.path.join(directory, name)
    image.save(image_path)