text: # -*- coding: utf-8 -*-
  """Protein Atlas.ipynb
  
  Automatically generated by Colaboratory.
  
  Original file is located at
      https://colab.research.google.com/drive/1PO3vaqDxcXa04QJ289jKC6J4koRG9bAD
  
  ## Imports and constants
  """
  
  import tensorflow as tf
  import numpy as np
  import os
  import zipfile
  import pandas as pd
  
  
  def extract_datazip(zipfile_path=None, validation_start_index=None, train_percent=0.95, extract_path='data', save_filelist=True):
  	"""
  	Extract dataset train and test zip file present in 'zipfile_path' path to train, valid and test directory in 'extract_path' path.
  	Creating directory structure:
  		# TODO: find directory structure
  		
  	return:
  		list, list, list : list of files in train, valid and test directories
  		dict : dict of paths 
  	
  	"""
  	
  	assert os.path.isfile(os.path.join(zipfile_path, 'train.zip')), "{} file not found".format(os.path.join(zipfile_path,'train.zip'))
  	assert os.path.isfile(os.path.join(zipfile_path,'test.zip')), "{} file not found".format(os.path.join(zipfile_path,'test.zip'))
  	 
  	train_zfile = zipfile.ZipFile(os.path.join(zipfile_path,'train.zip'))
  	test_zfile = zipfile.ZipFile(os.path.join(zipfile_path,'test.zip'))
  	train_filelist = train_zfile.namelist()
  	train_filelist = pd.DataFrame(train_filelist, columns=['fileid'])
  	train_filelist = train_filelist.sort_values(by='fileid')
  	train_filelist = train_filelist.reset_index(drop=True)
  	
  	filelist = pd.DataFrame(train_filelist.fileid.str.split("_",1).tolist(),columns=["id","color"] )
  	
  	
  	if validation_start_index ==None:
  		validation_start_index = int(train_filelist.shape[0]//4*train_percent)*4
  	
  	if save_filelist:
  		save_list_validation_start_index = validation_start_index//4
  		if validation_start_index == None:
  			save_list_validation_start_index = int(save_list.shape[0]*train_percent)
  		
  		save_list = pd.DataFrame(train_filelist.fileid.str.split('_', 1).tolist(), columns=['fileid','color']).drop(columns='color').drop_duplicates()		
  		save_list[ : save_list_validation_start_index].to_csv(os.path.join(extract_path,'train/train_filelist.csv'), sep=',', index=False)
  		save_list[save_list_validation_start_index : ].to_csv(os.path.join(extract_path,'valid/valid_filelist.csv'), sep=',', index=False)
  		test_savelist = pd.DataFrame(test_zfile.namelist(), columns=["fileid"]).sort_values(by="fileid").reset_index(drop=True)
  		test_savelist = pd.DataFrame(test_savelist.fileid.str.split("_",1).tolist(),columns=["id","color"] ).drop(columns='color').drop_duplicates()
  		test_savelist.to_csv(os.path.join(extract_path,'test/test_filelist.csv'), sep=',', index=False)
  	
  	valid_filelist = np.squeeze(train_filelist[validation_start_index : ].values)
  	train_filelist = np.squeeze(train_filelist[ : validation_start_index].values)
  
  
  	print('Extracting train images at {}'.format(extract_path), end='\t')
  	train_zfile.extractall(os.path.join(extract_path, 'train'), train_filelist)
  	print('done')
  	print('Extracting validation images at {}'.format(extract_path), end='\t')
  	train_zfile.extractall(os.path.join(extract_path, 'valid'), valid_filelist)
  	print('done')
  	print('Extracting test images at {}'.format(extract_path), end='\t')
  	test_zfile.extractall(os.path.join(extract_path, 'test'))
  	print('done')
  	extract_path = os.path.dirname(os.path.abspath(extract_path))
  	return {'train': os.path.join(extract_path,'train'), 'valid': os.path.join(extract_path,'valid'), 'test': os.path.join(extract_path,'test')}
  
  def preprocess_label(label_file=None, save_file=None):
  	"""
  	Converts label csv file into onehot format, requires pandas dataframe
  	"""
  	
  	new_lbl = pd.DataFrame(data=None,	columns=['Id']+ [i for i in range(28)])
  	tr_list = pd.read_csv(label_file, sep=',')
  	for index, row in tr_list.iterrows():
  		n_hot = np.array([0]*28)
  		n_hot[list(map(int,row['Target'].split()))] = 1
  		new_lbl.loc[index]= [row["Id"]]+ n_hot.tolist()
  	if save_file:
  		new_lbl.to_csv(save_file, index=False)
  	return new_lbl
  
  def split_label(label, validation_start_index=None, train_percent=0.95, save_listfile=True, save_path='data'):
  	"""
  	Split label file into train and validation 
  	
  	"""
  	lbl = pd.read_csv(label)
  	if validation_start_index == None:
  		validation_start_index = int(lbl.shape[0] * train_percent)
  	if save_listfile:
  		lbl[ : validation_start_index].to_csv(os.path.join(save_path, 'train/train_label.csv'), index=False)
  		lbl[validation_start_index : ].to_csv(os.path.join(save_path, 'valid/valid_label.csv'), index=False)
  	return np.squeeze(lbl[ : validation_start_index].values), np.squeeze(lbl[validation_start_index : ].values)
  
  
  def get_batch(img_file_path='data', batch_type=None, lbl_dir_path ='./', image_dim=[512, 512, 1], batch_size=32):
  	'''
  	Get iterator of dataset
  	
  	Parameters:
  	img_file_path - directory path containing train, valid and test image directories
  	batch_type - either 'train', 'valid' or 'test'
  	lbl_dir_path - directory path containing train, valid and test label csv file
  	image_dim - Three dimensions of image [width, height, channel]
  	batch_size - batch size of dataset
  	
  	Returns:
  	train and validation iterator if batch type is train
  	else test iterator
  	'''
  	
  	def decode_img(file_path, filename, image_dim):
  		r = tf.reshape(tf.image.decode_png(tf.read_file(file_path+'/'+filename + '_red.png'), channels=1), image_dim)
  		g = tf.reshape(tf.image.decode_png(tf.read_file(file_path+'/'+filename + '_green.png'), channels=1), image_dim)
  		b = tf.reshape(tf.image.decode_png(tf.read_file(file_path+'/'+filename + '_blue.png'), channels=1), image_dim)
  		y = tf.reshape(tf.image.decode_png(tf.read_file(file_path+'/'+filename + '_yellow.png'), channels=1), image_dim)
  		img = tf.concat([r,g,b,y], axis=2)
  		return img/ tf.reduce_max(img)
  	
  	assert batch_type in ['train', 'valid', 'test'], 'type should be train or test'
  	if batch_type in ['train','valid']:
  		train_img_path=os.path.join(img_file_path, batch_type)
  		train_lbl_path=os.path.join(lbl_dir_path, '')
  		tr_lbl_mat = pd.read_csv(train_lbl_path+'/{}_label.csv'.format(batch_type), sep=',').drop('Id', axis=1).values		
  		tr_img = tf.data.TextLineDataset(filenames=[train_img_path+'/{}_filelist.csv'.format(batch_type)]).skip(1)
  		tr_lbl = tf.data.Dataset.from_tensor_slices(tr_lbl_mat)
  		tr_dataset = tf.data.Dataset.zip((tr_img,tr_lbl))
  		tr_dataset = tr_dataset.shuffle(buffer_size=50000).map(lambda x,y: (x,decode_img(train_img_path, x, image_dim), y))
  		tr_dataset = tr_dataset.batch(batch_size).prefetch(batch_size*2)
  		return tr_dataset.make_initializable_iterator()
  	else:
  		test_img_path=os.path.join(img_file_path, 'test')
  		te_dataset = tf.data.TextLineDataset(filenames=[test_img_path+'/test_filelist.csv']).skip(1)
  		te_dataset = te_dataset.map(lambda x: [x,decode_img(test_img_path, x, image_dim)])
  		te_dataset = te_dataset.batch(batch_size).prefetch(batch_size*2)
  		return te_dataset.make_initializable_iterator()
  
  
  def get_model(X=None, Y=None, output_units=28, mode='Train', sigmoid_threshold=0.5):
  	"""
  	returns tensorflow model network 
  	
  	Args:
  	   X (tensor): input to the network
         Y (tensor): target/labels for the model
  	   output_units (int): class size of labels
         mode (str): one of ['train', 'valid', 'test']
         sigmoid_threshold (float): Sigmoid threshold for generating class targets from logits
         
      Returns:
          Dict : if mode is test then dictionary of logits and class targets generated from logits
                 else a dictionary of loss, accuracy, tf accuracy update function, and it's initializer
  	"""
  	
  	
  	sig_cond = tf.Variable([sigmoid_threshold]*output_units)
  	
  	# TODO: make new model
  	layer = tf.layers.conv2d(inputs=X, filters=16, kernel_size=[3,3], kernel_initializer= tf.initializers.random_uniform(),	)
  	layer = tf.layers.max_pooling2d(inputs = layer, pool_size=[2,2], strides=2)
  	layer = tf.nn.relu(layer)
  	layer = tf.layers.conv2d(inputs=layer, filters=64, kernel_size=[3,3], kernel_initializer= tf.initializers.random_uniform(), )
  	layer = tf.layers.max_pooling2d(inputs = layer, pool_size=[2,2], strides=2)
  	layer = tf.nn.relu(layer)
  	layer = tf.layers.flatten(inputs = layer,)
  	layer = tf.layers.dense(inputs = layer, units=256, kernel_initializer=tf.initializers.random_uniform(), )
  	layer = tf.nn.relu(layer)
  	layer = tf.layers.dense(inputs= layer, units = 128, kernel_initializer=tf.initializers.random_uniform(), )
  	layer = tf.nn.relu(layer)
  	layer =	tf.layers.dense(inputs= layer, units = output_units, kernel_initializer=tf.initializers.random_uniform(), )
  	
  	logit = layer
  	Y_ = tf.nn.sigmoid(logit,)
  	Y_ = tf.cast(x = tf.greater(tf.cast(Y_, tf.float32), sig_cond), dtype=tf.int32,	name = 'predictions')
  	
  	label = tf.stop_gradient(Y)
  	
  	if mode in ['train', 'valid']:
  		loss = tf.reduce_mean( tf.nn.sigmoid_cross_entropy_with_logits(labels=tf.cast(label, tf.float32), logits=logit))
  		accuracy, update_op = tf.metrics.accuracy( labels=label, predictions=Y_, name='accuracy')
  		accuracy_op_init = tf.variables_initializer(tf.get_default_graph().get_collection('local_variables',), name='metrics_initializer')
  		# Model Summaries
  		tf.summary.scalar('Loss', loss)
  		tf.summary.scalar('Accuracy', accuracy)		
  		return {'loss':loss,'accuracy':accuracy,'acc_update_op':update_op, 'acc_initializer':accuracy_op_init }
  	else:
  		return {'logits':logit, 'predictions':Y_}
  
  
  def get_graph(graph_type='train', summaries_dir=None, **kwargs):
  	"""
      Creates an instance of tf.graph and session for train/test purpose.
      
      If graph type is train then argument learning rate lr is also required.
      
      Args:
          graph_type (str): One of [train, valid, test]
          summaries_dir (str) : path of directory to save tf summaries
      Conditional Args : If graph_type is train (kwargs):
          lr (float) : Initial learning rate
          decay_lr (Boolean): If true use decaying learning rate, requiring kwargs:
              decay_rate (float): Decay rate of lr
              decay_step (int): Decay the lr every decay_step steps
      
      Returns:
          Dict : Dictionary containing a graph, session, optimizer, tf merge_all function, summary_writer, saver, max_to_keep number of checkpoint files  
      
      """
  
      # argument validation
  	assert graph_type.lower() in ['train', 'valid', 'test'], 'Invalid graph type'
  	assert summaries_dir !=None, 'summary directory not defined'
  	if graph_type=='train':
  		assert 'lr' in kwargs.keys(), 'training graph must have learing rate' 
  	
  	graph = tf.Graph()	
  	with graph.as_default():
  		with tf.variable_scope('{}_graph'.format(graph_type)):
  
  			if graph_type == 'train':
  				#optimiser
  				global_step = tf.Variable(0, name='global_step')
  				#decaying learning rate
  				if 'decay_lr' in kwargs.keys() and kwargs['decay_lr']:
  					lr = tf.train.exponential_decay(kwargs['lr'], global_step=global_step, decay_rate=kwargs['decay_rate'], decay_steps=kwargs['decay_steps'], name='lr_decay' )
  					tf.summary.scalar('Learning Rate', lr)
  				else: 
  					lr = kwargs['lr']
  				optimizer = tf.train.AdamOptimizer( learning_rate=lr, )
  			else:
  				global_step = None
  				optimizer = None
  
  			summaries = tf.summary.merge_all
  			summary_writer = tf.summary.FileWriter( summaries_dir, graph )
  			config_proto = tf.ConfigProto(allow_soft_placement=True)
  			session = tf.Session( config= config_proto )
  			if graph_type.lower()=='valid':
  				max_to_keep=1
  			else:
  				max_to_keep=5	
  			return {'graph':graph, 
  							'session':session,
  							'global_step': global_step,
  							'optimizer':optimizer,
  							'summaries':summaries,
  							'summary_writer':summary_writer,
  							'saver':tf.train.Saver,
  							'max_to_keep':max_to_keep
  							 }
  
  def train(epochs=1000, **kwargs ):
  	"""
  	Train the model
  
  	Args:
  		epochs (int): Number of training loop
  		image_dir (str): Path to directory containing train, validation and test image directories
  		label_dir (str): Path containing train and validation labels 
  		batch_size (int): batch size used in training
  		lr (float) : Initial learning rate
  		label_size : class size of label (defaults to 28)
  		checkpoint_dir : directory to save model and summaries, 
  		break_patience : Early stopping patience value 
  		decay_lr (Boolean): If true use decaying learning rate, requiring kwargs:
  		decay_rate (float): Decay rate of lr
  		decay_step (int): Decay the lr every decay_step steps
  	"""
  	with tf.device("/device:GPU:0"):
  		with tf.name_scope('training_loop'):
  			# build graph and insert model, iterator
  			tr_summaries_dir = os.path.join(kwargs['checkpoint_dir'], 'train')
  			va_summaries_dir = os.path.join(kwargs['checkpoint_dir'], 'validation')
  			tr_checkpoint_prefix = os.path.join( kwargs['checkpoint_dir'], 'model.ckpt')
  			best_checkpoint_prefix = os.path.join( kwargs['checkpoint_dir'], 'best_model/model.ckpt')
  
  			train_graph = get_graph( graph_type='train', summaries_dir=tr_summaries_dir, **kwargs ) 
  			valid_graph = get_graph( graph_type='valid', summaries_dir=va_summaries_dir, **kwargs )
  			with train_graph['graph'].as_default():
  				with tf.variable_scope('train_graph'):
  					tr_iter = get_batch( img_file_path=kwargs['image_dir'], batch_type='train', lbl_dir_path=kwargs['label_dir'], batch_size=kwargs['batch_size'] )
  					tr_img_ids, tr_x, tr_y = tr_iter.get_next()
  					train_model = get_model(X=tr_x, Y=tr_y, mode='train', output_units=kwargs['label_size'])
  					train_op = train_graph['optimizer'].minimize(train_model['loss'], global_step=train_graph['global_step'])
  					# merge all summaries to be saved during training
  					tr_merge_all_summaries = train_graph['summaries']
  					tr_all_summaries = tr_merge_all_summaries()
  					train_graph['session'].run(tf.global_variables_initializer())
  					train_graph['saver'] = train_graph['saver'](max_to_keep = train_graph['max_to_keep'])
  					# save initial graph
  					train_graph['saver'].save(save_path=tr_checkpoint_prefix, sess=train_graph['session'], global_step=train_graph['global_step'])
  
  			with valid_graph['graph'].as_default():
  				with tf.variable_scope('valid_graph'):
  					# naive early stopping parameters
  					prev_loss = best_loss = 100000000000.0
  					cur_patience = 0
  					va_iter = get_batch( img_file_path=kwargs['image_dir'], batch_type='valid', lbl_dir_path=kwargs['label_dir'], batch_size=kwargs['batch_size'] )
  					va_img_ids, va_x, va_y = va_iter.get_next()
  					valid_model = get_model(X=va_x, Y=va_y, mode='valid', output_units=kwargs['label_size'])
  					va_merge_all_summaries = valid_graph['summaries']
  					va_all_summaries = va_merge_all_summaries()
  					valid_graph['saver'] = valid_graph['saver'](max_to_keep = valid_graph['max_to_keep'])
  				
  			for epoc in range(epochs):
  				# one complete pass of training data
  				print("epoc: {}".format(epoc) )
  				
  				with train_graph['graph'].as_default():
  					with tf.variable_scope('train_graph'):
  						# reinitialize train_batch iterator
  						train_graph['session'].run(tr_iter.initializer)
  						train_graph['session'].run(train_model['acc_initializer'])
  						while True:
  							try:
  								#train model
  								loss, _, accuracy, _ , tr_summaries = train_graph['session'].run([train_model['loss'], train_model['acc_update_op'], train_model['accuracy'], train_op, tr_all_summaries])
  								train_graph['summary_writer'].add_summary(tr_summaries, train_graph['global_step'].eval())
  								train_graph['saver'].save(save_path=tr_checkpoint_prefix, sess=train_graph['session'], global_step=train_graph['global_step'].eval())
  							except tf.errors.OutOfRangeError:
  								break
  								
  				with valid_graph['graph'].as_default():
  					with tf.variable_scope('valid_graph'):
  						valid_graph['session'].run(va_iter.initializer)
  						valid_graph['session'].run(valid_graph['acc_initializer'])
  						valid_graph['saver'].restore(sess=valid_graph['session'], save_path= tf.train.latest_checkpoint(kwargs['checkpoint_dir']) )
  						while True:
  							try:
  								#validate model
  								val_loss, _, val_accuracy, val_summaries = valid_graph['session'].run([valid_model['loss'], valid_model['acc_update_op'], valid_model['accuracy'], va_all_summaries])
  							except tf.errors.OutOfRangeError:
  								valid_graph['summary_writer'].add_summary(val_summaries, epoc)
  								print("global_step: {:6}, loss: {:13.6f}, accuracy ={:.6f}, val_loss: {:13.6f}, val_accuracy: {:.6f}".format(valid_graph['global_step'].eval(), loss, accuracy, val_loss, val_accuracy))
  								break
  						# save checkpoint of loss be better than previous one
  						if float("{:.2f}".format(val_loss)) < float("{:.2f}".format(best_loss)):
  							best_loss = val_loss
  							cur_patience = 0
  							valid_graph['saver'].save(save_path=best_checkpoint_prefix, sess=valid_graph['session'], global_step=valid_graph['global_step'].eval())
  						else:
  							cur_patience += 1
  
  					if cur_patience == break_patience:
  						print('\n############ Early stopping ############')
  						train_graph['session'].close()
  						valid_graph['session'].close()
  						break
  
  
  def make_predictions(**kwargs):
  	# testing step
  	
  	with tf.device("/device:GPU:0"):
  			with tf.name_scope('test_loop'):
  				with graph.as_default():
  					with tf.variable_scope('test_graph'):    
  						te_iter = get_batch( img_file_path=kwargs['image_dir'], batch_type='test', lbl_dir_path=kwargs['label_dir'], batch_size=kwargs['test_batch_size'] )
  						te_img_ids, te_x, te_y = te_iter.get_next()
  						
  						graph.run(te_iter.initializer)
  						while True:
  							try:
  								#validate model
  								predictions = graph.run(Y_, {X:tr_x})
  								all_predictions.append( zip(te_img_ids,predictions) )
  							except tf.errors.OutOfRangeError:
  								with open(os.path.join(checkpoint_dir, 'predictions.csv')) as fp:
  									fp.write(all_predictions)
  								break
  
  
  
  