"""Budget optimization using PyMC-Marketing's built-in optimizer."""

import logging

import numpy as np
import pandas as pd
from pymc_marketing.mmm import MMM

logger = logging.getLogger(__name__)

# Special-case display names that .title() gets wrong
_DISPLAY_NAME_OVERRIDES = {
    "tiktok": "TikTok",
    "tv": "TV",
    "seo": "SEO",
    "cpc": "CPC",
    "cpm": "CPM",
    "roi": "ROI",
    "dv360": "DV360",
}


def format_channel_name(raw_col: str) -> str:
    """Convert a column name like 'tiktok_spend' to a display name like 'TikTok'."""
    name = raw_col.replace("_spend", "").replace("_", " ").strip()
    lower = name.lower()
    if lower in _DISPLAY_NAME_OVERRIDES:
        return _DISPLAY_NAME_OVERRIDES[lower]
    return name.title()


def optimize_budget(
    mmm: MMM,
    X: pd.DataFrame,
    channel_columns: list[str],
    total_budget: float | None = None,
    budget_bounds: dict[str, tuple[float, float]] | None = None,
    contributions=None,
) -> dict:
    """Optimize budget allocation across channels.

    Args:
        mmm: Fitted MMM model
        X: Feature DataFrame used for fitting
        channel_columns: Channel spend column names
        total_budget: Total weekly budget to optimize. If None, uses current total.
        budget_bounds: Optional per-channel (min, max) bounds.
        contributions: Pre-computed contributions xarray (avoids recomputation).

    Returns:
        Dictionary with current_allocation, optimal_allocation, expected_lift_pct
    """
    # Calculate current weekly averages
    n_weeks = len(X)
    current_allocation = {}
    for col in channel_columns:
        display_name = format_channel_name(col)
        current_allocation[display_name] = float(X[col].sum() / n_weeks)

    if total_budget is None:
        total_budget = sum(current_allocation.values())

    logger.info(f"Optimizing budget: ${total_budget:,.0f}/week across {len(channel_columns)} channels")

    # Reuse pre-computed contributions if available, otherwise compute
    if contributions is None:
        contributions = mmm.compute_channel_contribution_original_scale()
    # xarray DataArray -> numpy: mean over chain+draw, keep (date, channel)
    mean_contributions = contributions.mean(dim=["chain", "draw"]).values

    # Use PyMC-Marketing's optimize_budget if available, otherwise use simple gradient approach
    try:
        optimal = _optimize_with_pymc(mmm, channel_columns, total_budget, budget_bounds)
    except (AttributeError, NotImplementedError, TypeError):
        logger.warning("PyMC-Marketing optimize_budget not available, using gradient approach")
        optimal = _optimize_gradient(
            mmm, X, channel_columns, total_budget, mean_contributions=mean_contributions
        )

    # Calculate expected lift
    current_total_response = _estimate_response(
        mmm, X, channel_columns, current_allocation, mean_contributions=mean_contributions
    )
    optimal_total_response = _estimate_response(
        mmm, X, channel_columns, optimal, mean_contributions=mean_contributions
    )

    if current_total_response > 0:
        lift_pct = ((optimal_total_response - current_total_response) / current_total_response) * 100
    else:
        lift_pct = 0.0

    return {
        "current_allocation": current_allocation,
        "optimal_allocation": optimal,
        "total_budget": total_budget,
        "expected_lift_pct": round(lift_pct, 1),
    }


def _optimize_with_pymc(
    mmm: MMM,
    channel_columns: list[str],
    total_budget: float,
    budget_bounds: dict | None = None,
) -> dict[str, float]:
    """Use PyMC-Marketing's built-in budget optimizer."""
    result = mmm.optimize_budget(
        total_budget=total_budget,
        budget_bounds=budget_bounds,
    )

    optimal = {}
    for i, col in enumerate(channel_columns):
        display_name = format_channel_name(col)
        optimal[display_name] = float(result[col])

    return optimal


