"""Tests for impurity criteria."""

import numpy as np
import pytest

from surrogate_forest._impurity import (
    gini_impurity,
    entropy_impurity,
    mse_impurity,
    mae_impurity,
    impurity_improvement,
)


class TestGini:
    def test_pure_node(self):
        counts = np.array([10.0, 0.0])
        assert gini_impurity(counts) == 0.0

    def test_balanced_binary(self):
        counts = np.array([5.0, 5.0])
        assert abs(gini_impurity(counts) - 0.5) < 1e-10

    def test_imbalanced(self):
        counts = np.array([9.0, 1.0])
        expected = 1.0 - (0.9**2 + 0.1**2)
        assert abs(gini_impurity(counts) - expected) < 1e-10

    def test_empty(self):
        assert gini_impurity(np.array([0.0, 0.0])) == 0.0

    def test_multiclass(self):
        counts = np.array([1.0, 1.0, 1.0])
        expected = 1.0 - 3 * (1/3)**2
        assert abs(gini_impurity(counts) - expected) < 1e-10


class TestEntropy:
    def test_pure_node(self):
        assert entropy_impurity(np.array([10.0, 0.0])) == 0.0

    def test_balanced_binary(self):
        counts = np.array([5.0, 5.0])
        assert abs(entropy_impurity(counts) - 1.0) < 1e-10

    def test_empty(self):
        assert entropy_impurity(np.array([0.0, 0.0])) == 0.0


class TestMSE:
    def test_constant(self):
        y = np.array([5.0, 5.0, 5.0])
        assert mse_impurity(y) == 0.0

    def test_simple(self):
        y = np.array([1.0, 3.0])
        assert abs(mse_impurity(y) - 1.0) < 1e-10

    def test_weighted(self):
        y = np.array([1.0, 3.0])
        w = np.array([1.0, 1.0])
        assert abs(mse_impurity(y, w) - 1.0) < 1e-10

    def test_empty(self):
        assert mse_impurity(np.array([])) == 0.0


class TestMAE:
    def test_constant(self):
        y = np.array([5.0, 5.0, 5.0])
        assert mae_impurity(y) == 0.0

    def test_simple(self):
        y = np.array([1.0, 3.0])
        assert abs(mae_impurity(y) - 1.0) < 1e-10

    def test_empty(self):
        assert mae_impurity(np.array([])) == 0.0


class TestImpurityImprovement:
    def test_perfect_split(self):
        # Parent: [5, 5] → Gini 0.5. Left: [5, 0] → 0. Right: [0, 5] → 0.
        gain = impurity_improvement(10, 0.5, 5, 0.0, 5, 0.0)
        assert abs(gain - 0.5) < 1e-10

    def test_no_improvement(self):
        gain = impurity_improvement(10, 0.5, 5, 0.5, 5, 0.5)
        assert abs(gain) < 1e-10
