# Path: pynet/model.py
# class AutoEncoder(MLP):
# 
#     def __init__(self, **kwargs):
# 
#         self.encode_layers = []
#         self.decode_layers = []
# 
#         super(AutoEncoder, self).__init__(**kwargs)
# 
# 
#     def add_encode_layer(self, layer):
#         self.add_layer(layer)
#         self.encode_layers.append(layer)
# 
#     def add_decode_layer(self, layer):
#         self.add_layer(layer)
#         self.decode_layers.append(layer)
# 
#     def pop_encode_layer(self, index=0):
#         layer = self.encode_layers.pop(index)
#         self.layers.remove(layer)
#         return layer
# 
#     def pop_decode_layer(self, index=0):
#         layer = self.decode_layers.pop(index)
#         self.layers.remove(layer)
#         return layer
# 
#     def _fprop(self, layers, input_state):
#         for layer in layers:
#             input_state = layer._test_fprop(input_state)
#         return input_state
# 
#     def encode(self, input_values):
#         input_state = T.matrix('X', dtype=floatX)
#         output_state = self._fprop(self.encode_layers, input_state)
#         f = theano.function([input_state], output_state)
#         return f(input_values)
# 
#     def decode(self, input_values):
#         input_state = T.matrix('X', dtype=floatX)
#         output_state = self._fprop(self.decode_layers, input_state)
#         f = theano.function([input_state], output_state)
#         return f(input_values)


import argparse
import cPickle
import theano.tensor as T
import numpy as np
import copy
import theano
import pynet.layer as layers
from theano.sandbox.cuda.var import CudaNdarraySharedVariable
from pynet.model import AutoEncoder


floatX = theano.config.floatX

parser = argparse.ArgumentParser(description='''Convert gpu pickle pynet model to cpu pickle pynet model''')
parser.add_argument('--gpu_model', metavar='Path', required=True, help='the path to the gpu model pickle file')
parser.add_argument('--cpu_model', metavar='Path', required=True, help='''path to save the cpu model pickle file''')
args = parser.parse_args()

print ('loading gpu autoencoder..')
fin = open(args.gpu_model)
gpu_model = cPickle.load(fin)

ae = AutoEncoder(input_dim=gpu_model.input_dim)
for layer in gpu_model.encode_layers:
    layerW = T._shared(np.array(layer.W.get_value(), floatX),
                        name=layer.W.name, borrow=False)
    layerb = T._shared(np.array(layer.b.get_value(), floatX),
                        name=layer.b.name, borrow=False)
    encode_layer = getattr(layers, layer.__class__.__name__)(dim=layer.dim, name=layer.name,
                                                            W=layerW, b=layerb)
    ae.add_encode_layer(encode_layer)
    print 'encode layer', encode_layer.name, encode_layer.dim
print 'encode layers', ae.encode_layers

for ae_layer, gpu_layer in zip(reversed(ae.encode_layers), gpu_model.decode_layers):
    gpu_decode_layer_b = T._shared(np.array(gpu_layer.b.get_value(), floatX),
                        name=gpu_layer.b.name, borrow=False)
    decode_layer = getattr(layers, gpu_layer.__class__.__name__)(name=gpu_layer.name, dim=gpu_layer.dim,
                                                                W=ae_layer.W.T, b=gpu_decode_layer_b)

# Path: pynet/model.py
# class AutoEncoder(MLP):
# 
#     def __init__(self, **kwargs):
# 
#         self.encode_layers = []
#         self.decode_layers = []
# 
#         super(AutoEncoder, self).__init__(**kwargs)
# 
# 
#     def add_encode_layer(self, layer):
#         self.add_layer(layer)
#         self.encode_layers.append(layer)
# 
#     def add_decode_layer(self, layer):
#         self.add_layer(layer)
#         self.decode_layers.append(layer)
# 
#     def pop_encode_layer(self, index=0):
#         layer = self.encode_layers.pop(index)
#         self.layers.remove(layer)
#         return layer
# 
#     def pop_decode_layer(self, index=0):
#         layer = self.decode_layers.pop(index)
#         self.layers.remove(layer)
#         return layer
# 
#     def _fprop(self, layers, input_state):
#         for layer in layers:
#             input_state = layer._test_fprop(input_state)
#         return input_state
# 
#     def encode(self, input_values):
#         input_state = T.matrix('X', dtype=floatX)
#         output_state = self._fprop(self.encode_layers, input_state)
#         f = theano.function([input_state], output_state)
#         return f(input_values)
# 
#     def decode(self, input_values):
#         input_state = T.matrix('X', dtype=floatX)
#         output_state = self._fprop(self.decode_layers, input_state)
#         f = theano.function([input_state], output_state)
#         return f(input_values)


# Path: scripts/gpu_ae_to_cpu_ae.py
import argparse
import cPickle
import theano.tensor as T
import numpy as np
import copy
import theano
import pynet.layer as layers
from theano.sandbox.cuda.var import CudaNdarraySharedVariable
from pynet.model import AutoEncoder



floatX = theano.config.floatX

parser = argparse.ArgumentParser(description='''Convert gpu pickle pynet model to cpu pickle pynet model''')
parser.add_argument('--gpu_model', metavar='Path', required=True, help='the path to the gpu model pickle file')
parser.add_argument('--cpu_model', metavar='Path', required=True, help='''path to save the cpu model pickle file''')
args = parser.parse_args()

print ('loading gpu autoencoder..')
fin = open(args.gpu_model)
gpu_model = cPickle.load(fin)

ae = AutoEncoder(input_dim=gpu_model.input_dim)
for layer in gpu_model.encode_layers:
    layerW = T._shared(np.array(layer.W.get_value(), floatX),
                        name=layer.W.name, borrow=False)
    layerb = T._shared(np.array(layer.b.get_value(), floatX),
                        name=layer.b.name, borrow=False)
    encode_layer = getattr(layers, layer.__class__.__name__)(dim=layer.dim, name=layer.name,
                                                            W=layerW, b=layerb)
    ae.add_encode_layer(encode_layer)
    print 'encode layer', encode_layer.name, encode_layer.dim
print 'encode layers', ae.encode_layers

for ae_layer, gpu_layer in zip(reversed(ae.encode_layers), gpu_model.decode_layers):
    gpu_decode_layer_b = T._shared(np.array(gpu_layer.b.get_value(), floatX),
                        name=gpu_layer.b.name, borrow=False)
    decode_layer = getattr(layers, gpu_layer.__class__.__name__)(name=gpu_layer.name, dim=gpu_layer.dim,
                                                                W=ae_layer.W.T, b=gpu_decode_layer_b)

    ae.add_decode_layer(decode_layer)