"""Wrappers for sklearn regressors."""
from __future__ import annotations

from dataclasses import dataclass
from typing import Any, Optional

import numpy as np

from .base import BaseRegressor, as_1d_array, as_2d_array
from .registry import register_model


@dataclass
class SklearnRegressor(BaseRegressor):
    """Generic wrapper around a scikit-learn estimator."""

    estimator: Any
    name: str = "sklearn"

    def fit(self, X, y):
        self.estimator.fit(as_2d_array(X), as_1d_array(y))
        return self

    def predict(self, X, return_std: bool = False):
        X_np = as_2d_array(X)
        if return_std:
            try:
                y_pred, y_std = self.estimator.predict(X_np, return_std=True)
                return y_pred, y_std
            except TypeError:
                y_pred = self.estimator.predict(X_np)
                return y_pred, None
        return self.estimator.predict(X_np)

    def get_params(self, deep: bool = True):
        return self.estimator.get_params(deep=deep)

    def set_params(self, **params):
        self.estimator.set_params(**params)
        return self


@register_model("linear")
def build_linear(**params):
    from sklearn.linear_model import LinearRegression

    return SklearnRegressor(LinearRegression(**params), name="linear")


@register_model("ridge")
def build_ridge(alpha: float = 1.0, **params):
    from sklearn.linear_model import Ridge

    return SklearnRegressor(Ridge(alpha=alpha, **params), name="ridge")


@register_model("lasso")
def build_lasso(alpha: float = 1.0, max_iter: int = 10000, **params):
    from sklearn.linear_model import Lasso

    return SklearnRegressor(Lasso(alpha=alpha, max_iter=max_iter, **params), name="lasso")


@register_model("elastic")
def build_elastic(alpha: float = 1.0, l1_ratio: float = 0.5, max_iter: int = 10000, **params):
    from sklearn.linear_model import ElasticNet

    return SklearnRegressor(ElasticNet(alpha=alpha, l1_ratio=l1_ratio, max_iter=max_iter, **params), name="elastic")


@register_model("mlp")
def build_mlp(hidden_layer_sizes=(20, 20), max_iter: int = 2000, random_state: Optional[int] = 0, **params):
    from sklearn.neural_network import MLPRegressor

    return SklearnRegressor(MLPRegressor(hidden_layer_sizes=hidden_layer_sizes, max_iter=max_iter, random_state=random_state, **params), name="mlp")


@register_model("rf")
@register_model("random_forest")
def build_random_forest(n_estimators: int = 200, random_state: Optional[int] = 0, **params):
    from sklearn.ensemble import RandomForestRegressor

    return SklearnRegressor(RandomForestRegressor(n_estimators=n_estimators, random_state=random_state, **params), name="random_forest")


@register_model("gpr")
@register_model("gp")
def build_gpr(alpha: float = 1.0e-6, kernel: Any = None, n_restarts_optimizer: int = 5, **params):
    from sklearn.gaussian_process import GaussianProcessRegressor
    from sklearn.gaussian_process.kernels import ConstantKernel as C, RBF, WhiteKernel

    if kernel is None:
        kernel = C(1.0, (1.0e-3, 1.0e3)) * RBF(1.0, (1.0e-4, 1.0e2)) + WhiteKernel(noise_level=alpha)
    model = GaussianProcessRegressor(kernel=kernel, alpha=alpha, n_restarts_optimizer=n_restarts_optimizer, **params)
    wrapped = SklearnRegressor(model, name="gpr")
    wrapped.supports_std = True
    return wrapped


@register_model("svr")
def build_svr(C: float = 1.0, epsilon: float = 0.1, kernel: str = "rbf", **params):
    from sklearn.svm import SVR

    return SklearnRegressor(SVR(C=C, epsilon=epsilon, kernel=kernel, **params), name="svr")