"""Core MMM engine: build, fit, and extract results using PyMC-Marketing."""

import logging
import time
from dataclasses import dataclass, field

import arviz as az
import numpy as np
import pandas as pd
import warnings as _warnings

from mixlift.modeling.optimizer import format_channel_name

from pymc_marketing.mmm import GeometricAdstock, LogisticSaturation

# pymc-marketing ≥0.18 deprecates `from pymc_marketing.mmm import MMM` in favour
# of the multidimensional variant.  Try the new path first, then fall back.
try:
    from pymc_marketing.mmm.multidimensional import MMM
except ImportError:
    with _warnings.catch_warnings():
        _warnings.filterwarnings("ignore", category=FutureWarning)
        from pymc_marketing.mmm import MMM  # noqa: F811

from mixlift.modeling.defaults import DEFAULT_MCMC_CONFIG, DEFAULT_SEASONALITY

logger = logging.getLogger(__name__)


@dataclass
class ModelResults:
    """Container for fitted model results."""

    channel_contributions: dict[str, float]  # channel -> mean contribution
    channel_contributions_ci: dict[str, tuple[float, float]]  # channel -> (low, high)
    roas_estimates: dict[str, float]  # channel -> average ROAS
    roas_ci: dict[str, tuple[float, float]]  # channel -> (low, high)
    marginal_roas_estimates: dict[str, float]  # channel -> marginal ROAS
    marginal_roas_ci: dict[str, tuple[float, float]]  # channel -> (low, high)
    total_spend: dict[str, float]  # channel -> total spend
    duration_secs: float
    mmm: MMM  # the fitted model object
    convergence_warnings: list[str]  # MCMC convergence warnings
    posterior_predictive_warning: str = ""  # warning if posterior predictive failed
    raw_contributions: object = None  # xarray DataArray, reusable by optimizer


def build_model(
    channel_columns: list[str],
    adstock_l_max: int = 8,
    yearly_seasonality: int | None = None,
) -> MMM:
    """Build a PyMC-Marketing MMM with opinionated defaults.

    Args:
        channel_columns: List of channel spend column names in X
        adstock_l_max: Maximum adstock lag in weeks
        yearly_seasonality: Number of Fourier terms for yearly seasonality
    """
    if yearly_seasonality is None:
        yearly_seasonality = DEFAULT_SEASONALITY["yearly_seasonality"]

    mmm = MMM(
        date_column="date",
        adstock=GeometricAdstock(l_max=adstock_l_max),
        saturation=LogisticSaturation(),
        yearly_seasonality=yearly_seasonality,
        channel_columns=channel_columns,
    )

    logger.info(
        f"Built MMM with {len(channel_columns)} channels, "
        f"adstock_l_max={adstock_l_max}, seasonality={yearly_seasonality}"
    )

    return mmm


def fit_model(
    mmm: MMM,
    X: pd.DataFrame,
    y: pd.Series,
    mcmc_config: dict | None = None,
) -> tuple[MMM, str]:
    """Fit the MMM model using MCMC sampling.

    Args:
        mmm: Built MMM model
        X: Feature DataFrame with channel spend columns + date + controls
        y: Target series (revenue or conversions)
        mcmc_config: Override MCMC settings

    Returns:
        Tuple of (fitted MMM model, posterior predictive warning string)
    """
    config = {**DEFAULT_MCMC_CONFIG, **(mcmc_config or {})}

    logger.info(
        f"Starting model fitting: {config['chains']} chains, "
        f"{config['tune']} tune, {config['draws']} draws"
    )

    start = time.time()

    X = X.reset_index(drop=True)
    y = y.reset_index(drop=True)

    fit_kwargs = dict(
        X=X,
        y=y,
        chains=config["chains"],
        tune=config["tune"],
        draws=config["draws"],
        target_accept=config["target_accept"],
        cores=1,  # Avoid multiprocessing fork deadlocks on macOS ARM
    )
    # Only pass random_seed if explicitly provided (backward compat for tests)
    seed = config.get("random_seed")
    if seed is not None:
        fit_kwargs["random_seed"] = seed

    mmm.fit(**fit_kwargs)

    # Generate posterior predictive for model fit metrics (R-squared, MAPE)
    pp_warning = ""
    try:
        mmm.sample_posterior_predictive(X, extend_idata=True, combined=False)
    except Exception as e:
        pp_warning = (
            f"Model fit metrics could not be validated — interpret results with caution. Error: {e}"
        )
        logger.warning(f"Could not generate posterior predictive: {e}")

    elapsed = time.time() - start
    logger.info(f"Model fitting completed in {elapsed:.1f}s")

    return mmm, pp_warning


