import pandas as pd

from tqdm import tqdm

import tensorflow as tf
from tensorflow.keras import optimizers, backend

from utils import data, evaluate

from acrbm.functional import squared_l2_loss
from acrbm.models import ShallowACRBM


def run_on_mnist(n_clusters, visible, batch_size, epochs, lr=1e-3,
                 seed=None, summary=True, verbose=1, ret=None, **kwargs):

    tf.random.set_seed(seed=seed)

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

    model = ShallowACRBM(input_shape=(784, ), hidden_dim=n_clusters, visible=visible, **kwargs)

    optimizer = optimizers.Adam(learning_rate=lr, beta_1=0.01, beta_2=0.999)

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

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

    if summary:

        model.summary()

    _history = model.fit(x_train, y_train, batch_size=batch_size, epochs=epochs, shuffle=False, verbose=verbose).history

    outputs = model.predict(x_test)

    if not visible:

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

    else:

        score = evaluate.clustering_accuracy(y_test, outputs[0].argmax(axis=-1))

    if ret == 'score':

        return outputs[0].max(axis=-1).mean(), score

    elif ret == 'history':

        return _history

    elif ret == 'output':

        return outputs


def run_test_case(tid):

    if tid == 0:

        _, score = run_on_mnist(n_clusters=10, visible=False, epochs=20, batch_size=256, lr=1e-2,
                                use_bh=True, sigma_h='softmax', seed=3, ret='score')

    elif tid == 1:

        _, score = run_on_mnist(n_clusters=30, visible=False, epochs=50, batch_size=256, lr=1e-2,
                                use_bh=True, sigma_h='softmax', seed=1, ret='score')

    elif tid == 2:

        _, score = run_on_mnist(n_clusters=10, visible=True, epochs=20, batch_size=256, lr=1e-2,
                                reconstruction_loss=squared_l2_loss, use_bh=True, use_bv=True,
                                sigma_h='softmax', sigma_v='sigmoid', seed=0, ret='score')

    elif tid == 3:

        _, score = run_on_mnist(n_clusters=30, visible=True, epochs=20, batch_size=256, lr=1e-2,
                                reconstruction_loss=squared_l2_loss, use_bh=True, use_bv=True,
                                sigma_h='softmax', sigma_v='sigmoid', seed=0, ret='score')

    else:

        raise ValueError(f'Invalid test-case id, tid={tid}')

    print('clustering accuracy : {}'.format(score))


def confidence_benchmark(n_clusters):

    accuracy = []
    confidence = []

    niter = 5
    lr = 1e-2

    seed = 2019

    for nc in tqdm(n_clusters):

        c, score = run_on_mnist(n_clusters=nc, visible=False, epochs=niter, batch_size=256, lr=lr,
                                use_bh=True, sigma_h='softmax', summary=False, verbose=1, ret='score', seed=seed)

        print('n: {}, accuracy : {}, confidence: {}'.format(nc, score, c))

        confidence.append(c)
        accuracy.append(score)

        seed += 1

        backend.clear_session()

    benchmark = pd.DataFrame({'n_clusters': n_clusters, 'accuracy': accuracy, 'confidence': confidence})

    benchmark.to_csv('acrbm_confidence_rate.csv', index=False, encoding='utf-8')


run_test_case(0)
# confidence_benchmark(n_clusters=[10, 300, 600, 1000, 2000])