def _optimize_gradient(
    mmm: MMM,
    X: pd.DataFrame,
    channel_columns: list[str],
    total_budget: float,
    n_steps: int = 100,
    mean_contributions: np.ndarray | None = None,
) -> dict[str, float]:
    """Simple gradient-based budget optimization.

    Iteratively moves budget from lowest marginal ROAS channel to highest.
    """
    from mixlift.modeling.engine import extract_adstock_params, extract_saturation_params

    n_channels = len(channel_columns)
    allocation = np.array([total_budget / n_channels] * n_channels)
    step_size = total_budget * 0.02  # 2% of budget per step

    # Extract params once outside the loop
    sat_params = extract_saturation_params(mmm, channel_columns)
    adstock_params = extract_adstock_params(mmm, channel_columns)

    # Pre-compute mean contributions if not provided
    if mean_contributions is None:
        contributions = mmm.compute_channel_contribution_original_scale()
        mean_contributions = contributions.mean(dim=["chain", "draw"]).values

    for _ in range(n_steps):
        # Estimate marginal ROAS for each channel
        marginal_roas = np.zeros(n_channels)
        for i in range(n_channels):
            alloc_up = allocation.copy()
            alloc_up[i] += step_size
            alloc_down = allocation.copy()
            alloc_down[i] = max(0, alloc_down[i] - step_size)

            response_up = _estimate_response_from_array(
                mmm, X, channel_columns, alloc_up,
                _sat_params=sat_params,
                _adstock_params=adstock_params,
                _mean_contributions=mean_contributions,
            )
            response_down = _estimate_response_from_array(
                mmm, X, channel_columns, alloc_down,
                _sat_params=sat_params,
                _adstock_params=adstock_params,
                _mean_contributions=mean_contributions,
            )

            delta_spend = alloc_up[i] - alloc_down[i]
            if delta_spend > 0:
                marginal_roas[i] = (response_up - response_down) / delta_spend

        # Move from worst to best
        worst = np.argmin(marginal_roas)
        best = np.argmax(marginal_roas)

        if worst != best and allocation[worst] > step_size:
            move = min(step_size, allocation[worst])
            allocation[worst] -= move
            allocation[best] += move

    # Normalize to exact budget
    allocation = allocation * (total_budget / allocation.sum())

    optimal = {}
    for i, col in enumerate(channel_columns):
        display_name = format_channel_name(col)
        optimal[display_name] = round(float(allocation[i]), 2)

    return optimal


def estimate_scenario_impact_samples(
    mmm: MMM,
    X: pd.DataFrame,
    channel_columns: list[str],
    current_allocation: dict[str, float],
    modified_allocation: dict[str, float],
    n_posterior_samples: int = 200,
) -> np.ndarray:
    """Estimate scenario impact for each posterior sample (for uncertainty propagation).

    Samples from the joint posterior of (saturation_lam, saturation_beta, adstock_alpha)
    and computes the response ratio for each sample. Returns an array of impact percentages.
    """
    from mixlift.modeling.engine import geometric_adstock_np, logistic_saturation

    posterior = mmm.idata.posterior
    lam_all = posterior["saturation_lam"].values.reshape(-1, len(channel_columns))
    beta_all = posterior["saturation_beta"].values.reshape(-1, len(channel_columns))
    alpha_all = posterior["adstock_alpha"].values.reshape(-1, len(channel_columns))

    n_total = lam_all.shape[0]
    idx = np.random.choice(n_total, size=min(n_posterior_samples, n_total), replace=False)

    l_max = mmm.adstock.l_max
    normalize = mmm.adstock.normalize

    impact_samples = []
    for i in idx:
        current_resp = 0.0
        modified_resp = 0.0
        for j, col in enumerate(channel_columns):
            display_name = format_channel_name(col)
            lam = float(lam_all[i, j])
            alpha = float(alpha_all[i, j])

            cur_spend = current_allocation.get(display_name, 0)
            mod_spend = modified_allocation.get(display_name, 0)

            # Adstock then saturate for both allocations
            cur_series = np.full(len(X), cur_spend)
            mod_series = np.full(len(X), mod_spend)

            cur_adstocked = geometric_adstock_np(cur_series, alpha, l_max=l_max, normalize=normalize).mean()
            mod_adstocked = geometric_adstock_np(mod_series, alpha, l_max=l_max, normalize=normalize).mean()

            cur_sat = logistic_saturation(np.array([cur_adstocked]), lam)[0]
            mod_sat = logistic_saturation(np.array([mod_adstocked]), lam)[0]

            # beta scales both equally, so for ratio it cancels, but we need
            # absolute contributions to sum across channels
            beta = float(beta_all[i, j])
            current_resp += beta * cur_sat
            modified_resp += beta * mod_sat

        if current_resp > 0:
            impact_samples.append((modified_resp - current_resp) / current_resp * 100)

    return np.array(impact_samples)