def extract_channel_contributions(
    mmm: MMM,
    X: pd.DataFrame,
    channel_columns: list[str],
    contributions=None,
) -> tuple[dict[str, float], dict[str, tuple[float, float]]]:
    """Extract channel contribution estimates from fitted model.

    Args:
        contributions: Pre-computed contributions xarray (avoids recomputation).

    Returns:
        Tuple of (mean contributions dict, credible interval dict)
    """
    # Get posterior channel contributions (xarray DataArray)
    if contributions is None:
        contributions = mmm.compute_channel_contribution_original_scale()

    means = {}
    cis = {}

    for channel in channel_columns:
        display_name = format_channel_name(channel)

        # Select this channel, sum over date, stack chain+draw into samples
        channel_contrib = contributions.sel(channel=channel).sum(dim="date")
        samples = channel_contrib.values.flatten()
        means[display_name] = float(np.mean(samples))
        cis[display_name] = (
            float(np.percentile(samples, 5)),
            float(np.percentile(samples, 95)),
        )

    return means, cis


def extract_roas(
    mmm: MMM,
    X: pd.DataFrame,
    channel_columns: list[str],
    contributions=None,
) -> tuple[dict[str, float], dict[str, tuple[float, float]]]:
    """Extract average ROAS (total contribution / total spend) per channel.

    This is average ROAS, not marginal. Average ROAS tells you the overall
    return across all spend; marginal ROAS tells you the return on the next
    dollar spent. See extract_marginal_roas for the marginal variant.

    Args:
        contributions: Pre-computed contributions xarray (avoids recomputation).

    Returns:
        Tuple of (mean average ROAS dict, credible interval dict)
    """
    if contributions is None:
        contributions = mmm.compute_channel_contribution_original_scale()

    roas_means = {}
    roas_cis = {}

    for channel in channel_columns:
        display_name = format_channel_name(channel)

        total_spend = X[channel].sum()
        if total_spend < 1e-6:
            roas_means[display_name] = 0.0
            roas_cis[display_name] = (0.0, 0.0)
            continue

        channel_contrib = contributions.sel(channel=channel).sum(dim="date")
        samples = channel_contrib.values.flatten()
        roas_samples = samples / total_spend

        roas_means[display_name] = float(np.mean(roas_samples))
        roas_cis[display_name] = (
            float(np.percentile(roas_samples, 5)),
            float(np.percentile(roas_samples, 95)),
        )

    return roas_means, roas_cis


def extract_marginal_roas(
    mmm: MMM,
    X: pd.DataFrame,
    channel_columns: list[str],
    delta_pct: float = 0.01,
) -> tuple[dict[str, float], dict[str, tuple[float, float]]]:
    """Extract marginal ROAS per channel via finite-difference approximation.

    Marginal ROAS = d(contribution) / d(spend) evaluated at current spend level.
    Computed by bumping spend by delta_pct and measuring the change in contribution.

    Args:
        mmm: Fitted MMM model
        X: Feature DataFrame
        channel_columns: Channel column names
        delta_pct: Fractional bump to apply (default 1%)

    Returns:
        Tuple of (mean marginal ROAS dict, credible interval dict)
    """
    # Use channel_contribution_forward_pass for finite-difference marginal ROAS.
    # This takes preprocessed channel data as numpy and returns numpy arrays.
    channel_data_base = X[channel_columns].values  # (n_obs, n_channels)
    contributions_base = mmm.channel_contribution_forward_pass(channel_data_base)

    mroas_means = {}
    mroas_cis = {}

    for i, channel in enumerate(channel_columns):
        display_name = format_channel_name(channel)

        mean_spend = float(X[channel].mean())
        if mean_spend < 1e-6:
            mroas_means[display_name] = 0.0
            mroas_cis[display_name] = (0.0, 0.0)
            continue

        delta = mean_spend * delta_pct
        channel_data_bumped = channel_data_base.copy()
        channel_data_bumped[:, i] += delta

        contributions_bumped = mmm.channel_contribution_forward_pass(channel_data_bumped)

        # Sum over time for each posterior sample
        base_total = contributions_base[:, :, i].sum(axis=1)
        bumped_total = contributions_bumped[:, :, i].sum(axis=1)

        total_delta_spend = delta * len(X)
        mroas_samples = (bumped_total - base_total) / total_delta_spend

        mroas_means[display_name] = float(np.mean(mroas_samples))
        mroas_cis[display_name] = (
            float(np.percentile(mroas_samples, 5)),
            float(np.percentile(mroas_samples, 95)),
        )

    return mroas_means, mroas_cis


