# -*- coding: utf-8 -*-

from __future__ import absolute_import
from __future__ import division

import numpy
import sklearn.neural_network
import sklearn.pipeline

import submissions
from data import *


def label_to_feature(y):
    feature = [0] * 10
    feature[y] = 1
    return feature

target = target.astype(numpy.uint8)
train = train.astype(numpy.float32)
test = test.astype(numpy.float32)

train /= 255
test /= 255

rbm_1 = sklearn.neural_network.BernoulliRBM(n_components=500, learning_rate=0.06, n_iter=20, verbose=1)
rbm_2 = sklearn.neural_network.BernoulliRBM(n_components=500, learning_rate=0.06, n_iter=20, verbose=1)
pl = sklearn.pipeline.Pipeline([("rbm_1", rbm_1), ("rbm_2", rbm_2)])
pl.fit(train, target)

new_features = []
for example, label in zip(train, target):
    train_trans = pl.transform(example)[0]
    new_features.append(numpy.concatenate((train_trans, label_to_feature(label))))

rbm_3 = sklearn.neural_network.BernoulliRBM(n_components=2000, learning_rate=0.06, n_iter=20, verbose=1)
rbm_3.fit(new_features, target)

test_trans = pl.transform(test)
test_trans = numpy.concatenate((test_trans, [[0] * 10] * len(test_trans)), axis=1)

rbm_aux = sklearn.neural_network.BernoulliRBM()
rbm_aux.intercept_hidden_ = rbm_3.intercept_visible_
rbm_aux.intercept_visible_ = rbm_3.intercept_hidden_
rbm_aux.components_ = numpy.transpose(rbm_3.components_)
result = rbm_aux.transform(rbm_3.transform(test_trans))[:, -10:]
pred = numpy.argmax(result, axis=1)

submissions.save_csv(pred, "deep_bernoulli_restricted_boltzmann_machine.csv")
