import os  
import sys
import cPickle  
import pickle  
import numpy as np  
from numpy import array, append  
from PIL import Image  
#import Image  
#Lecture des batchs
def read_batch(batch_path, data_size):  
    in_file = open(batch_path, 'r+')  
    xx = cPickle.load(in_file)  
    in_file.close()  
    T_datas = xx['data']  
    datas = T_datas.T  
    c = np.zeros((1, data_size*data_size), dtype=np.float64)  
    i  = 0  
    for data in datas:  
        i += 1  
        c = c + data  
    return i, c 

def add_all(data_size, path):  
    count = 0 
    idx = 1 
    totalc = np.zeros((1, data_size*data_size), dtype=np.float64)  
    for elements in os.listdir(path):  
    	if elements.count('data_batch') == 1:
		print 'reading '+elements
		path1 = ''  
		path1 = path+elements 
		curcount, curc = read_batch(path1, data_size)  
		print curcount
		print curc
		count += curcount  
		totalc = totalc + curc  
		idx += 1
  
    return count, totalc  

def write_data(data_size, path):  
    cout, total = add_all(data_size, path)  
    a  = []  
    for i in range(0, len(total[0])):  
        c = total[0][i] / cout  
        a.append( [c])  
    a_array = array(a, dtype = np.float64)  
    return a_array  



def main():
        input_dir = sys.argv[1]		#Dossier des batchs
	data_size = int ( sys.argv[2] )	#Taille des images (ex: 128)
	batch_size = int ( sys.argv[3] )#Taille des batchs

	data_mean = write_data(data_size, input_dir)
	label_names = ['cover', 'stego']   
	dic = { 'num_cases_per_batch':batch_size, 'label_names':label_names, 'num_vis':data_size*data_size, 'data_mean':data_mean}  
	out_file = open(input_dir+'batches.meta', 'w+')  
	cPickle.dump(dic, out_file, 1)  
	out_file.close() 
	

	

if __name__ == '__main__':
	main()