import gc

import numpy as np
import lightgbm as lgb
from lightgbm.callback import _format_eval_result
from psutil import cpu_count

from utils import fast_auc


__all__ = ['LightGBM']


def eval_auc(preds, dtrain):
    labels = dtrain.get_label()
    return 'auc', fast_auc(labels, preds), True


def log_evaluation(logger, period=1, show_stdv=True):
    def _callback(env):
        if (period > 0 and env.evaluation_result_list and
            (env.iteration + 1) % period == 0):

            result = '\t'.join([_format_eval_result(x, show_stdv)
                               for x in env.evaluation_result_list])
            logger.info(f'[{env.iteration+1}]\t{result}')

    _callback.order = 11
    return _callback


class LightGBM():
    def __init__(self):
        pass

    def train_and_validate(self, x_train, x_valid, y_train, y_valid, params, logger):
        train_data = lgb.Dataset(x_train, label=y_train)
        valid_data = lgb.Dataset(x_valid, label=y_valid)

        del x_train, y_train, y_valid
        gc.collect()

        params['model_params']['num_threads'] = int(cpu_count() / 2)

        self.model = lgb.train(
            params['model_params'],
            train_data,
            valid_sets=[train_data, valid_data],
            **params['train_params'],
            feval=eval_auc,
            callbacks=[log_evaluation(logger, period=20)],
        )

        del train_data, valid_data
        gc.collect()

        y_pred_valid = self.model.predict(
                            x_valid,
                            num_iteration=self.model.best_iteration
                        )
        return y_pred_valid

    def predict(self, x_test):
        y_pred_test = self.predict_chunk(x_test)
        return y_pred_test


    def predict_chunk(self, x_test):
        initial_idx = 0
        chunk_size = 1000000
        current_pred = np.zeros(len(x_test))

        while initial_idx < x_test.shape[0]:
            final_idx = min(initial_idx + chunk_size, x_test.shape[0])
            idx = range(initial_idx, final_idx)
            current_pred[idx] = self.model.predict(
                                    x_test.iloc[idx],
                                    num_iteration=self.model.best_iteration
                                )
            initial_idx = final_idx

        return current_pred

