"""Tests for the budget optimizer module.

These tests use mocked models to verify optimizer logic without MCMC fitting.
"""

from unittest.mock import MagicMock, PropertyMock, patch

import numpy as np
import pandas as pd
import pytest
import xarray as xr

from mixlift.modeling.optimizer import (
    _estimate_response_from_array,
    _optimize_gradient,
)


def _make_mock_mmm(
    channel_columns,
    sat_lam_values,
    sat_beta_values,
    adstock_alpha_values,
    contributions_array,
    l_max=8,
    normalize=True,
):
    """Create a mock MMM with controlled parameters and contributions.

    Args:
        channel_columns: List of channel column names
        sat_lam_values: List of saturation lam values per channel
        sat_beta_values: List of saturation beta values per channel
        adstock_alpha_values: List of adstock alpha values per channel
        contributions_array: (n_samples, n_obs, n_channels) array
        l_max: Adstock maximum lag
        normalize: Whether adstock normalizes weights
    """
    mmm = MagicMock()

    n_channels = len(channel_columns)
    chains, draws = 2, 50

    # Build posterior
    posterior = xr.Dataset({
        "saturation_lam": xr.DataArray(
            np.broadcast_to(
                np.array(sat_lam_values).reshape(1, 1, -1),
                (chains, draws, n_channels),
            ).copy(),
            dims=["chain", "draw", "channel"],
        ),
        "saturation_beta": xr.DataArray(
            np.broadcast_to(
                np.array(sat_beta_values).reshape(1, 1, -1),
                (chains, draws, n_channels),
            ).copy(),
            dims=["chain", "draw", "channel"],
        ),
        "adstock_alpha": xr.DataArray(
            np.broadcast_to(
                np.array(adstock_alpha_values).reshape(1, 1, -1),
                (chains, draws, n_channels),
            ).copy(),
            dims=["chain", "draw", "channel"],
        ),
    })

    mmm.idata = MagicMock()
    mmm.idata.posterior = posterior
    # Wrap contributions in xarray DataArray with proper dims so .mean(dim=...) works
    # contributions_array is (n_samples, n_obs, n_channels) — broadcast to (chains, draws, n_obs, n_channels)
    n_obs = contributions_array.shape[1]
    fill_value = contributions_array[0, 0, 0]  # uniform value used in tests
    contributions_xr = xr.DataArray(
        np.full((chains, draws, n_obs, n_channels), fill_value),
        dims=["chain", "draw", "date", "channel"],
    )
    mmm.compute_channel_contribution_original_scale = MagicMock(
        return_value=contributions_xr
    )

    # Mock adstock component
    mmm.adstock = MagicMock()
    mmm.adstock.l_max = l_max
    mmm.adstock.normalize = normalize

    # Make optimize_budget raise so gradient approach is used
    mmm.optimize_budget = MagicMock(side_effect=AttributeError)

    return mmm


def _make_test_X(n_weeks, channel_columns, spend_values):
    """Create a test DataFrame.

    Args:
        n_weeks: Number of weeks
        channel_columns: List of channel column names
        spend_values: Dict of column_name -> constant weekly spend
    """
    data = {"date": pd.date_range("2024-01-01", periods=n_weeks, freq="W")}
    for col in channel_columns:
        data[col] = [spend_values.get(col, 100.0)] * n_weeks
    data["trend"] = list(range(n_weeks))
    return pd.DataFrame(data)


