"""Tests for categorical feature support."""
import numpy as np
import pytest

from georf import GeoRFClassifier, GeoRFRegressor

_FAST = dict(n_estimators=8, n_samples_per_tree=150, random_state=0)


def _mixed_dataset(n=200, seed=0):
    """3 continuous + 1 categorical (integer 0-3) features."""
    rng = np.random.default_rng(seed)
    X_cont = rng.standard_normal((n, 3))
    X_cat = rng.integers(0, 4, size=(n, 1)).astype(float)
    X = np.hstack([X_cont, X_cat])
    y = (X_cont[:, 0] + X_cat[:, 0] > 0).astype(int)
    return X, y


# ── 1. Mixed features fit/predict ────────────────────────────────────────────

def test_mixed_features_classifier():
    X, y = _mixed_dataset()
    feature_types = ["continuous", "continuous", "continuous", "categorical"]
    clf = GeoRFClassifier(**_FAST)
    clf.fit(X, y, feature_types=feature_types)
    preds = clf.predict(X[:10])
    assert preds.shape == (10,)


def test_mixed_features_regressor():
    rng = np.random.default_rng(1)
    X_cont = rng.standard_normal((200, 3))
    X_cat = rng.integers(0, 4, size=(200, 1)).astype(float)
    X = np.hstack([X_cont, X_cat])
    y = X_cont[:, 0] + X_cat[:, 0].astype(float)
    feature_types = ["continuous", "continuous", "continuous", "categorical"]
    reg = GeoRFRegressor(**_FAST)
    reg.fit(X, y, feature_types=feature_types)
    preds = reg.predict(X[:10])
    assert preds.shape == (10,)


# ── 2. feature_types=None (all continuous) ───────────────────────────────────

def test_feature_types_none_no_crash():
    X, y = _mixed_dataset()
    clf = GeoRFClassifier(**_FAST)
    clf.fit(X, y, feature_types=None)  # should not raise
    clf.predict(X[:5])
