"""Unit tests for engine math functions, saturation, adstock, and convergence diagnostics.

These tests do NOT require MCMC fitting and run fast.
"""

import numpy as np
import pytest


class TestLogisticSaturation:
    """Test that our logistic_saturation matches PyMC-Marketing's formula."""

    def test_zero_spend_gives_zero(self):
        from mixlift.modeling.engine import logistic_saturation

        result = logistic_saturation(np.array([0.0]), lam=1.0)
        assert result[0] == pytest.approx(0.0, abs=1e-10)

    def test_large_spend_approaches_one(self):
        from mixlift.modeling.engine import logistic_saturation

        result = logistic_saturation(np.array([100.0]), lam=1.0)
        assert result[0] == pytest.approx(1.0, abs=1e-6)

    def test_formula_matches_pymc_marketing(self):
        """Verify our NumPy formula matches PyMC-Marketing's PyTensor formula."""
        from mixlift.modeling.engine import logistic_saturation
        from pymc_marketing.mmm.transformers import (
            logistic_saturation as pymc_logistic_saturation,
        )

        test_values = np.array([0.0, 0.1, 0.5, 1.0, 2.0, 5.0, 10.0])
        for lam in [0.25, 0.5, 1.0, 2.0, 4.0]:
            our_result = logistic_saturation(test_values, lam=lam)
            pymc_result = pymc_logistic_saturation(test_values, lam=lam).eval()
            np.testing.assert_allclose(
                our_result,
                pymc_result,
                rtol=1e-7,
                err_msg=f"Mismatch at lam={lam}",
            )

    def test_monotonically_increasing(self):
        from mixlift.modeling.engine import logistic_saturation

        spend = np.linspace(0, 10, 50)
        result = logistic_saturation(spend, lam=1.0)
        diffs = np.diff(result)
        assert np.all(diffs >= 0), "Saturation function should be monotonically increasing"

    def test_concave(self):
        """Second derivative should be negative (diminishing returns) for positive spend."""
        from mixlift.modeling.engine import logistic_saturation

        spend = np.linspace(0.1, 10, 50)
        result = logistic_saturation(spend, lam=1.0)
        second_diffs = np.diff(np.diff(result))
        assert np.all(second_diffs <= 1e-10), "Saturation function should be concave"

    def test_higher_lam_saturates_faster(self):
        from mixlift.modeling.engine import logistic_saturation

        spend = np.array([1.0])
        low_lam = logistic_saturation(spend, lam=0.5)[0]
        high_lam = logistic_saturation(spend, lam=2.0)[0]
        assert high_lam > low_lam, "Higher lambda should saturate faster"

    def test_vectorized(self):
        from mixlift.modeling.engine import logistic_saturation

        spend = np.array([0.0, 1.0, 2.0, 5.0])
        result = logistic_saturation(spend, lam=1.0)
        assert result.shape == (4,)


class TestGeometricAdstockNp:
    """Test that our NumPy adstock matches PyMC-Marketing's geometric adstock."""

    def test_no_decay_passes_through(self):
        """With alpha=0, no carryover, so output = input (if normalize=False)."""
        from mixlift.modeling.engine import geometric_adstock_np

        x = np.array([100.0, 200.0, 0.0, 50.0])
        result = geometric_adstock_np(x, alpha=0.0, l_max=4, normalize=False)
        np.testing.assert_allclose(result, x, rtol=1e-7)

    def test_full_decay_smooths_heavily(self):
        """With alpha close to 1, carryover should spread the effect."""
        from mixlift.modeling.engine import geometric_adstock_np

        x = np.array([100.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0])
        result = geometric_adstock_np(x, alpha=0.9, l_max=8, normalize=False)
        # The first value should be 100
        assert result[0] == pytest.approx(100.0, rel=1e-5)
        # Subsequent values should be decaying carryover
        assert result[1] > 0
        assert result[1] < result[0]
        # Each subsequent should decay
        for i in range(2, len(result)):
            assert result[i] < result[i - 1]

    def test_output_length_matches_input(self):
        from mixlift.modeling.engine import geometric_adstock_np

        x = np.array([10.0, 20.0, 30.0, 40.0, 50.0])
        result = geometric_adstock_np(x, alpha=0.5, l_max=4)
        assert len(result) == len(x)

    def test_normalized_weights_sum_to_one(self):
        """With normalized weights, a constant spend should yield approximately the same value."""
        from mixlift.modeling.engine import geometric_adstock_np

        # After enough time, a constant spend of 100 with normalized weights
        # should converge to 100 (since weights sum to 1)
        x = np.full(50, 100.0)
        result = geometric_adstock_np(x, alpha=0.5, l_max=8, normalize=True)
        # After the first l_max periods, should converge to 100
        np.testing.assert_allclose(result[-1], 100.0, rtol=0.01)

    def test_matches_pymc_marketing(self):
        """Verify our NumPy adstock matches PyMC-Marketing's geometric adstock."""
        from pymc_marketing.mmm.transformers import geometric_adstock as pymc_geometric

        from mixlift.modeling.engine import geometric_adstock_np

        x = np.array([100.0, 50.0, 200.0, 0.0, 75.0, 30.0, 150.0, 80.0, 0.0, 60.0])

        for alpha in [0.0, 0.3, 0.5, 0.7, 0.9]:
            for normalize in [True, False]:
                our_result = geometric_adstock_np(
                    x, alpha=alpha, l_max=8, normalize=normalize
                )
                pymc_result = pymc_geometric(
                    x, alpha=alpha, l_max=8, normalize=normalize
                ).eval()
                np.testing.assert_allclose(
                    our_result,
                    pymc_result,
                    rtol=1e-5,
                    err_msg=f"Mismatch at alpha={alpha}, normalize={normalize}",
                )


