import numpy as np
import pandas as pd

from sklearn.ensemble import ExtraTreesRegressor, HistGradientBoostingRegressor, RandomForestRegressor
from sklearn.multioutput import MultiOutputRegressor
from sklearn.neighbors import KNeighborsRegressor
from sklearn.pipeline import Pipeline
from sklearn.preprocessing import StandardScaler

from config import RANDOM_STATE

try:
    from xgboost import XGBClassifier, XGBRegressor
except ImportError:
    XGBClassifier = None
    XGBRegressor = None


def xgb_regressor_params(use_gpu=False, deep=False):
    params = {
        "n_estimators": 1200 if deep else 900,
        "max_depth": 8 if deep else 7,
        "learning_rate": 0.02 if deep else 0.025,
        "subsample": 0.90,
        "colsample_bytree": 0.90,
        "reg_lambda": 1.5,
        "reg_alpha": 0.05,
        "objective": "reg:squarederror",
        "random_state": RANDOM_STATE,
        "n_jobs": -1,
    }

    if use_gpu:
        params.update({
            "tree_method": "gpu_hist",
            "predictor": "gpu_predictor",
        })
    else:
        params.update({
            "tree_method": "hist",
        })

    return params


def make_xgboost_regressor(use_gpu=False, deep=False):
    if XGBRegressor is None:
        raise ImportError("未安装 xgboost，请先运行 pip install xgboost")

    base = XGBRegressor(**xgb_regressor_params(use_gpu=use_gpu, deep=deep))
    return MultiOutputRegressor(base)


def make_xgboost_classifier(num_classes, use_gpu=False):
    if XGBClassifier is None:
        raise ImportError("未安装 xgboost，请先运行 pip install xgboost")

    params = {
        "n_estimators": 700,
        "max_depth": 6,
        "learning_rate": 0.035,
        "subsample": 0.90,
        "colsample_bytree": 0.90,
        "reg_lambda": 1.5,
        "reg_alpha": 0.05,
        "objective": "multi:softprob",
        "num_class": num_classes,
        "random_state": RANDOM_STATE,
        "n_jobs": -1,
        "eval_metric": "mlogloss",
    }

    if use_gpu:
        params.update({
            "tree_method": "gpu_hist",
            "predictor": "gpu_predictor",
        })
    else:
        params.update({
            "tree_method": "hist",
        })

    return XGBClassifier(**params)


def get_baseline_models(use_gpu=False):
    models = {
        "RandomForest": RandomForestRegressor(
            n_estimators=300,
            max_depth=30,
            min_samples_leaf=2,
            random_state=RANDOM_STATE,
            n_jobs=-1,
        ),
        "ExtraTrees": ExtraTreesRegressor(
            n_estimators=500,
            max_depth=None,
            min_samples_leaf=2,
            random_state=RANDOM_STATE,
            n_jobs=-1,
        ),
        "HistGradientBoosting": MultiOutputRegressor(
            HistGradientBoostingRegressor(
                max_iter=500,
                learning_rate=0.04,
                max_leaf_nodes=31,
                l2_regularization=0.01,
                random_state=RANDOM_STATE,
            )
        ),
        "KNN": Pipeline([
            ("scaler", StandardScaler()),
            ("knn", KNeighborsRegressor(n_neighbors=15, weights="distance"))
        ]),
    }

    if XGBRegressor is not None:
        models["XGBoost"] = make_xgboost_regressor(use_gpu=use_gpu, deep=False)
        models["XGBoost_Deep"] = make_xgboost_regressor(use_gpu=use_gpu, deep=True)

    return models
