from aix import Dataset
import aix.constants as C
import os
import numpy as np
import cv2

def load_image(path):

    img = cv2.imread(path, cv2.IMREAD_GRAYSCALE)
    array = np.reshape(img, (img.shape[0], img.shape[1], -1)) / 255

    return array

def temp_numpy_masks(training_fold):

    image_rows = C.IMG_WIDTH
    image_cols = C.IMG_HEIGHT

    X_train = np.ndarray((len(training_fold[0]), image_rows, image_cols, 1), dtype = np.float64)
    y_train = np.ndarray((len(training_fold[1]), image_rows, image_cols, 1), dtype = np.float64)
    X_test = np.ndarray((len(training_fold[2]), image_rows, image_cols, 1), dtype = np.float64)
    y_test = np.ndarray((len(training_fold[3]), image_rows, image_cols, 1), dtype = np.float64)

    for i, obj in enumerate(training_fold[0]):

        img = load_image(obj)
        img = cv2.resize(img, (image_cols, image_rows), interpolation = cv2.INTER_AREA)
        img = np.expand_dims(img, axis = 2)
        X_train[i] = img

    for i, obj in enumerate(training_fold[1]):

        mask = load_image(obj)
        mask = cv2.resize(mask, (image_cols, image_rows), interpolation = cv2.INTER_AREA)
        mask = np.expand_dims(mask, axis = 2)
        y_train[i] = mask

    for i, obj in enumerate(training_fold[2]):

        img = load_image(obj)
        img = cv2.resize(img, (image_cols, image_rows), interpolation = cv2.INTER_AREA)
        img = np.expand_dims(img, axis = 2)
        X_test[i] = img

    for i, obj in enumerate(training_fold[3]):

        mask = load_image(obj)
        mask = cv2.resize(mask, (image_cols, image_rows), interpolation = cv2.INTER_AREA)
        mask = np.expand_dims(mask, axis = 2)
        y_test[i] = mask

    return X_train, y_train, X_test, y_test

def load_training_batch(name = 'average_model', oocytes_list = [], images_path = '', annotations_path = ''):

    dataset = Dataset(name, oocytes_list, images_path, annotations_path)

    gen = dataset.train_test_iterator()

    j = 0
    folds = []
    for x_train, x_test, y_train, y_test in gen:

        j += 1
        print(j)
        folds.append(temp_numpy_masks([x_train, y_train, x_test, y_test]))

        # save temporary numpies??

        # DO MODEL stuff in here

    return folds