"""Stress tests for missing data at various rates and mechanisms."""

import numpy as np
import pytest

from surrogate_forest import (
    SurrogateDecisionTreeClassifier,
    SurrogateRandomForestClassifier,
    SurrogateDecisionTreeRegressor,
    SurrogateRandomForestRegressor,
)


def _inject_mcar(X, rate, rng):
    """Inject Missing Completely At Random."""
    X = X.copy()
    mask = rng.rand(*X.shape) < rate
    X[mask] = np.nan
    return X


def _inject_mar(X, rate, rng):
    """Inject Missing At Random — missingness in feature j depends on feature 0."""
    X = X.copy()
    n, p = X.shape
    for j in range(1, p):
        prob = rate * (1 + X[:, 0])  # Higher X[:,0] → more missing
        prob = np.clip(prob, 0, 1)
        mask = rng.rand(n) < prob
        X[mask, j] = np.nan
    return X


def _inject_mnar(X, rate, rng):
    """Inject Missing Not At Random — missingness depends on the value itself."""
    X = X.copy()
    n, p = X.shape
    for j in range(p):
        threshold = np.nanpercentile(X[:, j], 100 * (1 - rate))
        X[X[:, j] > threshold, j] = np.nan
    return X


class TestMissingRates:
    @pytest.mark.parametrize("rate", [0.1, 0.3, 0.5, 0.8])
    def test_classifier_mcar_rates(self, rate):
        rng = np.random.RandomState(42)
        X = rng.randn(300, 5)
        y = (X[:, 0] + X[:, 1] > 0).astype(int)
        X_miss = _inject_mcar(X, rate, rng)

        clf = SurrogateDecisionTreeClassifier(max_depth=5, random_state=42)
        clf.fit(X_miss, y)
        preds = clf.predict(X_miss)

        # Should not crash and should produce valid predictions
        assert not np.any(np.isnan(preds.astype(float)))
        assert len(np.unique(preds)) >= 1

    @pytest.mark.parametrize("rate", [0.1, 0.3, 0.5, 0.8])
    def test_regressor_mcar_rates(self, rate):
        rng = np.random.RandomState(42)
        X = rng.randn(300, 5)
        y = X[:, 0] + 0.5 * X[:, 1] + rng.randn(300) * 0.1
        X_miss = _inject_mcar(X, rate, rng)

        reg = SurrogateDecisionTreeRegressor(max_depth=5, random_state=42)
        reg.fit(X_miss, y)
        preds = reg.predict(X_miss)
        assert not np.any(np.isnan(preds))

    @pytest.mark.parametrize("rate", [0.1, 0.3, 0.5])
    def test_forest_mcar_rates(self, rate):
        rng = np.random.RandomState(42)
        X = rng.randn(200, 5)
        y = (X[:, 0] + X[:, 1] > 0).astype(int)
        X_miss = _inject_mcar(X, rate, rng)

        rf = SurrogateRandomForestClassifier(
            n_estimators=10, max_depth=5, random_state=42
        )
        rf.fit(X_miss, y)
        preds = rf.predict(X_miss)
        assert not np.any(np.isnan(preds.astype(float)))


class TestMissingMechanisms:
    def test_mar(self):
        rng = np.random.RandomState(42)
        X = rng.randn(300, 5)
        y = (X[:, 0] > 0).astype(int)
        X_miss = _inject_mar(X, 0.3, rng)

        clf = SurrogateRandomForestClassifier(
            n_estimators=10, max_depth=5, random_state=42
        )
        clf.fit(X_miss, y)
        preds = clf.predict(X_miss)
        assert not np.any(np.isnan(preds.astype(float)))

    def test_mnar(self):
        rng = np.random.RandomState(42)
        X = rng.randn(300, 5)
        y = (X[:, 0] > 0).astype(int)
        X_miss = _inject_mnar(X, 0.3, rng)

        clf = SurrogateRandomForestClassifier(
            n_estimators=10, max_depth=5, random_state=42
        )
        clf.fit(X_miss, y)
        preds = clf.predict(X_miss)
        assert not np.any(np.isnan(preds.astype(float)))


class TestGracefulDegradation:
    def test_accuracy_decreases_with_more_missing(self):
        """Accuracy should generally decrease as missing rate increases."""
        rng = np.random.RandomState(42)
        X = rng.randn(500, 5)
        y = (X[:, 0] + X[:, 1] > 0).astype(int)

        accuracies = []
        for rate in [0.0, 0.2, 0.5]:
            X_miss = _inject_mcar(X, rate, rng) if rate > 0 else X.copy()
            clf = SurrogateRandomForestClassifier(
                n_estimators=20, max_depth=5, random_state=42
            )
            clf.fit(X_miss, y)
            acc = (clf.predict(X_miss) == y).mean()
            accuracies.append(acc)

        # Clean should be best or close
        assert accuracies[0] >= accuracies[-1] - 0.05

    def test_all_nan_column(self):
        """Should handle a column that is entirely NaN."""
        rng = np.random.RandomState(42)
        X = rng.randn(100, 3)
        y = (X[:, 0] > 0).astype(int)
        X[:, 2] = np.nan  # Entire column missing

        clf = SurrogateDecisionTreeClassifier(max_depth=3, random_state=42)
        clf.fit(X, y)
        preds = clf.predict(X)
        assert not np.any(np.isnan(preds.astype(float)))

    def test_all_nan_row(self):
        """Should handle a row that is entirely NaN."""
        rng = np.random.RandomState(42)
        X = rng.randn(100, 3)
        y = (X[:, 0] > 0).astype(int)

        clf = SurrogateDecisionTreeClassifier(max_depth=3, random_state=42)
        clf.fit(X, y)

        # Predict on a row with all NaN
        X_test = np.full((1, 3), np.nan)
        preds = clf.predict(X_test)
        assert len(preds) == 1
