"""Tests for interpretability methods."""
import numpy as np
import pytest
from sklearn.datasets import make_classification

from georf import GeoRFClassifier, GeoRFRegressor

_CLF_FAST = dict(n_estimators=10, n_samples_per_tree=200, random_state=0)
_REG_FAST = dict(n_estimators=10, n_samples_per_tree=200, random_state=0)


def _fitted_clf():
    X, y = make_classification(n_samples=150, n_features=5, random_state=0)
    clf = GeoRFClassifier(**_CLF_FAST)
    clf.fit(X, y)
    return clf


def _fitted_reg():
    rng = np.random.default_rng(0)
    X = rng.standard_normal((150, 5))
    y = X[:, 0] + X[:, 1]
    reg = GeoRFRegressor(**_REG_FAST)
    reg.fit(X, y)
    return reg


# ── 1. feature_importances ────────────────────────────────────────────────────

def test_feature_importances_clf():
    clf = _fitted_clf()
    fi = clf.feature_importances()
    assert isinstance(fi, dict)
    total = sum(fi.values())
    assert abs(total - 1.0) < 0.05, f"Importances should sum to ~1.0, got {total:.4f}"


def test_feature_importances_named():
    clf = _fitted_clf()
    names = [f"col_{i}" for i in range(5)]
    fi = clf.feature_importances(feature_names=names)
    assert set(fi.keys()) == set(names)


def test_feature_importances_reg():
    reg = _fitted_reg()
    fi = reg.feature_importances()
    assert isinstance(fi, dict)
    total = sum(fi.values())
    assert abs(total - 1.0) < 0.05


# ── 2. tree_quality_scores ────────────────────────────────────────────────────

def test_tree_quality_scores_clf():
    clf = _fitted_clf()
    scores = clf.tree_quality_scores()
    assert scores.shape == (clf.n_estimators,)
    assert np.all(scores >= 0.0), "AUC values must be non-negative"
    assert np.all(scores <= 1.0), "AUC values must be <= 1.0"


def test_tree_quality_scores_reg():
    reg = _fitted_reg()
    scores = reg.tree_quality_scores()
    assert scores.shape == (reg.n_estimators,)
    # R² can be negative for bad trees; just check it's a finite array
    assert np.all(np.isfinite(scores))


# ── 3. diversity_score ────────────────────────────────────────────────────────

def test_diversity_score_clf():
    clf = _fitted_clf()
    d = clf.diversity_score()
    assert isinstance(d, float)
    assert d >= 0.0


def test_diversity_score_reg():
    reg = _fitted_reg()
    d = reg.diversity_score()
    assert isinstance(d, float)
    assert d >= 0.0


# ── 4. provenance ─────────────────────────────────────────────────────────────

def test_provenance_clf():
    clf = _fitted_clf()
    prov = clf.provenance()
    required = {
        "n_real_samples",
        "n_synthetic_per_tree",
        "n_total_synthetic",
        "n_estimators",
        "expansion_ratio",
    }
    assert required <= set(prov.keys()), f"Missing keys: {required - set(prov.keys())}"
    assert prov["n_total_synthetic"] == clf.n_estimators * clf.n_samples_per_tree
    assert prov["n_real_samples"] == 150


def test_provenance_reg():
    reg = _fitted_reg()
    prov = reg.provenance()
    assert prov["n_total_synthetic"] == reg.n_estimators * reg.n_samples_per_tree
