"""Tests for scenario mode confidence intervals."""

import json
from unittest.mock import MagicMock, patch

import numpy as np
import pytest


class TestScenarioCIOutput:
    """Test that _run_scenario includes CI data when posterior is available."""

    def test_scenario_includes_ci_keys(self):
        """Scenario result should include impact_ci_80 and direction_note."""
        from mixlift_mcp.server import _run_scenario

        # Mock model_results with posterior samples
        mock_mmm = MagicMock()
        mock_posterior = MagicMock()

        n_chains, n_draws, n_channels = 4, 250, 2
        mock_posterior.__getitem__ = lambda self, key: MagicMock(
            values=np.random.uniform(0.1, 2.0, (n_chains, n_draws, n_channels))
        )
        mock_mmm.idata.posterior = mock_posterior
        mock_mmm.adstock.l_max = 8
        mock_mmm.adstock.normalize = True

        mock_results = MagicMock()
        mock_results.mmm = mock_mmm

        import pandas as pd
        X = pd.DataFrame({
            "meta_spend": np.full(30, 5000.0),
            "google_spend": np.full(30, 3000.0),
        })

        optimization = {
            "current_allocation": {"Meta": 5000.0, "Google": 3000.0},
        }

        with patch(
            "mixlift.modeling.optimizer._estimate_response", return_value=100.0
        ):
            result = _run_scenario(
                "cut Meta by 20%",
                optimization,
                mock_results,
                X,
                ["meta_spend", "google_spend"],
            )

        assert result is not None
        assert "projected_impact_pct" in result
        assert "impact_ci_80" in result
        assert "low" in result["impact_ci_80"]
        assert "high" in result["impact_ci_80"]
        assert "direction_note" in result

    def test_scenario_parse_failure(self):
        """Unparseable scenario returns error, no CI."""
        from mixlift_mcp.server import _run_scenario

        result = _run_scenario(
            "do something weird",
            {"current_allocation": {"Meta": 5000}},
            MagicMock(),
            MagicMock(),
            ["meta_spend"],
        )
        assert "error" in result


class TestEstimateScenarioImpactSamples:
    """Test the posterior sampling function directly."""

    def test_returns_array(self):
        from mixlift.modeling.optimizer import estimate_scenario_impact_samples

        mock_mmm = MagicMock()
        n_chains, n_draws, n_channels = 2, 50, 2

        mock_posterior = MagicMock()
        mock_posterior.__getitem__ = lambda self, key: MagicMock(
            values=np.random.uniform(0.1, 1.0, (n_chains, n_draws, n_channels))
        )
        mock_mmm.idata.posterior = mock_posterior
        mock_mmm.adstock.l_max = 4
        mock_mmm.adstock.normalize = True

        import pandas as pd
        X = pd.DataFrame({
            "meta_spend": np.full(20, 5000.0),
            "google_spend": np.full(20, 3000.0),
        })

        samples = estimate_scenario_impact_samples(
            mock_mmm, X, ["meta_spend", "google_spend"],
            current_allocation={"Meta": 5000.0, "Google": 3000.0},
            modified_allocation={"Meta": 4000.0, "Google": 3000.0},
            n_posterior_samples=50,
        )

        assert isinstance(samples, np.ndarray)
        assert len(samples) > 0
        assert len(samples) <= 50
