import gc
import numpy as np
from functools import partial
from nn_utils.manager import NNManager
from NN.embedding_nn import embedded_mlp_multi
from nn_utils.neural import prep_for_embedding, df_to_list
from config import conf


def NN_cv(
    X_train,
    y_train,
    folds,
    logger,
    cv_path,
    X_test=None,
    optional_data=None,
    prep=True,
    split_conf=None
):
    preds = []

    meta = np.zeros_like(y_train).astype("float64")
    if split_conf is None:
        X_tr, X_te, main_conf, main_gen = prep_for_embedding(
            X_train, X_test, conf, prep=prep
        )
        X_train, X_test = X_tr, X_te
    else:
        main_conf = split_conf
        main_gen = partial(df_to_list, conf=split_conf)

    for num_fold, (tr_ind, tes_ind) in enumerate(folds):
        if num_fold > 0:
            break
        logger.info(f"fold_{num_fold}")

        X_cv_train, X_cv_test = X_train.iloc[tr_ind], X_train.iloc[tes_ind]
        y_cv_train, y_cv_test = y_train.iloc[tr_ind], y_train.iloc[tes_ind]

        fold_path = cv_path / f"fold{num_fold}/"
        seed_path = fold_path

        mlp = embedded_mlp_multi(main_conf, 10)
        # mlp = sparse_mlp(X_train.shape[1])
        nnm = NNManager(
            seed_path, mlp, 2048, X_train.columns, logger, proc_per_gen=main_gen,
            onehot_size=10
        )
        nnm.learn(X_cv_train, y_cv_train, X_cv_test, y_cv_test, epochs=25)
        gc.collect()

        if X_test is not None:
            pred = nnm.predict_generator(X_test)
            pred = pred[:, 0]
            np.save(f"{seed_path}pred.npy", pred)

        train_oof = nnm.predict_generator(X_cv_test)
        train_oof = train_oof[:, 0]
        np.save(f"{seed_path}train_oof.npy", train_oof)
        np.save(f"{fold_path}tes_ind.npy", tes_ind)
        meta[tes_ind] += train_oof
        del X_cv_train, y_cv_train, X_cv_test, y_cv_test

        if X_test is not None:
            preds.append(pred)

    return meta
