from skimage.io import imread
from skimage.transform import resize
import numpy as np
import tensorflow as tf
import keras
from skimage.transform import resize
from tf_record_utility import TFRecordUtility
from configuration import DatasetName, DatasetType, \
    AffectnetConf, D300wConf, W300Conf, InputDataSize, LearningConfig
from numpy import save, load, asarray


class CustomHeatmapGenerator(keras.utils.Sequence):

    def __init__(self, is_single, image_filenames, label_filenames, batch_size, n_outputs, accuracy=100):
        self.image_filenames = image_filenames
        self.label_filenames = label_filenames
        self.batch_size = batch_size
        self.n_outputs = n_outputs
        self.is_single = is_single
        self.accuracy = accuracy

    def __len__(self):
        _len = np.ceil(len(self.image_filenames) // float(self.batch_size))
        return int(_len)

    def __getitem__(self, idx):
        img_path = D300wConf.train_images_dir
        tr_path_85 = D300wConf.train_hm_dir_85
        tr_path_90 = D300wConf.train_hm_dir_90
        tr_path_97 = D300wConf.train_hm_dir_97
        tr_path = D300wConf.train_hm_dir

        batch_x = self.image_filenames[idx * self.batch_size:(idx + 1) * self.batch_size]
        batch_y = self.label_filenames[idx * self.batch_size:(idx + 1) * self.batch_size]

        img_batch = np.array([imread(img_path + file_name) for file_name in batch_x])

        if self.is_single:
            if self.accuracy == 85:
                lbl_batch = np.array([load(tr_path_85 + file_name) for file_name in batch_y])
            elif self.accuracy == 90:
                lbl_batch = np.array([load(tr_path_90 + file_name) for file_name in batch_y])
            elif self.accuracy == 97:
                lbl_batch = np.array([load(tr_path_97 + file_name) for file_name in batch_y])
            else:
                lbl_batch = np.array([load(tr_path + file_name) for file_name in batch_y])

            lbl_out_array = lbl_batch
        else:
            lbl_batch_85 = np.array([load(tr_path_85 + file_name) for file_name in batch_y])
            lbl_batch_90 = np.array([load(tr_path_90 + file_name) for file_name in batch_y])
            lbl_batch_97 = np.array([load(tr_path_97 + file_name) for file_name in batch_y])
            lbl_out_array = [lbl_batch_85, lbl_batch_90, lbl_batch_97]

        return img_batch, lbl_out_array