import os
import sys
import time
import numpy
import theano
import theano.tensor as T
from theano.tensor.shared_randomstreams import RandomStreams
from pprint import pprint
from theano.tensor.signal import downsample
from theano.tensor.nnet import conv
from theano.tensor.nnet import conv3d2d
from scipy.io import matlab
import re
import math
from theano import shared
from collections import OrderedDict
from store_pytable import *
import random
from layers import *

def shared_dataset(data_x, data_y, borrow=True):

	shared_x = theano.shared(numpy.asarray(data_x,
										   dtype=theano.config.floatX),
							 borrow=borrow)
	shared_y = theano.shared(numpy.asarray(data_y,
										   dtype=theano.config.floatX),
							 borrow=borrow)

	return shared_x, T.cast(shared_y, 'int32')


def load_ATP_data():

	dataName = "data";
	labelName = "label";
	labelShape = [];

	pos_or_neg = 'pos'
	ID = 'test'+'_'+'pos'
	filename_train_pos = "../data/pytables_atp/"+ID+".pytables";
	h5file_train = tables.openFile(filename_train_pos, mode="r")
	dataColumn_train = getH5column(h5file_train, dataName);
	labelColumn_train = getH5column(h5file_train, labelName);
	X_test_pos=dataColumn_train[:]
	y_test_pos=labelColumn_train[:]
	
	pos_or_neg = 'neg'
	ID = 'test'+'_'+'neg'
	filename_train_neg = "../data/pytables_atp/"+ID+".pytables";
	h5file_train = tables.openFile(filename_train_neg, mode="r")
	dataColumn_train = getH5column(h5file_train, dataName);
	labelColumn_train = getH5column(h5file_train, labelName);
	X_test_neg=dataColumn_train[:]
	y_test_neg=labelColumn_train[:]

	num_of_pos_test = X_test_pos.shape[0]
	num_of_neg_test = X_test_neg.shape[0]

	Xt=numpy.concatenate((X_test_pos, X_test_neg), axis=0)
	yt=numpy.concatenate((y_test_pos, y_test_neg), axis=0)

	test_set_x, test_set_y = shared_dataset(Xt, yt)
	test_set_x=test_set_x.dimshuffle(0,4,1,2,3)
   
	return [test_set_x, test_set_y, Xt.shape[0]]


def load_weights_pickle(file_name,dropout_rate):
	keys=numpy.load(file_name).keys()
	print keys

	W0=numpy.load(file_name)[keys[0]]
	W1=numpy.load(file_name)[keys[1]]
	W2=numpy.load(file_name)[keys[2]]
	W3=numpy.load(file_name)[keys[3]]
	W4=numpy.load(file_name)[keys[4]]
	W5=numpy.load(file_name)[keys[5]]
	b0=numpy.load(file_name)[keys[6]]
	b1=numpy.load(file_name)[keys[7]]
	b2=numpy.load(file_name)[keys[8]]
	b3=numpy.load(file_name)[keys[9]]
	b4=numpy.load(file_name)[keys[10]]
	b5=numpy.load(file_name)[keys[11]]

	print W0.shape
	print W1.shape
	print W2.shape
	print W3.shape
	print W4.shape
	print W5.shape
	print b0.shape
	print b1.shape
	print b2.shape
	print b3.shape
	print b4.shape
	print b5.shape

	W0=theano.shared(value=W0*(1 - dropout_rate), name='W0', borrow=True)
	W1=theano.shared(value=W1*(1 - dropout_rate), name='W1', borrow=True)
	W2=theano.shared(value=W2*(1 - dropout_rate), name='W2', borrow=True)
	W3=theano.shared(value=W3*(1 - dropout_rate), name='W3', borrow=True)
	W4=theano.shared(value=W4*(1 - dropout_rate), name='W4', borrow=True)
	W5=theano.shared(value=W5, name='W5', borrow=True)

	b0=theano.shared(value=b0*(1 - dropout_rate), name='b0', borrow=True)
	b1=theano.shared(value=b1*(1 - dropout_rate), name='b1', borrow=True)
	b2=theano.shared(value=b2*(1 - dropout_rate), name='b2', borrow=True)
	b3=theano.shared(value=b3*(1 - dropout_rate), name='b3', borrow=True)
	b4=theano.shared(value=b4*(1 - dropout_rate), name='b4', borrow=True)
	b5=theano.shared(value=b5, name='b5', borrow=True)
	return [W0,W1,W2,W3,W4,W5,b0,b1,b2,b3,b4,b5]



