from __future__ import absolute_import
from __future__ import division
from __future__ import print_function

import tensorflow as tf
import tensorflow.contrib.slim as slim

import numpy as np

from skimage import data, io, filters, util

import argparse
import sys
from os import listdir
from os.path import isfile, join

import time

def readImages(imageDir, numImages):
    files = [join(imageDir, f) for f in listdir(imageDir) if isfile(join(imageDir, f))]

    if (numImages < 0):
      numImages = len(files)

    images = np.array([np.array(util.img_as_float(io.imread((files[i]))), dtype=np.float32) for i in range(numImages)])
    return images


def readSubject(imageDir, numImages):
    tmpDir = imageDir + '/Correct'
    correctImages = readImages(tmpDir, numImages)
    correctLabel = np.resize([1., 0., 0.], (correctImages.shape[0], 3))
    print('Finished loading correct images ...')

    tmpDir = imageDir + '/OverSegmentation'
    overSegImages = readImages(tmpDir, numImages)
    overSegLabel = np.resize([0., 1., 0.], (overSegImages.shape[0], 3))
    print('Finished loading over segmentation images ...')

    tmpDir = imageDir + '/UnderSegmentation'
    underSegImages = readImages(tmpDir, numImages)
    underSegLabel = np.resize([0., 0., 1.], (underSegImages.shape[0], 3))
    print('Finished loading under segmentation images ...')

    images = np.concatenate((correctImages, overSegImages), axis=0)
    images = np.concatenate((images, underSegImages), axis=0)

    label = np.concatenate((correctLabel, overSegLabel), axis=0)
    label = np.concatenate((label, underSegLabel), axis=0)

    return images, label

def readUnclassifiedImages(imageDir, numImages):
    files = [join(imageDir, f) for f in listdir(imageDir) if isfile(join(imageDir, f))]

    if (numImages < 0):
        numImages = len(files)

    images = np.array([np.array(util.img_as_float(io.imread((files[i]))), dtype=np.float32) for i in range(numImages)])

    label = np.resize([1., 0., 0.], (images.shape[0], 3))
    print('Finished loading images ...')

    return images, label, files


def tfTest():
    rootDir = '/home/daten_bartschat/Dataset/Training'

    trainImages, trainLabel = readSubject(rootDir + '/EGT3_002-L/')

    images = tf.convert_to_tensor(trainImages)
    images = tf.reshape(images, trainImages.shape)

    label = tf.convert_to_tensor(trainLabel)
    label = tf.reshape(label, trainLabel.shape)

    # step 4: Batching
    image_batch, label_batch = tf.train.batch([trainImages, trainLabel], batch_size=8)

    return trainImages, trainLabel

