import matplotlib.pyplot as plt
import numpy as np

import tensorflow as tf

from tensorflow.keras import optimizers
from tensorflow.keras.layers import Input, Dense, Activation, Concatenate, GaussianNoise
from tensorflow.keras.models import Model

from utils import data, evaluate, plot, op

from acrbm.models import ShallowACRBM, AEBasedACRBM

tf.random.set_seed(2023)

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


def build_encoder(hidden_dim, latent_dim):

    # --------------------------------------

    inputs = Input(shape=(784, ))

    z = GaussianNoise(1.0)(inputs)

    for i in range(2):

        z = Dense(hidden_dim)(z)
        z = Activation('relu')(z)

    _latent = Dense(latent_dim, name='latent')(z)

    _model = Model(inputs, _latent)

    return _model


def build_acrbm(input_dim):

    # def _compile(acrbm_layer):
    #
    #     _optimizer = optimizers.SGD(learning_rate=1e-2)
    #
    #     acrbm_layer.compile(optimizer=_optimizer)

    shallow_acrbm = ShallowACRBM(input_shape=(input_dim, ), hidden_dim=10, visible=False)

    # _compile(shallow_acrbm)

    return shallow_acrbm


def merge(inputs):

    return Concatenate(name='merge')(inputs)


def build_decoder(input_dim, hidden_dim):

    inputs = Input(shape=(input_dim, ))

    z = inputs

    for i in range(2):

        z = Dense(hidden_dim)(z)
        z = Activation('relu')(z)

    _outputs = Dense(784)(z)
    _outputs = Activation('sigmoid', name='outputs')(_outputs)

    _model = Model(inputs, _outputs)

    return _model


def scheduler(epoch, lr):

    if epoch % 200 == 0:

        return lr

    else:

        return lr * 0.1


encoder = build_encoder(100, 90)
acrbm = build_acrbm(90)
decoder = build_decoder(100, 100)

autoencoder = AEBasedACRBM(input_shape=(784, ), encoder=encoder, acrbm=acrbm, merge=merge, decoder=decoder,
                           feedback=False)

optimizer = optimizers.Adam(learning_rate=1e-3, beta_1=0.05, beta_2=0.999)

autoencoder.compile(optimizer=optimizer, metrics=evaluate.ClusteringAccuracy())

autoencoder.summary()

reset_state = op.RestRandomState(random_state=2023, steps=50)
reset_opt = op.RestOptimizerState(identifier='adam', steps=50)
reset_metric = op.RestMetricState(steps=50)

autoencoder.load_weights('autoencoder.h5')

# # epochs=300+200
# _history = autoencoder.fit(x_train, y_train, batch_size=256, epochs=200, verbose=1,
#                            callbacks=[reset_state, reset_opt, reset_metric]).history

# autoencoder.save('autoencoder.h5')

generated, latent, h = autoencoder.predict(x_test)

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

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

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

sample_size = 4

indices = op.sample_indices(len(x_test), sample_size, 0)

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

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