def post_fine_S_CNN_dA(dropout_rate,best_epoch):

	[test_set_x, test_set_y, test_size] = load_ATP_data()
	batch_size = 100

	[W0,W1,W2,W3,W4,W5,b0,b1,b2,b3,b4,b5] = load_weights_pickle(file_name='../results/weights/weight_3DCNN_ATP_epoch_'+str(best_epoch)+'.zip',dropout_rate=dropout_rate)

	
	Weights = [W0,W1,W2,W3,W4,W5]
	Bias = [b0,b1,b2,b3,b4,b5]

	numpy_rng = numpy.random.RandomState(89677)
	print '... building the model'

	flt_time = 3
	in_channels = 4
	flt_height = 3
	flt_width = 3
	in_time = 20
	in_height = 20
	in_width = 20

	print (test_size/batch_size)
	for ind in range (0,test_size/batch_size):

		x = test_set_x[ind*batch_size:(ind+1)*batch_size]
		y = test_set_y[ind*batch_size:(ind+1)*batch_size]   

		rng = numpy.random.RandomState(23455)

		in_channels = 4
		filter_w = 3

		######################
		# BUILD ACTUAL MODEL #
		######################
		print '... building the model'
		# image sizes
		batchsize     = batch_size
		in_time       = 20
		in_width      = 20
		in_height     = 20
		#filter sizes
		flt_channels  = 100
		flt_time      = filter_w
		flt_width     = filter_w
		flt_height    = filter_w

		layer0_w = 20
		layer0_h = 20
		layer0_d = 20

		# ====== net1 =======

		layer1_w = (layer0_w-3+1) #14
		layer1_h = (layer0_h-3+1)
		layer1_d = (layer0_d-3+1)

		layer2_w = (layer1_w-3+1)/2 #14
		layer2_h = (layer1_h-3+1)/2
		layer2_d = (layer1_d-3+1)/2

		layer3_w = (layer2_w-3+1)/2
		layer3_h = (layer2_h-3+1)/2
		layer3_d = (layer2_d-3+1)/2

		signals_shape0 = (batchsize, in_time, in_channels, in_height, in_width)
		filters_shape0 = (flt_channels, 3, in_channels, 3, 3)
		signals_shape1 = (batchsize, layer1_d, flt_channels, layer1_h, layer1_w)
		filters_shape1 = (flt_channels*2, 3, flt_channels, 3, 3)
		signals_shape2 = (batchsize, layer2_d, flt_channels*2, layer2_h, layer2_w)
		filters_shape2 = (flt_channels*4, 3, flt_channels*2, 3, 3)

			
		layer0_input = x.reshape(signals_shape0)*(1 - dropout_rate) #20

		layer0 = Conv_3d_Layer(rng, input=layer0_input, #18
				image_shape=signals_shape0,
				filter_shape=filters_shape0, W=W0, b=b0)

		layer1 = Conv_3d_Layer(rng, input=layer0.output, #8
				image_shape=signals_shape1,
				filter_shape=filters_shape1, W=W1, b=b1)

		layer1_pool = PoolLayer3D(input=layer1.output.dimshuffle(0,2,1,3,4), pool_shape=(2,2,2)) #4

		layer2 = Conv_3d_Layer(rng, input=layer1_pool.output.dimshuffle(0,2,1,3,4), 
				image_shape=signals_shape2,
				filter_shape=filters_shape2, W=W2, b=b2)

		layer2_pool = PoolLayer3D(input=layer2.output.dimshuffle(0,2,1,3,4), pool_shape=(2,2,2)) #4

		layer3_input = layer2_pool.output.dimshuffle(0,2,1,3,4).flatten(2) 

		layer3 = HiddenLayer(rng, input=layer3_input, n_in=(flt_channels*4*layer3_d*layer3_w*layer3_h), 
							 n_out=1000, activation=relu, W=W3, b=b3)
		layer4 = HiddenLayer(rng, input=layer3.output, n_in=1000, 
							 n_out=100, activation=relu, W=W4, b=b4)

		layer5 = LogisticRegression(input=layer4.output, n_in=100, n_out=20, W=W5, b=b5) 

		pred = layer5.y_pred.eval()
		pos_prob = layer5.p_y_given_x[:,1].eval()
   

		if ind ==0:
			y_true=numpy.array(y.eval())
			y_pred=numpy.array(pred)
			y_prob=numpy.array(pos_prob)
		else:
			y_true=numpy.concatenate((y_true,numpy.array(y.eval())),axis=0)
			y_pred=numpy.concatenate((y_pred,numpy.array(pred)),axis=0)
			y_prob=numpy.concatenate((y_prob,numpy.array(pos_prob)),axis=0)


	return y_true, y_prob

	   
		


if __name__ == '__main__':
	# THEANO_FLAGS=floatX=float32 python post_train_eval_3DCNN_drop.py
	dropout_rate = 0.2
	best_epoch = 4
	
	[y_true, pos_prob] = post_fine_S_CNN_dA(dropout_rate,best_epoch)

	pos_prob.dump("../results/prob_scores/3DCNN/prob_y.dat")
	y_true.dump("../results/prob_scores/3DCNN/true_y.dat")