"""
pre-processing of images of plankton.
The amount of images are counted in the folder_location, so it is known which shape the 3d array should be.
A 3d array (amount_of_images, 428, 428) is then created and is filled with the grayscale data of the images.

Author: Lin Chang @ l.z.chang@st.hanze.nl
"""

import numpy as np
import glob
import sys
import os
import time
from matplotlib import pyplot as plt
import skimage.transform
import skimage.io


def count_files(location, folder_name):
    """Count the amount of images to create an 3d array with."""
    total_images = 0

    for folder in os.listdir(location):
        if folder.endswith(folder_name):
            total_images += (len(os.listdir(location + folder)))

    return total_images


def create_array_with_images(folder, folder_name):
    """Create an 3d Array and put the data of a single image in every row."""

    # Count the amount of files and create an array with amount_of_files * image_length * image_width
    total_images = count_files(folder, folder_name)
    array3d = np.ndarray([total_images, 428, 428])

    counter = 0

    for item in glob.glob(folder + "*/"):
        if item.endswith(folder_name + "/"):
            for file in glob.glob(item + "*"):
                # Create an empty array to paste the image on
                wall = np.zeros((428, 428), dtype=np.float)

                # Read image in as array in grayscale
                img = skimage.io.imread(file, as_grey=True)
                img[img == 255] = 0

                # Calculate the middle point of the image to paste on
                x = (wall.shape[0] - img.shape[0]) // 2
                y = (wall.shape[1] - img.shape[1]) // 2
                wall[x:x + img.shape[0], y:y + img.shape[1]] = img

                # Normalize to values between -0.5 and 0.5
                # normalized = (wall - 255 / 2) / 255  # weird output after augmentation
                normalized = wall / 255
                array3d[counter] = normalized

                counter += 1

    return array3d


def main(argv=None):
    if sys.argv is None:
        sys.argv = argv

    folder_location = "/commons/student/2017-2018/Thema11/Lin_Sylt/"

    start = time.time()

    train_data = create_array_with_images(folder_location, "train+")
    valid_data = create_array_with_images(folder_location, "valid+")
    test_data = create_array_with_images(folder_location, "test+")

    # Reshape into 2d array if needed

    a = np.reshape(train_data, (train_data.shape[0], train_data.shape[1] * train_data.shape[2]))

    print("took: {:.2f} seconds".format(time.time() - start))

    res = skimage.transform.rotate(valid_data[26], 45)
    plt.imshow(res)
    # plt.imshow(valid_data[26])
    plt.show()

    return 0


if __name__ == "__main__":
    sys.exit(main())