import tarfile
import keras
import numpy as np

def raw_generator_with_label_from_tar(tar_fn, files, labels):
	'''
	# Functionality
		Generate (data, label) pairs and feed into keras fit_generator
	# Arguments
		tar_fn: string. The local file location of the tar file. (e.g., 102flowers.tgz)
		files:  list. list of file names in the tar
		labels: list.
	# Returns
		A generator.
	'''

	label_by_fn = dict(zip(files, labels))
	with tarfile.open(tar_fn) as f:
		while True:
			m = f.next()
			if m is None:
				break
			if m.name in label_by_fn:
				yield f.extractfile(m).read(), label_by_fn[m.name]



def batch_generator(items, batch_size):
	'''
	# Functionality
		Implement batch generator that yields items in batches of size batch_size.
    	There's no need to shuffle input items, just chop them into batches.
	# Arguments
		items: any iterable (list, generator, ...). You should do `for item in items: ...`
    		   In case of generator you can pass through your items only once!
		batch_size:  integer.
	# Returns
		A generator.
	'''

	current_batch = []
	counter = 0

	for item in items:
	    no_remainder = False
	    counter = counter + 1

	    if counter == batch_size:
	        current_batch.append(item)
	        yield current_batch
	        current_batch = []
	        counter = 0
	        no_remainder = True
	    else:
	        current_batch.append(item)

	if not no_remainder:
		yield current_batch



def train_and_test_generator(items, batch_size, preprocess_func, num_classes):
	'''
    # Functionality
        Implement train or test generator that add preprocessing to batches generated by batch_generator
    # Arguments
        items: any iterable (list, generator, ...). You should do `for item in items: ...`
    		   In case of generator you can pass through your items only once!
		batch_size:  integer.
        preprocess_func:  Python function. preprocessing function to apply to each image.
        num_classes: integer.
    # Returns
        A generator.
    '''
	batch_gen = batch_generator(items, batch_size)
	while True:  # so that Keras can loop through this as long as it wants
		for batch in batch_gen:
	        # prepare batch images
			batch_imgs = []
			batch_targets = []
			for raw, label in batch:
				img = preprocess_func(raw)
				batch_imgs.append(img)
				batch_targets.append(label)
	        # stack images into 4D tensor [batch_size, img_size, img_size, 3]
			batch_imgs = np.stack(batch_imgs, axis=0)
	        # convert targets into 2D tensor [batch_size, num_classes]
			batch_targets = keras.utils.np_utils.to_categorical(batch_targets, num_classes)
			yield batch_imgs, batch_targets
