"""Tests for GeoRFRegressor."""
import numpy as np
import pytest
from sklearn.metrics import r2_score
from sklearn.model_selection import train_test_split

from georf import GeoRFRegressor

_FAST = dict(n_estimators=10, n_samples_per_tree=200, random_state=0)


# ── 1. Smoke test ──────────────────────────────────────────────────────────────

def test_smoke():
    rng = np.random.default_rng(0)
    X = rng.standard_normal((150, 4))
    y = rng.standard_normal(150)
    reg = GeoRFRegressor(**_FAST)
    reg.fit(X, y)
    preds = reg.predict(X[:20])
    assert preds.shape == (20,)


# ── 2. Clean linear ───────────────────────────────────────────────────────────

def test_linear():
    rng = np.random.default_rng(1)
    X = rng.standard_normal((300, 4))
    y = 2.0 * X[:, 0] + X[:, 1]
    X_tr, X_te, y_tr, y_te = train_test_split(X, y, test_size=0.25, random_state=1)

    reg = GeoRFRegressor(n_estimators=20, n_samples_per_tree=500, random_state=1)
    reg.fit(X_tr, y_tr)
    r2 = r2_score(y_te, reg.predict(X_te))
    assert r2 > 0.7, f"Expected R² > 0.7 on linear data, got {r2:.3f}"


# ── 3. Nonlinear ──────────────────────────────────────────────────────────────

def test_nonlinear():
    rng = np.random.default_rng(2)
    X = rng.uniform(-2, 2, size=(300, 3))
    y = np.sin(X[:, 0]) + X[:, 1] ** 2
    X_tr, X_te, y_tr, y_te = train_test_split(X, y, test_size=0.25, random_state=2)

    reg = GeoRFRegressor(n_estimators=20, n_samples_per_tree=500, random_state=2)
    reg.fit(X_tr, y_tr)
    r2 = r2_score(y_te, reg.predict(X_te))
    assert r2 > 0.5, f"Expected R² > 0.5 on nonlinear data, got {r2:.3f}"


# ── 4. n_samples_per_tree scaling ────────────────────────────────────────────

def test_n_samples_scaling():
    """More samples per tree should not hurt R²."""
    rng = np.random.default_rng(3)
    X = rng.standard_normal((200, 3))
    y = 3 * X[:, 0] - X[:, 1] + rng.standard_normal(200) * 0.3
    X_tr, X_te, y_tr, y_te = train_test_split(X, y, test_size=0.25, random_state=3)

    r2_low = r2_score(
        y_te,
        GeoRFRegressor(n_estimators=15, n_samples_per_tree=100, random_state=3)
        .fit(X_tr, y_tr)
        .predict(X_te),
    )
    r2_high = r2_score(
        y_te,
        GeoRFRegressor(n_estimators=15, n_samples_per_tree=1000, random_state=3)
        .fit(X_tr, y_tr)
        .predict(X_te),
    )
    assert r2_high >= r2_low - 0.05, (
        f"Higher n_samples_per_tree degraded R²: {r2_low:.3f} vs {r2_high:.3f}"
    )
