import numpy as np
import pandas as pd

from sklearn.metrics import accuracy_score
from keras_uncertainty.utils import numpy_negative_log_likelihood

import h5py

def ensemble_predictions(models, data):
    preds = []

    for model in models:
        preds.append(np.expand_dims(model.predict(data, verbose=0), axis=0))

    preds = np.concatenate(preds, axis=0)

    mean_pred = np.mean(preds, axis=0)
    mean_pred = mean_pred / np.sum(mean_pred, axis=1, keepdims=True)

    return mean_pred

def save_predictions(filename, data):
    output = h5py.File(filename)
    output.create_dataset("preds", data=data)
    output.close()

def evaluate_subensemble_layers(experiment_name, num_ensembles, x_train, y_train, x_test, y_test, ood_data, train_model_fn, train_submodel_fn, max_layers, subensemble_layers):
    results = pd.DataFrame(columns = ['num_ensembles', 'error', 'nll'])

    ensemble_models = []
    base_model = None

    for num_ens in num_ensembles:

        # First member of the ensemble, make base model
        if num_ens == 1:
            model, base_model = train_model_fn(x_train, y_train, max_layers - subensemble_layers, subensemble_layers)
        else:
            model = train_submodel_fn(base_model, x_train, y_train, max_layers - subensemble_layers, subensemble_layers)

        #Add one member to the ensemble    
        ensemble_models.append(model)
            
        #Evaluate ensemble
        mean_pred = ensemble_predictions(ensemble_models, x_test)
        class_true = np.argmax(y_test, axis=1)
        class_pred = np.argmax(mean_pred, axis=1)

        save_predictions("{}-num_ens-{}-preds.hdf5".format(experiment_name, num_ens), mean_pred)

        #Evaluate ensembe on OOD data
        ood_preds = ensemble_predictions(ensemble_models, ood_data)
        save_predictions("{}-num_ens-{}-ood-preds.hdf5".format(experiment_name, num_ens), ood_preds)

        acc = accuracy_score(class_true, class_pred)
        err = 100.0 * (1.0 - acc)
        nll = numpy_negative_log_likelihood(y_test, mean_pred)

        print("{} Ensembles, Error: {:.4f} NLL {:.4f}".format(num_ens, err, nll))

        results = results.append({'num_ensembles': num_ens, 'error': err, 'nll': nll}, ignore_index=True)

    return results