import numpy as np
from sklearn.preprocessing import StandardScaler  # type: ignore
from sklearn.metrics import mean_squared_error, r2_score  # type: ignore
from sklearn.utils import resample  # type: ignore
from sklearn.model_selection import KFold  # type: ignore


def get_RNG() -> np.random.Generator:
    return np.random.default_rng(314)


def polynomial_features(x: np.ndarray, p: int, intercept: bool = True) -> np.ndarray:
    """Generates a design matrix with polynomial features up to degree p.
    Args:
        x: The input data vector of shape (n_samples,).
        p: The maximum polynomial degree.
        intercept: Whether to include the intercept term (degree 0).
    Returns:
        The design matrix of shape (n_samples, p + 1) if intercept is True,
        else (n_samples, p).
    """
    if intercept:
        P = np.arange(p + 1)
    else:
        P = np.arange(1, p + 1)
    X = np.power(x[:, np.newaxis], P)
    return X


def runge_function(x: np.ndarray) -> np.ndarray:
    """Computes the Runge function values for the input x.
    Args:
        x: The input data vector of shape (n_samples,).
    Returns:
        The Runge function values of shape (n_samples,).
    """
    return 1 / (1 + 25 * np.square(x))


def scale_data(
    X: np.ndarray, X_test: np.ndarray | None = None
) -> tuple[np.ndarray, np.ndarray | None]:
    """Scales the input data using StandardScaler.
    Args:
        X: The training data matrix of shape (n_samples, n_features).
        X_test: The test data matrix of shape (n_samples_test, n_features) or None.
    Returns:
        A tuple containing the scaled training data and the scaled test data (or None if X_test is None).
    """
    scaler = StandardScaler()
    single_feature = False
    if len(X.shape) == 1:
        X = X.reshape(-1, 1)
        single_feature = True
    if X_test is not None and single_feature:
        X_test = X_test.reshape(-1, 1)
    X_scaled = scaler.fit_transform(X)
    if single_feature:
        X_scaled = X_scaled.ravel()
    if X_test is not None:
        X_test_scaled = scaler.transform(X_test)
        if single_feature:
            X_test_scaled = X_test_scaled.ravel()
        return X_scaled, X_test_scaled
    return X_scaled, None


def noise_data(y: np.ndarray, noise_level: float = 1.0) -> np.ndarray:
    """Adds Gaussian noise to the target data.
    Args:
        y: The target vector of shape (n_samples,).
        noise_level: The standard deviation of the Gaussian noise.
    Returns:
        The noisy target vector of shape (n_samples,).
    """
    noise = get_RNG().normal(0, noise_level, size=y.shape)
    return y + noise


def evaluate_model(y_true: np.ndarray, y_pred: np.ndarray) -> tuple[float, float]:
    """Evaluates the model using Mean Squared Error (MSE) and R^2 score.
    Args:
        y_true: The true target vector of shape (n_samples,).
        y_pred: The predicted target vector of shape (n_samples,).
    Returns:
        A tuple containing the MSE and R^2 score.
    """
    return mean_squared_error(y_true, y_pred), r2_score(y_true, y_pred)


def bootstrap_resample(
    X: np.ndarray, y: np.ndarray, n_resamples: int
) -> list[tuple[np.ndarray, np.ndarray]]:
    """Generates bootstrap resamples of the dataset.
    Args:
        X: The input data matrix of shape (n_samples, n_features).
        y: The target vector of shape (n_samples,).
        n_resamples: The number of bootstrap resamples to generate.
    Returns:
        A list of tuples, each containing a resampled (X_resample, y_resample).
    """
    resamples = []
    for _ in range(n_resamples):
        X_resample, y_resample = resample(
            X, y, random_state=get_RNG().integers(0, int(1e6))
        )
        resamples.append((X_resample, y_resample))
    return resamples


def k_fold_split(
    X: np.ndarray, y: np.ndarray, k: int
) -> list[tuple[np.ndarray, np.ndarray, np.ndarray, np.ndarray]]:
    """Splits the dataset into k folds for cross-validation.
    Args:
        X: The input data matrix of shape (n_samples, n_features).
        y: The target vector of shape (n_samples,).
        k: The number of folds.
    Returns:
        A list of tuples, each containing (X_train, y_train, X_val, y_val) for each fold.
    """
    kf = KFold(n_splits=k, shuffle=True, random_state=get_RNG().integers(0, int(1e6)))
    folds = []
    for train_index, val_index in kf.split(X):
        X_train, X_val = X[train_index], X[val_index]
        y_train, y_val = y[train_index], y[val_index]
        folds.append((X_train, y_train, X_val, y_val))
    return folds
