"""Tests for surrogate split discovery."""

import numpy as np
import pytest

from surrogate_forest._surrogate import find_surrogates


class TestSurrogates:
    def test_perfectly_correlated_feature(self):
        """Feature 1 = Feature 0 → lambda should be near 1.0."""
        rng = np.random.RandomState(42)
        n = 200
        X = np.zeros((n, 3))
        X[:, 0] = rng.randn(n)
        X[:, 1] = X[:, 0]  # Perfect copy
        X[:, 2] = rng.randn(n)  # Independent

        indices = np.arange(n, dtype=np.intp)
        surrogates = find_surrogates(
            X, indices,
            primary_feature=0, primary_threshold=0.0,
            primary_missing_direction=-1, max_surrogates=5,
        )

        assert len(surrogates) > 0
        # Feature 1 should be the best surrogate with lambda ≈ 1.0
        best = surrogates[0]
        assert best.feature == 1
        assert best.lambda_ > 0.95
        assert not best.mirror

    def test_negatively_correlated_feature(self):
        """Feature 1 = -Feature 0 → lambda near 1.0, mirror=True."""
        rng = np.random.RandomState(42)
        n = 200
        X = np.zeros((n, 3))
        X[:, 0] = rng.randn(n)
        X[:, 1] = -X[:, 0]  # Perfect negative correlation
        X[:, 2] = rng.randn(n)

        indices = np.arange(n, dtype=np.intp)
        surrogates = find_surrogates(
            X, indices,
            primary_feature=0, primary_threshold=0.0,
            primary_missing_direction=-1, max_surrogates=5,
        )

        assert len(surrogates) > 0
        best = surrogates[0]
        assert best.feature == 1
        assert best.lambda_ > 0.95
        assert best.mirror

    def test_independent_features_low_lambda(self):
        """Independent features should have lambda near 0."""
        rng = np.random.RandomState(42)
        n = 1000
        X = rng.randn(n, 3)  # All independent

        indices = np.arange(n, dtype=np.intp)
        surrogates = find_surrogates(
            X, indices,
            primary_feature=0, primary_threshold=0.0,
            primary_missing_direction=-1, max_surrogates=5,
        )

        for s in surrogates:
            assert s.lambda_ < 0.15

    def test_max_surrogates_limit(self):
        """Should return at most max_surrogates."""
        rng = np.random.RandomState(42)
        n = 200
        X = np.zeros((n, 10))
        X[:, 0] = rng.randn(n)
        for i in range(1, 10):
            X[:, i] = X[:, 0] + rng.randn(n) * 0.1 * i

        indices = np.arange(n, dtype=np.intp)
        surrogates = find_surrogates(
            X, indices,
            primary_feature=0, primary_threshold=0.0,
            primary_missing_direction=-1, max_surrogates=3,
        )
        assert len(surrogates) <= 3

    def test_zero_max_surrogates(self):
        rng = np.random.RandomState(42)
        X = rng.randn(50, 3)
        indices = np.arange(50, dtype=np.intp)
        surrogates = find_surrogates(
            X, indices, 0, 0.0, -1, max_surrogates=0,
        )
        assert len(surrogates) == 0
