import sys
import os
from shutil import copyfile
GLOBAL_PATH='/home/casp13/DNCON2/scripts/';
sys.path.insert(0, GLOBAL_PATH)

from keras import backend as K
#from keras.utils.conv_utils import convert_kernel

from libcnnpredict import *

def convert_kernel(kernel):
    """Converts a Numpy kernel matrix from Theano format to TensorFlow format.
    Also works reciprocally, since the transformation is its own inverse.
    # Arguments
        kernel: Numpy array (3D, 4D or 5D).
    # Returns
        The converted kernel.
    # Raises
        ValueError: in case of invalid kernel shape or invalid data_format.
    """
    kernel = np.asarray(kernel)
    if not 3 <= kernel.ndim <= 5:
        raise ValueError('Invalid kernel shape:', kernel.shape)
    slices = [slice(None, None, -1) for _ in range(kernel.ndim)]
    no_flip = (slice(None, None), slice(None, None))
    slices[-2:] = no_flip
    return np.copy(kernel[slices])

if len(sys.argv) != 5:
          print 'please input the right parameters'
          sys.exit(1)


dir_config     = sys.argv[1]  # /home/casp13/DNCON2//model-config-n-weights
weights   = sys.argv[2]  # stage1-10A.hdf5
fileX          = sys.argv[3]
outputdir   = sys.argv[4]

file_weights  = dir_config + '/' + weights
file_weights_converted  = outputdir + '/' + weights
model_arch = read_model_arch(dir_config + '/model-arch.config')

print ''
print('SCRIPT        : ' + sys.argv[0]   )
print('dir_config    : ' + dir_config    )
print('file_weights  : ' + file_weights  )
print('file_weights_converted  : ' + file_weights_converted  )
print('fileX         : ' + fileX         )
print('outputdir  : ' + outputdir  )
print ''

# Need to make X slightly bigger than L x L by padding zeros
# Building a model with L x L decreases performance
L = 0
with open(fileX) as f:
	for line in f:
		if line.startswith('#'):
			continue
		L = line.strip().split()
		L = int(round(math.exp(float(L[0]))))
		break
LMAX = L + 10
x = getX(fileX, LMAX)
F = len(x[0, 0, :])
X = np.zeros((1, LMAX, LMAX, F))
X[0, :, :, :] = x

# Predict at (L+10) x (L+10) and trim it back to L x L
model = build_model_for_this_input_shape(model_arch, X)
model.load_weights(file_weights)

for layer in model.layers:
	if layer.__class__.__name__ in ['Convolution1D', 'Convolution2D']:
		original_w = K.get_value(layer.W)
		converted_w = convert_kernel(original_w)
		K.set_value(layer.W, converted_w)

model.save_weights(file_weights_converted)

print "model converted to ",file_weights_converted;