import optunity

import numpy as np

from lifelines.utils import concordance_index

def prepare_data(x, e, t):

    return x.astype("float32"), e.astype("int32"), t.astype("float32")

def sort4minibatches(xvals, evals, tvals, batchsize):

    ntot = len(xvals)

    indices = np.arange(ntot)

    np.random.shuffle(indices)

    start_idx = 0

    esall = []

    for end_idx in list(range(batchsize, batchsize*(ntot//batchsize)+1, batchsize))+[ntot]:

        excerpt = indices[start_idx:end_idx]

        sort_idx = np.argsort(tvals[excerpt])[::-1]

        es = excerpt[sort_idx]

        esall += list(es)

        start_idx = end_idx

    return xvals[esall], evals[esall], tvals[esall], esall

# 1. Hyperparameter search for Deep Learning model

def hypersearch_nn(x_data, y_data, method, nfolds, nevals, batch_size, num_epochs, backend: str,

                   model_kwargs: dict, **hypersearch):

    @optunity.cross_validated(x=x_data, y=y_data, num_folds=nfolds)

    def modelrun(x_train, y_train, x_test, y_test, **hypersearch):

        cv_log = train_nn(

            backend=backend, xtr=x_train, ytr=y_train, batch_size=batch_size, n_epochs=num_epochs,

            **model_kwargs, **hypersearch

        )

        cv_preds = cv_log.predict(x_test, batch_size=1)[1]

        cv_C = concordance_index(y_test[:, 1], -cv_preds, y_test[:, 0])

        return cv_C

    optimal_pars, searchlog, _ = optunity.maximize(

        modelrun, num_evals=nevals, solver_name=method, **hypersearch

    )

    print('Optimal hyperparameters : ' + str(optimal_pars))

    print('Cross-validated C after tuning: %1.3f' % searchlog.optimum)

    return optimal_pars, searchlog

def train_nn(backend: str, xtr, ytr, batch_size, n_epochs, model_name, lr_exp, alpha, weight_decay_exp, **model_kwargs):

    if backend == "tf":

        from survival4D.nn.tf import train_nn

    elif backend == "torch":

        from survival4D.nn.torch import train_nn

    else:

        raise ValueError("Backend {} not supported. Only tf or torch. ".format(backend))

    return train_nn(xtr, ytr, batch_size, n_epochs, model_name, lr_exp, alpha, weight_decay_exp, **model_kwargs)