"""Base classes for tkmlr regressors.

The design follows the scikit-learn style where possible:
    fit(X, y)
    predict(X, return_std=False)
    get_params()
    set_params(**params)

Bayesian optimizers should be implemented separately.  A PHYSBO Gaussian-process
surrogate can live here as a regressor, while ask/tell adaptive search belongs
in a separate optimizer module/package.
"""
from __future__ import annotations

from dataclasses import dataclass, field
from typing import Any, Dict, Optional, Tuple

import numpy as np


ArrayLike = Any


class BaseRegressor:
    """Small sklearn-like base class.

    Subclasses only need to override ``fit`` and ``predict``.  This class is kept
    intentionally light so that wrappers for sklearn, PHYSBO, GPy, GPflow, etc.
    can share the same external interface without inheriting from sklearn.
    """

    supports_std: bool = False
    supports_score: bool = False
    name: str = "base"

    def fit(self, X: ArrayLike, y: ArrayLike) -> "BaseRegressor":
        raise NotImplementedError

    def predict(self, X: ArrayLike, return_std: bool = False):
        raise NotImplementedError

    def get_params(self, deep: bool = True) -> Dict[str, Any]:
        params: Dict[str, Any] = {}
        for key, value in self.__dict__.items():
            if key.endswith("_"):
                continue
            if key.startswith("_"):
                continue
            params[key] = value
        return params

    def set_params(self, **params: Any) -> "BaseRegressor":
        for key, value in params.items():
            setattr(self, key, value)
        return self

    def score(self, X: ArrayLike, y: ArrayLike) -> float:
        from .metrics import r2

        y_pred = self.predict(X)
        return r2(y, y_pred)


def as_2d_array(X: ArrayLike) -> np.ndarray:
    X_np = np.asarray(X, dtype=float)
    if X_np.ndim == 1:
        X_np = X_np.reshape(-1, 1)
    return X_np


def as_1d_array(y: ArrayLike) -> np.ndarray:
    y_np = np.asarray(y, dtype=float)
    if y_np.ndim == 2 and y_np.shape[1] == 1:
        y_np = y_np[:, 0]
    return y_np