class TestCheckConvergence:
    """Test convergence diagnostic checks using mock inference data."""

    def _make_mock_mmm_with_idata(self, rhat_values, n_divergent=0):
        """Create a mock MMM with controlled R-hat and divergence data."""
        from unittest.mock import MagicMock

        import xarray as xr

        mmm = MagicMock()

        # Build a mock posterior with one variable
        posterior = xr.Dataset(
            {"param_a": xr.DataArray(rhat_values, dims=["chain", "draw"])}
        )

        # Build mock sample stats
        if n_divergent > 0:
            chains, draws = rhat_values.shape
            div = np.zeros((chains, draws), dtype=bool)
            div.flat[:n_divergent] = True
            sample_stats = xr.Dataset(
                {"diverging": xr.DataArray(div, dims=["chain", "draw"])}
            )
        else:
            chains, draws = rhat_values.shape
            sample_stats = xr.Dataset(
                {
                    "diverging": xr.DataArray(
                        np.zeros((chains, draws), dtype=bool), dims=["chain", "draw"]
                    )
                }
            )

        # Build InferenceData mock
        import arviz as az

        idata = az.InferenceData(posterior=posterior, sample_stats=sample_stats)
        mmm.idata = idata

        return mmm

    def test_good_convergence_no_warnings(self):
        from mixlift.modeling.engine import check_convergence

        # 2 chains, 100 draws, values drawn from similar distributions (good convergence)
        rng = np.random.default_rng(42)
        values = rng.normal(1.0, 0.1, size=(2, 100))
        mmm = self._make_mock_mmm_with_idata(values, n_divergent=0)

        warnings = check_convergence(mmm)
        # Good convergence should produce no R-hat warnings
        rhat_warnings = [w for w in warnings if "R-hat" in w]
        div_warnings = [w for w in warnings if "divergent" in w]
        assert len(rhat_warnings) == 0
        assert len(div_warnings) == 0

    def test_divergent_transitions_warned(self):
        from mixlift.modeling.engine import check_convergence

        rng = np.random.default_rng(42)
        values = rng.normal(1.0, 0.1, size=(2, 100))
        mmm = self._make_mock_mmm_with_idata(values, n_divergent=10)

        warnings = check_convergence(mmm)
        div_warnings = [w for w in warnings if "divergent" in w]
        assert len(div_warnings) == 1
        assert "10" in div_warnings[0]

    def test_bad_rhat_warned(self):
        from mixlift.modeling.engine import check_convergence

        # 2 chains with very different means (bad convergence)
        chain_1 = np.full(100, 0.0)
        chain_2 = np.full(100, 10.0)
        values = np.stack([chain_1, chain_2])
        mmm = self._make_mock_mmm_with_idata(values, n_divergent=0)

        warnings = check_convergence(mmm)
        rhat_warnings = [w for w in warnings if "R-hat" in w]
        assert len(rhat_warnings) >= 1