def _estimate_response(
    mmm: MMM,
    X: pd.DataFrame,
    channel_columns: list[str],
    allocation: dict[str, float],
    mean_contributions: np.ndarray | None = None,
) -> float:
    """Estimate total response for a given budget allocation."""
    alloc_array = np.array([
        allocation[format_channel_name(col)]
        for col in channel_columns
    ])
    return _estimate_response_from_array(
        mmm, X, channel_columns, alloc_array, _mean_contributions=mean_contributions
    )


def _estimate_response_from_array(
    mmm: MMM,
    X: pd.DataFrame,
    channel_columns: list[str],
    allocation: np.ndarray,
    _sat_params: dict | None = None,
    _adstock_params: dict | None = None,
    _mean_contributions: np.ndarray | None = None,
) -> float:
    """Estimate total response from allocation array using model's saturation curves.

    Applies geometric adstock to the spend before computing the saturation ratio,
    matching what the model does internally during fitting.

    The saturation function is: beta * (1 - exp(-lam * x)) / (1 + exp(-lam * x))
    where x is the adstock-transformed spend, not the raw spend.
    """
    from mixlift.modeling.engine import (
        extract_adstock_params,
        extract_saturation_params,
        geometric_adstock_np,
        logistic_saturation,
    )

    total = 0.0

    if _mean_contributions is None:
        contributions = mmm.compute_channel_contribution_original_scale()
        _mean_contributions = contributions.mean(dim=["chain", "draw"]).values

    if _sat_params is None:
        _sat_params = extract_saturation_params(mmm, channel_columns)

    if _adstock_params is None:
        _adstock_params = extract_adstock_params(mmm, channel_columns)

    # Get adstock l_max from the model's adstock component
    l_max = mmm.adstock.l_max
    normalize = mmm.adstock.normalize

    for i, col in enumerate(channel_columns):
        display_name = format_channel_name(col)
        current_spend = X[col].mean()
        if current_spend > 0:
            lam = _sat_params[display_name]["lam"]
            beta = _sat_params[display_name]["beta"]
            adstock_alpha = _adstock_params[display_name]["alpha"]

            # Apply geometric adstock to current spend time series
            current_adstocked = geometric_adstock_np(
                X[col].values, adstock_alpha, l_max=l_max, normalize=normalize
            )
            current_adstocked_mean = float(current_adstocked.mean())

            # Apply geometric adstock to the new constant-spend time series
            new_spend_series = np.full(len(X), allocation[i])
            new_adstocked = geometric_adstock_np(
                new_spend_series, adstock_alpha, l_max=l_max, normalize=normalize
            )
            new_adstocked_mean = float(new_adstocked.mean())

            # Compute saturation ratio using adstock-transformed values
            sat_new = logistic_saturation(np.array([new_adstocked_mean]), lam)[0]
            sat_current = logistic_saturation(np.array([current_adstocked_mean]), lam)[0]

            if sat_current > 0:
                # Scale the known contribution by the ratio of saturated values
                # beta cancels out in the ratio, so we don't need it here
                scaled = _mean_contributions[:, i].sum() * (sat_new / sat_current)
            else:
                scaled = 0.0
            total += scaled

    return float(total)
