import matplotlib.pyplot as plt
import numpy as np
from tensorflow.keras import optimizers

from utils import data, evaluate, plot

from acrbm.models import DeepACRBM

(x_train, y_train), (x_test, y_test) = data.mnist(flatten=True)


def acrbm_compile(acrbm_layer):

    optimizer = optimizers.Adam(learning_rate=1e-2, beta_1=0.01, beta_2=0.999)

    metrics = {'h': evaluate.ClusteringAccuracy()}

    acrbm_layer.compile(metrics=metrics, optimizer=optimizer)


deep_acrbm = DeepACRBM()

deep_acrbm.stack(input_shape=(784, ), hidden_dim=500, visible=False, sigma_h='softmax',
                 compile_callback=acrbm_compile)

deep_acrbm.stack(input_shape=(500, ), hidden_dim=100, visible=False, sigma_h='softmax',
                 compile_callback=acrbm_compile)

deep_acrbm.stack(input_shape=(100, ), hidden_dim=100, visible=False, sigma_h='softmax',
                 compile_callback=acrbm_compile)

deep_acrbm.stack(input_shape=(100, ), hidden_dim=10, visible=False, sigma_h='softmax',
                 compile_callback=acrbm_compile)

deep_acrbm.compile()

deep_acrbm.summary()

# deep_acrbm.load_weights('deep_acrbm.h5')

_history = deep_acrbm.fit(x_train, y_train, batch_size=256, epochs=30, shuffle=False, verbose=1).history

# deep_acrbm.save('deep_acrbm.h5')

outputs = deep_acrbm.predict(x_test)

score = evaluate.clustering_accuracy(y_test, outputs.argmax(axis=-1))
confidence = outputs.max(axis=-1).mean()

print('n: {}, accuracy : {}, confidence: {}'.format(outputs.shape[-1], score, confidence))

generated = deep_acrbm.predict_visible(x_test, batch_size=32).numpy()

plt.style.use('seaborn-dark')

sample_size = 4

sample = np.concatenate([x_test[:sample_size], generated[:sample_size]], axis=0)
sample = sample.reshape((-1, 28, 28))

fig, axs = plot.show_images(sample, nrows=2, ncols=sample_size)
plt.show()
