from lightgbm.callback import _format_eval_result, CallbackEnv
from optuna import Trial


class LogEvaluationCallback(object):
    def __init__(self, logger, period=1, show_stdv=True):
        self.logger = logger
        self.period = period
        self.show_stdv = show_stdv

    def __call__(self, env: CallbackEnv, *args, **kwargs):
        """
        print or logging out current learning status

        Args:
            env:
            *args:
            **kwargs:

        Returns:
            None
        """
        if self.period <= 0:
            return

        if not env.evaluation_result_list:
            return

        if (env.iteration + 1) % self.period != 0:
            return

        try:
            result = '\t'.join([_format_eval_result(x, self.show_stdv) for x in env.evaluation_result_list])
        except ValueError:
            scores = [f'{key} {score}' for key, score in env.evaluation_result_list]
            result = '\t'.join(scores)
        self.logger.info('[%d]\t%s' % (env.iteration + 1, result))


def get_boosting_parameter_suggestions(trial: Trial) -> dict:
    """
    Get parameter sample for Boosting (like XGBoost, LightGBM)

    Args:
        trial(trial.Trial):

    Returns:
        dict: parameter sample generated by trial obj
    """
    return {
        # L2 正則化
        'reg_lambda': trial.suggest_loguniform('reg_lambda', 1e-3, 1e3),
        # L1 正則化
        'reg_alpha': trial.suggest_loguniform('reg_alpha', 1e-3, 1e3),
        # 弱学習木ごとに使う特徴量の割合
        # 0.5 だと全体のうち半分の特徴量を最初に選んで, その範囲内で木を成長させる
        'colsample_bytree': trial.suggest_loguniform('colsample_bytree', .5, 1.),
        # 学習データ全体のうち使用する割合
        # colsample とは反対に row 方向にサンプルする
        'subsample': trial.suggest_loguniform('subsample', .5, 1.),
        # 木の最大の深さ
        # たとえば 5 の時各弱学習木のぶん機は最大でも5に制限される.
        'max_depth': trial.suggest_int('max_depth', low=3, high=8),
        # 末端ノードに含まれる最小のサンプル数
        # これを下回るような分割は作れなくなるため, 大きく設定するとより全体の傾向でしか分割ができなくなる
        # [NOTE]: 数であるのでデータセットの大きさ依存であることに注意
        'min_child_weight': trial.suggest_uniform('min_child_weight', low=.5, high=40)
    }
