from .enums import Algo
from .schemas import ModelConfig
import optuna


def get_params(trial: optuna.Trial, model_config: ModelConfig) -> dict:
    """
    This method returns the required parameters for optuna.

    Args:
        trial (optuna.trial.Trial): Optuna trial
        model_config (ModelConfig): Model configuration

    Returns:
        dict: Parameters to optimize
    """
    if model_config.algo == Algo.xgb.value:
        params = {
            "learning_rate": trial.suggest_float("learning_rate", 1e-2, 0.25, log=True),
            "reg_lambda": trial.suggest_float("reg_lambda", 1e-8, 100.0, log=True),
            "reg_alpha": trial.suggest_float("reg_alpha", 1e-8, 100.0, log=True),
            "subsample": trial.suggest_float("subsample", 0.1, 1.0),
            "colsample_bytree": trial.suggest_float("colsample_bytree", 0.1, 1.0),
            "max_depth": trial.suggest_int("max_depth", 1, 9),
            "early_stopping_rounds": trial.suggest_int(
                "early_stopping_rounds", 100, 500
            ),
            "n_estimators": trial.suggest_categorical(
                "n_estimators", [7000, 15000, 20000]
            ),
        }
        if model_config.use_gpu:
            params["tree_method"] = "gpu_hist"
            params["gpu_id"] = 0
            params["predictor"] = "gpu_predictor"
        else:
            params["tree_method"] = trial.suggest_categorical(
                "tree_method", ["exact", "approx", "hist"]
            )
            params["booster"] = trial.suggest_categorical(
                "booster", ["gbtree", "gblinear"]
            )
            if params["booster"] == "gbtree":
                params["gamma"] = trial.suggest_float("gamma", 1e-8, 1.0, log=True)
                params["grow_policy"] = trial.suggest_categorical(
                    "grow_policy", ["depthwise", "lossguide"]
                )

    elif model_config.algo == Algo.lgbm.value:
        params = {
            "learning_rate": trial.suggest_float("learning_rate", 1e-2, 0.25, log=True),
            "max_depth": trial.suggest_int("max_depth", 1, 9),
            "early_stopping_rounds": trial.suggest_int(
                "early_stopping_rounds", 100, 500
            ),
            "n_estimators": trial.suggest_categorical(
                "n_estimators", [7000, 15000, 20000]
            ),
            "lambda_l1": trial.suggest_float("lambda_l1", 1e-8, 10.0, log=True),
            "lambda_l2": trial.suggest_float("lambda_l2", 1e-8, 10.0, log=True),
            "num_leaves": trial.suggest_int("num_leaves", 2, 256),
            "feature_fraction": trial.suggest_float("feature_fraction", 0.4, 1.0),
            "bagging_fraction": trial.suggest_float("bagging_fraction", 0.4, 1.0),
            "bagging_freq": trial.suggest_int("bagging_freq", 1, 7),
            "min_child_samples": trial.suggest_int("min_child_samples", 5, 100),
        }
    else:
        raise NotImplementedError

    return params