def extract_saturation_params(mmm: MMM, channel_columns: list[str]) -> dict[str, dict[str, float]]:
    """Extract posterior mean saturation parameters (lam, beta) per channel.

    PyMC-Marketing's LogisticSaturation uses: beta * (1 - exp(-lam * x)) / (1 + exp(-lam * x))

    Returns:
        dict mapping display_name -> {"lam": float, "beta": float}
    """
    posterior = mmm.idata.posterior
    lam_samples = posterior["saturation_lam"].values  # (chains, draws, n_channels)
    beta_samples = posterior["saturation_beta"].values

    params = {}
    for i, col in enumerate(channel_columns):
        display_name = format_channel_name(col)
        params[display_name] = {
            "lam": float(np.mean(lam_samples[:, :, i])),
            "beta": float(np.mean(beta_samples[:, :, i])),
        }
    return params


def extract_adstock_params(mmm: MMM, channel_columns: list[str]) -> dict[str, dict[str, float]]:
    """Extract posterior mean adstock parameters (alpha) per channel.

    GeometricAdstock uses alpha as the retention rate.

    Returns:
        dict mapping display_name -> {"alpha": float}
    """
    posterior = mmm.idata.posterior
    alpha_samples = posterior["adstock_alpha"].values  # (chains, draws, n_channels)

    params = {}
    for i, col in enumerate(channel_columns):
        display_name = format_channel_name(col)
        params[display_name] = {
            "alpha": float(np.mean(alpha_samples[:, :, i])),
        }
    return params


def logistic_saturation(spend: np.ndarray, lam: float) -> np.ndarray:
    """Apply logistic saturation matching PyMC-Marketing's formula.

    Formula: (1 - exp(-lam * x)) / (1 + exp(-lam * x))

    This is the raw saturation curve without the beta scaling.
    To get the full channel response, multiply by beta.
    """
    # Clamp exponent to prevent exp underflow (exp(-36) ≈ 2e-16 → result ≈ 1.0 exactly,
    # which gives zero gradient in the optimizer). Clamping at 30 keeps a small gradient.
    exponent = np.clip(lam * spend, 0, 30)
    return (1 - np.exp(-exponent)) / (1 + np.exp(-exponent))


def geometric_adstock_np(
    x: np.ndarray, alpha: float, l_max: int = 8, normalize: bool = True
) -> np.ndarray:
    """Apply geometric adstock transformation (NumPy version).

    Convolves spend with geometrically decaying weights [1, alpha, alpha^2, ...].
    At time t, output[t] = sum(weights[k] * x[t-k]) for k=0..l_max-1,
    treating x[t] = 0 for t < 0.

    Args:
        x: 1-D spend time series.
        alpha: Retention rate (0 to 1).
        l_max: Maximum lag.
        normalize: Whether to normalize weights to sum to 1.

    Returns:
        Adstocked spend array of same length as x.
    """
    weights = np.power(alpha, np.arange(l_max))
    if normalize:
        weights = weights / weights.sum()
    # np.convolve(x, weights) computes sum(x[k] * weights[t-k]) = sum(weights[k] * x[t-k])
    # which is exactly the causal convolution we need.
    # Take only the first len(x) elements (causal part, no future leakage).
    conv = np.convolve(x, weights)
    return conv[: len(x)]


def extract_saturation_curves(
    mmm: MMM,
    X: pd.DataFrame,
    channel_columns: list[str],
    n_points: int = 100,
) -> dict[str, dict[str, list[float]]]:
    """Generate saturation curve data for each channel.

    Returns:
        dict mapping channel_name -> {"spend": [...], "response": [...]}
    """
    params = extract_saturation_params(mmm, channel_columns)
    curves = {}

    for col in channel_columns:
        display_name = format_channel_name(col)
        max_spend = float(X[col].max()) * 2
        spend = np.linspace(0, max_spend, n_points)
        lam = params[display_name]["lam"]
        beta = params[display_name]["beta"]
        response = beta * logistic_saturation(spend, lam)
        curves[display_name] = {
            "spend": spend.tolist(),
            "response": response.tolist(),
        }

    return curves


