"""Tests for GeoRFClassifier."""
import numpy as np
import pytest
from sklearn.datasets import make_classification
from sklearn.metrics import accuracy_score
from sklearn.model_selection import train_test_split

from georf import GeoRFClassifier

# Small fast defaults for most tests
_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((200, 5))
    y = rng.integers(0, 2, size=200)
    clf = GeoRFClassifier(**_FAST)
    clf.fit(X, y)
    preds = clf.predict(X[:20])
    assert preds.shape == (20,)


# ── 2. Binary accuracy ────────────────────────────────────────────────────────

def test_binary_accuracy():
    X, y = make_classification(
        n_samples=300, n_features=5, n_informative=3, random_state=1
    )
    X_tr, X_te, y_tr, y_te = train_test_split(X, y, test_size=0.25, random_state=1)
    clf = GeoRFClassifier(n_estimators=20, n_samples_per_tree=300, random_state=1)
    clf.fit(X_tr, y_tr)
    acc = accuracy_score(y_te, clf.predict(X_te))
    assert acc > 0.8, f"Expected accuracy > 0.8, got {acc:.3f}"


# ── 3. predict_proba shape (binary) ──────────────────────────────────────────

def test_predict_proba_binary():
    X, y = make_classification(n_samples=200, n_features=4, random_state=2)
    clf = GeoRFClassifier(**_FAST)
    clf.fit(X, y)
    proba = clf.predict_proba(X[:30])
    assert proba.shape == (30, 2)
    np.testing.assert_allclose(proba.sum(axis=1), 1.0, atol=1e-6)


# ── 4. Multiclass accuracy ────────────────────────────────────────────────────

def test_multiclass_accuracy():
    X, y = make_classification(
        n_samples=300, n_features=6, n_informative=4,
        n_classes=3, n_clusters_per_class=1, random_state=3,
    )
    X_tr, X_te, y_tr, y_te = train_test_split(X, y, test_size=0.25, random_state=3)
    clf = GeoRFClassifier(n_estimators=20, n_samples_per_tree=300, random_state=3)
    clf.fit(X_tr, y_tr)
    acc = accuracy_score(y_te, clf.predict(X_te))
    assert acc > 0.7, f"Expected accuracy > 0.7, got {acc:.3f}"


# ── 5. Multiclass predict_proba shape ────────────────────────────────────────

def test_predict_proba_multiclass():
    X, y = make_classification(
        n_samples=200, n_features=6, n_informative=4,
        n_classes=3, n_clusters_per_class=1, random_state=4,
    )
    clf = GeoRFClassifier(**_FAST)
    clf.fit(X, y)
    proba = clf.predict_proba(X[:30])
    assert proba.shape == (30, 3)
    np.testing.assert_allclose(proba.sum(axis=1), 1.0, atol=1e-6)


# ── 6. String labels ──────────────────────────────────────────────────────────

def test_string_labels():
    rng = np.random.default_rng(5)
    X = rng.standard_normal((150, 4))
    y = np.array(["cat", "dog"] * 75)
    clf = GeoRFClassifier(**_FAST)
    clf.fit(X, y)
    preds = clf.predict(X[:20])
    assert preds.dtype.kind in ("U", "O"), "Expected string predictions"
    assert set(preds).issubset({"cat", "dog"})


# ── 7. n_samples_per_tree scaling ────────────────────────────────────────────

def test_n_samples_scaling():
    """More synthetic samples per tree should not hurt accuracy."""
    X, y = make_classification(
        n_samples=200, n_features=5, n_informative=3, random_state=6
    )
    X_tr, X_te, y_tr, y_te = train_test_split(X, y, test_size=0.3, random_state=6)

    accs = []
    for n in (100, 300, 600):
        clf = GeoRFClassifier(n_estimators=15, n_samples_per_tree=n, random_state=6)
        clf.fit(X_tr, y_tr)
        accs.append(accuracy_score(y_te, clf.predict(X_te)))

    # Allow small noise: the best accuracy is at least as good as the worst
    assert max(accs) >= min(accs) - 0.05, f"Accuracy degraded badly: {accs}"
