# Path: morb/rbms.py
# class BinaryBinaryRBM(RBM): # the basic RBM, with binary visibles and binary hiddens
# class BinaryBinaryCRBM(BinaryBinaryRBM):
# class GaussianBinaryRBM(RBM): # Gaussian visible units
# class LearntPrecisionGaussianBinaryRBM(RBM):
# class LearntPrecisionSeparateGaussianBinaryRBM(RBM):
# class TruncExpBinaryRBM(RBM): # RBM with truncated exponential visibles and binary hiddens
#     def __init__(self, n_visible, n_hidden):
#     def _initial_W(self):
#     def _initial_bv(self):
#     def _initial_bh(self):
#     def __init__(self, n_visible, n_hidden, n_context):
#     def _initial_A(self):
#     def _initial_B(self):
#     def __init__(self, n_visible, n_hidden):
#     def _initial_W(self):
#     def _initial_bv(self):
#     def _initial_bh(self):
#     def __init__(self, n_visible, n_hidden):
#     def _initial_W(self):
#     def _initial_bias(self, n):
#     def __init__(self, n_visible, n_hidden_mean, n_hidden_precision):
#     def _initial_W(self, nv, nh):
#     def _initial_bias(self, n):
#     def __init__(self, n_visible, n_hidden):
#     def _initial_W(self):
#     def _initial_bv(self):
#     def _initial_bh(self):
# 
# Path: morb/stats.py
# def gibbs_step(rbm, vmap, units_list, mean_field_for_stats=[], mean_field_for_gibbs=[]):
# def cd_stats(rbm, v0_vmap, visible_units, hidden_units, context_units=[], k=1, mean_field_for_stats=[], mean_field_for_gibbs=[], persistent_vmap=None):
#     def gibbs_hvh(*args):
# 
# Path: morb/updaters.py
# class SelfUpdater(Updater):
# class MomentumUpdater(Updater):
# class CDUpdater(Updater):
# class SparsityUpdater(Updater):
# class BoundUpdater(Updater):
# class GradientUpdater(Updater):
#     def get_update(self):
#     def __init__(self, pu, momentum, variable_shape):
#     def get_update(self):
#     def get_theano_updates(self):
#     def __init__(self, rbm, variable, stats):
#     def get_update(self):
#     def __init__(self, rbm, variable, sparsity_targets, stats):
#     def get_update(self):        
#     def __init__(self, pu, bound=0, type='lower'):  
#     def get_update(self):
#     def get_theano_updates(self):
#     def __init__(self, objective, variable, theano_updates={}):
#     def get_update(self):
#     def get_theano_updates(self):
# 
# Path: morb/trainers.py
# class MinibatchTrainer(Trainer):
#     def compile_function(self, initial_vmap, monitors=[], name='func', mb_size=32, train=True, mode=None):
#         def func(dmap):
#         TF = theano.function([index], monitors,
#             updates = updates, givens = givens, name = name, mode = mode)    
# 
# Path: morb/monitors.py
# def reconstruction_mse(stats, u):
# def reconstruction_error_rate(stats, u):
# def reconstruction_crossentropy(stats, u):


import morb
import theano
import theano.tensor as T
import numpy as np
import gzip, cPickle
import matplotlib.pyplot as plt
from morb import rbms, stats, updaters, trainers, monitors
from utils import generate_data, get_context, plot_data
from theano import ProfileMode
mode = None


# load data
print ">> Loading dataset..."

f = gzip.open('datasets/mnist.pkl.gz','rb')
train_set, valid_set, test_set = cPickle.load(f)
f.close()

train_set_x, train_set_y = train_set
valid_set_x, valid_set_y = valid_set
test_set_x, test_set_y = test_set


# TODO DEBUG
train_set_x = train_set_x[:1000]
valid_set_x = valid_set_x[:100]


n_visible = train_set_x.shape[1]
n_hidden = 100


print ">> Constructing RBM..."
rbm = rbms.BinaryBinaryRBM(n_visible, n_hidden)
initial_vmap = { rbm.v: T.matrix('v') }

# try to calculate weight updates using CD-1 stats
print ">> Constructing contrastive divergence updaters..."

