"""Tests for split finding."""

import numpy as np
import pytest

from surrogate_forest._splitter import (
    find_best_split_classification,
    find_best_split_regression,
)


class TestClassificationSplit:
    def test_finds_obvious_split(self):
        # Feature 0 perfectly separates classes
        X = np.array([[0.0], [1.0], [2.0], [3.0], [4.0], [5.0], [6.0], [7.0]])
        y = np.array([0, 0, 0, 0, 1, 1, 1, 1])
        w = np.ones(8)
        indices = np.arange(8, dtype=np.intp)

        split = find_best_split_classification(
            X, y, w, indices, np.array([0]), n_classes=2
        )
        assert split.feature == 0
        assert 3.0 < split.threshold < 4.0
        assert split.gain > 0

    def test_nan_direction_learning(self):
        # NaN samples are class 1, so they should go right
        X = np.array([[0.0], [1.0], [2.0], [np.nan], [4.0], [5.0], [np.nan], [7.0]])
        y = np.array([0, 0, 0, 1, 1, 1, 1, 1])
        w = np.ones(8)
        indices = np.arange(8, dtype=np.intp)

        split = find_best_split_classification(
            X, y, w, indices, np.array([0]), n_classes=2
        )
        assert split.feature == 0
        assert split.gain > 0

    def test_no_valid_split(self):
        # All same class
        X = np.array([[0.0], [1.0], [2.0]])
        y = np.array([0, 0, 0])
        w = np.ones(3)
        indices = np.arange(3, dtype=np.intp)

        split = find_best_split_classification(
            X, y, w, indices, np.array([0]), n_classes=1
        )
        assert split.gain <= 0


class TestRegressionSplit:
    def test_finds_obvious_split(self):
        X = np.array([[0.0], [1.0], [2.0], [3.0], [10.0], [11.0], [12.0], [13.0]])
        y = np.array([0.0, 0.5, 0.3, 0.1, 10.0, 10.5, 10.3, 10.1])
        w = np.ones(8)
        indices = np.arange(8, dtype=np.intp)

        split = find_best_split_regression(
            X, y, w, indices, np.array([0]),
        )
        assert split.feature == 0
        assert split.gain > 0

    def test_nan_handling(self):
        X = np.array([[0.0], [1.0], [np.nan], [3.0], [10.0], [11.0], [np.nan], [13.0]])
        y = np.array([0.0, 0.5, 10.0, 0.1, 10.0, 10.5, 0.3, 10.1])
        w = np.ones(8)
        indices = np.arange(8, dtype=np.intp)

        split = find_best_split_regression(
            X, y, w, indices, np.array([0]),
        )
        assert split.feature == 0
        assert split.gain > 0
