import warnings

import numpy as np
import pandas as pd
from sklearn.exceptions import ConvergenceWarning

from femsurrogate.surrogates.registry import MODEL_NAMES, make_model


def _toy_regression_data():
    X = pd.DataFrame(
        {
            "L_m": np.linspace(1.0, 3.0, 12),
            "b_m": np.linspace(0.02, 0.08, 12),
            "h_m": np.linspace(0.04, 0.16, 12),
            "E_pa": np.linspace(100e9, 220e9, 12),
            "P_n": np.linspace(100.0, 2000.0, 12),
        }
    )
    y = -X["P_n"] * X["L_m"] ** 3 / (3.0 * X["E_pa"] * X["b_m"] * X["h_m"] ** 3)
    return X, y


def test_registry_exposes_expected_surrogate_names():
    assert MODEL_NAMES == ("rsm", "gpr", "random_forest", "gradient_boosting", "mlp")


def test_make_model_builds_estimators_that_fit_and_predict():
    X, y = _toy_regression_data()
    fast_overrides = {
        "random_forest": {"n_estimators": 5, "n_jobs": 1},
        "gradient_boosting": {"n_estimators": 5},
        "mlp": {"hidden_layer_sizes": (4,), "max_iter": 25, "early_stopping": False},
    }

    for model_name in MODEL_NAMES:
        model = make_model(
            model_name,
            random_state=20260521,
            **fast_overrides.get(model_name, {}),
        )
        with warnings.catch_warnings():
            warnings.filterwarnings("ignore", category=ConvergenceWarning)
            model.fit(X, y)
        predictions = model.predict(X.iloc[:3])

        assert predictions.shape == (3,)