# Path: morb/rbms.py
# class BinaryBinaryRBM(RBM): # the basic RBM, with binary visibles and binary hiddens
# class BinaryBinaryCRBM(BinaryBinaryRBM):
# class GaussianBinaryRBM(RBM): # Gaussian visible units
# class LearntPrecisionGaussianBinaryRBM(RBM):
# class LearntPrecisionSeparateGaussianBinaryRBM(RBM):
# class TruncExpBinaryRBM(RBM): # RBM with truncated exponential visibles and binary hiddens
#     def __init__(self, n_visible, n_hidden):
#     def _initial_W(self):
#     def _initial_bv(self):
#     def _initial_bh(self):
#     def __init__(self, n_visible, n_hidden, n_context):
#     def _initial_A(self):
#     def _initial_B(self):
#     def __init__(self, n_visible, n_hidden):
#     def _initial_W(self):
#     def _initial_bv(self):
#     def _initial_bh(self):
#     def __init__(self, n_visible, n_hidden):
#     def _initial_W(self):
#     def _initial_bias(self, n):
#     def __init__(self, n_visible, n_hidden_mean, n_hidden_precision):
#     def _initial_W(self, nv, nh):
#     def _initial_bias(self, n):
#     def __init__(self, n_visible, n_hidden):
#     def _initial_W(self):
#     def _initial_bv(self):
#     def _initial_bh(self):
# 
# Path: morb/stats.py
# def gibbs_step(rbm, vmap, units_list, mean_field_for_stats=[], mean_field_for_gibbs=[]):
# def cd_stats(rbm, v0_vmap, visible_units, hidden_units, context_units=[], k=1, mean_field_for_stats=[], mean_field_for_gibbs=[], persistent_vmap=None):
#     def gibbs_hvh(*args):
# 
# Path: morb/updaters.py
# class SelfUpdater(Updater):
# class MomentumUpdater(Updater):
# class CDUpdater(Updater):
# class SparsityUpdater(Updater):
# class BoundUpdater(Updater):
# class GradientUpdater(Updater):
#     def get_update(self):
#     def __init__(self, pu, momentum, variable_shape):
#     def get_update(self):
#     def get_theano_updates(self):
#     def __init__(self, rbm, variable, stats):
#     def get_update(self):
#     def __init__(self, rbm, variable, sparsity_targets, stats):
#     def get_update(self):        
#     def __init__(self, pu, bound=0, type='lower'):  
#     def get_update(self):
#     def get_theano_updates(self):
#     def __init__(self, objective, variable, theano_updates={}):
#     def get_update(self):
#     def get_theano_updates(self):
# 
# Path: morb/trainers.py
# class MinibatchTrainer(Trainer):
#     def compile_function(self, initial_vmap, monitors=[], name='func', mb_size=32, train=True, mode=None):
#         def func(dmap):
#         TF = theano.function([index], monitors,
#             updates = updates, givens = givens, name = name, mode = mode)    
# 
# Path: morb/monitors.py
# def reconstruction_mse(stats, u):
# def reconstruction_error_rate(stats, u):
# def reconstruction_crossentropy(stats, u):


# Path: examples/example_sparsity.py
import morb
import theano
import theano.tensor as T
import numpy as np
import gzip, cPickle
import matplotlib.pyplot as plt
from morb import rbms, stats, updaters, trainers, monitors
from utils import generate_data, get_context, plot_data
from theano import ProfileMode

mode = None


# load data
print ">> Loading dataset..."

f = gzip.open('datasets/mnist.pkl.gz','rb')
train_set, valid_set, test_set = cPickle.load(f)
f.close()

train_set_x, train_set_y = train_set
valid_set_x, valid_set_y = valid_set
test_set_x, test_set_y = test_set


# TODO DEBUG
train_set_x = train_set_x[:1000]
valid_set_x = valid_set_x[:100]


n_visible = train_set_x.shape[1]
n_hidden = 100


print ">> Constructing RBM..."
rbm = rbms.BinaryBinaryRBM(n_visible, n_hidden)
initial_vmap = { rbm.v: T.matrix('v') }

# try to calculate weight updates using CD-1 stats
print ">> Constructing contrastive divergence updaters..."

s = stats.cd_stats(rbm, initial_vmap, visible_units=[rbm.v], hidden_units=[rbm.h], k=1, mean_field_for_stats=[rbm.v], mean_field_for_gibbs=[rbm.v])