def check_convergence(mmm: MMM) -> list[str]:
    """Check MCMC convergence diagnostics and return warnings.

    Checks:
    - R-hat > 1.05 for any parameter (indicates chains haven't converged)
    - Divergent transitions (indicates sampling problems)

    Returns:
        List of warning strings (empty if all checks pass)
    """
    warnings = []
    idata = mmm.idata

    # Check R-hat (requires ≥2 chains)
    try:
        n_chains = idata.posterior.sizes.get("chain", 1)
    except Exception:
        n_chains = 1

    if n_chains < 2:
        logger.info("R-hat not computed (requires ≥2 chains). Using single-chain diagnostics.")
        warnings.append(
            "R-hat not computed (requires ≥2 chains). "
            "Consider using ≥2 chains for more reliable convergence diagnostics."
        )
    else:
        try:
            rhat = az.rhat(idata)
            for var_name in rhat.data_vars:
                values = rhat[var_name].values
                max_rhat = float(np.max(values))
                if max_rhat > 1.05:
                    warnings.append(
                        f"R-hat for '{var_name}' is {max_rhat:.3f} (> 1.05). "
                        f"Chains may not have converged. Consider increasing tune/draws."
                    )
        except Exception as e:
            warnings.append(f"Could not compute R-hat: {e}")

    # Check divergences
    try:
        if hasattr(idata, "sample_stats"):
            divergent = idata.sample_stats.get("diverging", None)
            if divergent is not None:
                n_divergent = int(np.sum(divergent.values))
                if n_divergent > 0:
                    total_samples = int(np.prod(divergent.shape))
                    pct = 100 * n_divergent / total_samples
                    warnings.append(
                        f"{n_divergent} divergent transitions ({pct:.1f}% of samples). "
                        f"Consider increasing target_accept or reparameterizing the model."
                    )
    except Exception as e:
        warnings.append(f"Could not check divergences: {e}")

    if warnings:
        for w in warnings:
            logger.warning(f"Convergence issue: {w}")
    else:
        logger.info("All convergence diagnostics passed.")

    return warnings


def run_full_pipeline(
    X: pd.DataFrame,
    y: pd.Series,
    channel_columns: list[str],
    mcmc_config: dict | None = None,
    adstock_l_max: int = 8,
) -> ModelResults:
    """Run the complete modeling pipeline: build -> fit -> extract.

    Args:
        X: Feature DataFrame
        y: Target series
        channel_columns: Channel spend column names
        mcmc_config: Optional MCMC overrides
        adstock_l_max: Max adstock lag

    Returns:
        ModelResults with all extracted information
    """
    start = time.time()

    mmm = build_model(channel_columns, adstock_l_max=adstock_l_max)
    mmm, pp_warning = fit_model(mmm, X, y, mcmc_config)

    # Check convergence
    convergence_warnings = check_convergence(mmm)

    # Compute contributions once (expensive posterior operation) and reuse
    raw_contributions = mmm.compute_channel_contribution_original_scale()

    contributions, contributions_ci = extract_channel_contributions(
        mmm, X, channel_columns, contributions=raw_contributions
    )
    roas, roas_ci = extract_roas(
        mmm, X, channel_columns, contributions=raw_contributions
    )
    marginal_roas, marginal_roas_ci = extract_marginal_roas(mmm, X, channel_columns)

    total_spend = {}
    for channel in channel_columns:
        display_name = format_channel_name(channel)
        total_spend[display_name] = float(X[channel].sum())

    duration = time.time() - start

    return ModelResults(
        channel_contributions=contributions,
        channel_contributions_ci=contributions_ci,
        roas_estimates=roas,
        roas_ci=roas_ci,
        marginal_roas_estimates=marginal_roas,
        marginal_roas_ci=marginal_roas_ci,
        total_spend=total_spend,
        duration_secs=duration,
        mmm=mmm,
        convergence_warnings=convergence_warnings,
        posterior_predictive_warning=pp_warning,
        raw_contributions=raw_contributions,
    )