class TestEstimateResponseFromArray:
    """Test that _estimate_response_from_array correctly uses adstock + saturation."""

    def test_higher_spend_gives_higher_response(self):
        """More spend should produce a higher (or equal) response."""
        channels = ["google_search_spend", "meta_spend"]
        n_weeks = 20
        n_samples = 10

        # Use small spend values so saturation curve is not flat
        X = _make_test_X(n_weeks, channels, {c: 1.0 for c in channels})

        # Contributions: shape (n_samples, n_weeks, n_channels)
        contributions = np.ones((n_samples, n_weeks, 2)) * 50.0

        mmm = _make_mock_mmm(
            channels,
            sat_lam_values=[0.5, 0.5],
            sat_beta_values=[1.0, 1.0],
            adstock_alpha_values=[0.3, 0.3],
            contributions_array=contributions,
        )

        low_alloc = np.array([0.5, 0.5])
        high_alloc = np.array([3.0, 3.0])

        response_low = _estimate_response_from_array(mmm, X, channels, low_alloc)
        response_high = _estimate_response_from_array(mmm, X, channels, high_alloc)

        assert response_high > response_low

    def test_zero_current_spend_gracefully_handled(self):
        """If a channel has zero historical spend, it should be skipped."""
        channels = ["google_search_spend", "meta_spend"]
        n_weeks = 20
        n_samples = 10

        X = _make_test_X(n_weeks, channels, {"google_search_spend": 1.0, "meta_spend": 0.0})

        contributions = np.ones((n_samples, n_weeks, 2)) * 50.0

        mmm = _make_mock_mmm(
            channels,
            sat_lam_values=[0.5, 0.5],
            sat_beta_values=[1.0, 1.0],
            adstock_alpha_values=[0.3, 0.3],
            contributions_array=contributions,
        )

        alloc = np.array([1.0, 1.0])
        # Should not raise
        response = _estimate_response_from_array(mmm, X, channels, alloc)
        assert response > 0

    def test_response_is_positive_for_positive_spend(self):
        """Response should be positive when there is positive spend and contributions."""
        channels = ["channel_a_spend"]
        n_weeks = 20
        n_samples = 10

        X = _make_test_X(n_weeks, channels, {"channel_a_spend": 1.0})
        contributions = np.ones((n_samples, n_weeks, 1)) * 100.0

        mmm = _make_mock_mmm(
            channels,
            sat_lam_values=[0.5],
            sat_beta_values=[2.0],
            adstock_alpha_values=[0.3],
            contributions_array=contributions,
        )

        alloc = np.array([1.0])
        response = _estimate_response_from_array(mmm, X, channels, alloc)
        assert response > 0


class TestOptimizeGradient:
    """Test the gradient-based optimizer produces reasonable allocations."""

    def test_budget_is_preserved(self):
        """Total budget should be preserved after optimization."""
        channels = ["google_search_spend", "meta_spend"]
        n_weeks = 20
        n_samples = 10
        total_budget = 4.0

        # Use small spend values in the non-saturated regime
        X = _make_test_X(n_weeks, channels, {c: 2.0 for c in channels})

        contributions = np.ones((n_samples, n_weeks, 2)) * 50.0

        mmm = _make_mock_mmm(
            channels,
            sat_lam_values=[0.1, 0.5],  # channel 0 less saturated
            sat_beta_values=[1.0, 1.0],
            adstock_alpha_values=[0.3, 0.3],
            contributions_array=contributions,
        )

        result = _optimize_gradient(mmm, X, channels, total_budget, n_steps=20)

        total = sum(result.values())
        assert total == pytest.approx(total_budget, rel=0.01)

    def test_optimizer_favors_less_saturated_channel(self):
        """Optimizer should allocate more to the channel with lower saturation.

        Channel 0 has lam=0.05 (far from saturation at spend~2) while
        channel 1 has lam=1.0 (nearly saturated). The optimizer should
        shift budget toward channel 0 since it has higher marginal returns.
        """
        channels = ["low_sat_spend", "high_sat_spend"]
        n_weeks = 30
        n_samples = 10
        total_budget = 4.0

        # Both channels start at spend=2.0
        X = _make_test_X(n_weeks, channels, {c: 2.0 for c in channels})

        contributions = np.ones((n_samples, n_weeks, 2)) * 50.0

        mmm = _make_mock_mmm(
            channels,
            sat_lam_values=[0.05, 1.0],  # channel 0 far from saturation, channel 1 nearly there
            sat_beta_values=[1.0, 1.0],
            adstock_alpha_values=[0.3, 0.3],
            contributions_array=contributions,
        )

        result = _optimize_gradient(mmm, X, channels, total_budget, n_steps=50)

        low_sat_alloc = result["Low Sat"]
        high_sat_alloc = result["High Sat"]

        # The less saturated channel should get more budget
        assert low_sat_alloc > high_sat_alloc

    def test_all_allocations_positive(self):
        """No channel should get negative allocation."""
        channels = ["a_spend", "b_spend", "c_spend"]
        n_weeks = 20
        n_samples = 10
        total_budget = 6.0

        X = _make_test_X(n_weeks, channels, {c: 2.0 for c in channels})
        contributions = np.ones((n_samples, n_weeks, 3)) * 30.0

        mmm = _make_mock_mmm(
            channels,
            sat_lam_values=[0.5, 0.5, 0.5],
            sat_beta_values=[1.0, 1.0, 1.0],
            adstock_alpha_values=[0.3, 0.3, 0.3],
            contributions_array=contributions,
        )

        result = _optimize_gradient(mmm, X, channels, total_budget, n_steps=50)

        for channel, alloc in result.items():
            assert alloc >= 0, f"Channel {channel} has negative allocation: {alloc}"
