import logging
import os
from pathlib import Path
from typing import Any

import numpy as np
from mcp.server.fastmcp import FastMCP

logger = logging.getLogger(__name__)

mcp = FastMCP("bluealpha-meridian")

# ---------------------------------------------------------------------------
# Model loading
# ---------------------------------------------------------------------------

_model: Any = None
_trace: Any = None
_posterior: Any = None
_loaded_model_name: str | None = None


def _reset_model() -> None:
    """Clear global model state so a new model can be loaded."""
    global _model, _trace, _posterior, _loaded_model_name
    _model = None
    _trace = None
    _posterior = None
    _loaded_model_name = None


def _get_cache_dir() -> Path:
    cache_dir = Path(os.environ.get("MERIDIAN_CACHE_DIR", "~/.cache/mcp-server-meridian"))
    cache_dir = cache_dir.expanduser()
    cache_dir.mkdir(parents=True, exist_ok=True)
    return cache_dir


def _download_from_s3(bucket: str, key: str, local_path: Path) -> None:
    """Download an object from S3 to a local path."""
    import boto3

    s3 = boto3.client("s3")
    s3.download_file(bucket, key, str(local_path))


def _load_model(path: str | None = None) -> None:
    global _model, _trace, _posterior, _loaded_model_name

    if path is None and _model is not None:
        return

    if path is None:
        model_path = os.environ.get("MERIDIAN_MODEL_PATH")
        if not model_path:
            raise RuntimeError(
                "MERIDIAN_MODEL_PATH environment variable is not set. "
                "Point it to a pickled Meridian model file."
            )
    else:
        model_path = path

    if not os.path.exists(model_path):
        raise RuntimeError(f"Model file not found: {model_path}")

    from meridian.model import model as meridian_model

    _reset_model()
    _model = meridian_model.load_mmm(model_path)
    _trace = _model._inference_data
    _posterior = _trace.posterior
    _loaded_model_name = os.path.basename(model_path)


def _get_model() -> Any:
    _load_model()
    return _model


def _get_posterior() -> Any:
    _load_model()
    return _posterior


def _get_trace() -> Any:
    _load_model()
    return _trace


def _get_channels() -> list[str]:
    posterior = _get_posterior()
    return list(posterior.coords["media_channel"].values)


def _filter_channels(channels: list[str] | None) -> list[str]:
    all_channels = _get_channels()
    if channels is None:
        return all_channels
    invalid = [c for c in channels if c not in all_channels]
    if invalid:
        raise ValueError(
            f"Unknown channels: {invalid}. Available: {all_channels}"
        )
    return channels


# ---------------------------------------------------------------------------
# Tool 1: get_model_summary
# ---------------------------------------------------------------------------


@mcp.tool()
def get_model_summary() -> dict:
    """Returns metadata about the loaded Meridian model — channels, time range,
    sample count, and convergence diagnostics. Call this first to understand
    what model is loaded."""
    trace = _get_trace()
    posterior = _get_posterior()
    model = _get_model()

    channels = list(posterior.coords["media_channel"].values)
    times = posterior.coords["time"].values

    # Convergence diagnostics from sample_stats
    convergence: dict[str, Any] = {}
    try:
        sample_stats = trace.sample_stats
        diverging = sample_stats["diverging"].values
        convergence = {
            "has_divergences": bool(diverging.any()),
            "n_divergences": int(diverging.sum()),
            "mean_accept_ratio": float(sample_stats["accept_ratio"].values.mean()),
        }
    except (KeyError, AttributeError):
        convergence = {"error": "sample_stats not available in trace"}

    created_at = posterior.attrs.get("created_at", "unknown")

    # Model config from validated paths
    model_spec = None
    try:
        model_spec = model._model_context._model_spec
    except AttributeError:
        pass

    result: dict[str, Any] = {
        "model_name": os.path.splitext(
            _loaded_model_name or os.path.basename(
                os.environ.get("MERIDIAN_MODEL_PATH", "unknown")
            )
        )[0],
        "created_at": str(created_at),
        "n_chains": int(posterior.sizes["chain"]),
        "n_draws": int(posterior.sizes["draw"]),
        "media_channels": channels,
        "n_channels": len(channels),
        "geo": list(posterior.coords["geo"].values) if "geo" in posterior.coords else "unknown",
        "time_range": {
            "start": str(times[0]),
            "end": str(times[-1]),
            "n_weeks": len(times),
        },
        "convergence": convergence,
    }

    if model_spec is not None:
        result["adstock_type"] = "geometric"
        result["max_lag"] = getattr(model_spec, "max_lag", "unknown")
        result["saturation_type"] = "hill"
        result["hill_before_adstock"] = getattr(model_spec, "hill_before_adstock", "unknown")
        result["media_prior_type"] = getattr(model_spec, "media_prior_type", "unknown")
        result["media_effects_dist"] = getattr(model_spec, "media_effects_dist", "unknown")

    return result


# ---------------------------------------------------------------------------
# Tool 2: get_channel_roi
# ---------------------------------------------------------------------------


@mcp.tool()
def get_channel_roi(
    channels: list[str] | None = None,
    credible_interval: float = 0.9,
) -> dict:
    """Returns the posterior ROI distribution for each channel — median, mean,
    and credible intervals. Core metric for channel efficiency.

    Uses roi_m from the posterior when available. Falls back to computing ROI
    from contribution_m (incremental KPI / total spend per channel)."""
    posterior = _get_posterior()
    model = _get_model()
    filtered = _filter_channels(channels)
    all_channels = _get_channels()
    ci = credible_interval

    # Determine ROI source
    if "roi_m" in posterior:
        roi_source = "roi_m"
    elif "contribution_m" in posterior:
        roi_source = "contribution_m"
    else:
        return {"error": "Neither roi_m nor contribution_m found in posterior."}

    results = []
    n_samples = 0

    if roi_source == "roi_m":
        roi = posterior["roi_m"]
        n_samples = int(roi.sizes["draw"] * roi.sizes["chain"])
        for channel in filtered:
            vals = roi.sel(media_channel=channel).values.flatten()
            results.append({
                "name": channel,
                "roi_median": round(float(np.median(vals)), 4),
                "roi_mean": round(float(np.mean(vals)), 4),
                "roi_ci_lower": round(float(np.percentile(vals, (1 - ci) / 2 * 100)), 4),
                "roi_ci_upper": round(float(np.percentile(vals, (1 + ci) / 2 * 100)), 4),
                "roi_std": round(float(np.std(vals)), 4),
                "prob_positive": round(float((vals > 0).mean()), 4),
            })
    else:
        cm = posterior["contribution_m"]
        n_samples = int(cm.sizes["draw"] * cm.sizes["chain"])
        spend_np = model.media_tensors.media_spend.numpy()
        total_kpi = float(np.sum(model.input_data.kpi.values))

        for channel in filtered:
            ch_idx = all_channels.index(channel)
            contrib_frac = cm.sel(media_channel=channel).values.flatten()
            total_spend = float(np.sum(spend_np[..., ch_idx]))
            vals = contrib_frac * total_kpi / (total_spend + 1e-10)
            results.append({
                "name": channel,
                "roi_median": round(float(np.median(vals)), 4),
                "roi_mean": round(float(np.mean(vals)), 4),
                "roi_ci_lower": round(float(np.percentile(vals, (1 - ci) / 2 * 100)), 4),
                "roi_ci_upper": round(float(np.percentile(vals, (1 + ci) / 2 * 100)), 4),
                "roi_std": round(float(np.std(vals)), 4),
                "prob_positive": round(float((vals > 0).mean()), 4),
            })

    return {
        "channels": results,
        "credible_interval_width": ci,
        "n_samples": n_samples,
        "roi_source": roi_source,
    }


# ---------------------------------------------------------------------------
# Tool 3: get_marginal_roi
# ---------------------------------------------------------------------------


@mcp.tool()
def get_marginal_roi(
    channels: list[str] | None = None,
    credible_interval: float = 0.9,
) -> dict:
    """Returns the marginal ROI (mROI) — the expected return on the next dollar
    spent. Channels with high mROI have headroom; channels with low mROI are
    saturated. Computed from the Hill curve derivative at current spend.

    For average ROI comparison, uses roi_m when available, otherwise computes
    from contribution_m."""
    posterior = _get_posterior()
    model = _get_model()
    filtered = _filter_channels(channels)
    all_channels = _get_channels()
    ci = credible_interval

    ec = posterior["ec_m"]
    slope_var = posterior["slope_m"]
    alpha = posterior["alpha_m"]
    beta = posterior["beta_gm"]

    # Get current spend and scale factors
    spend_np = model.media_tensors.media_spend.numpy()
    sf = model.media_tensors.media_transformer._scale_factors_gm.numpy()
    sf = np.squeeze(sf)
    kpi_stdev = float(model.kpi_transformer._population_scaled_stdev.numpy())

    # Resolve average ROI per channel for headroom comparison
    avg_roi_by_channel: dict[str, np.ndarray] = {}
    if "roi_m" in posterior:
        roi_source = "roi_m"
        roi = posterior["roi_m"]
        for channel in filtered:
            avg_roi_by_channel[channel] = roi.sel(media_channel=channel).values.flatten()
    elif "contribution_m" in posterior:
        roi_source = "contribution_m"
        cm = posterior["contribution_m"]
        total_kpi = float(np.sum(model.input_data.kpi.values))
        for channel in filtered:
            ch_idx = all_channels.index(channel)
            contrib_frac = cm.sel(media_channel=channel).values.flatten()
            total_spend = float(np.sum(spend_np[..., ch_idx]))
            avg_roi_by_channel[channel] = contrib_frac * total_kpi / (total_spend + 1e-10)
    else:
        return {"error": "Neither roi_m nor contribution_m found in posterior."}

    results = []
    for channel in filtered:
        ch_idx = all_channels.index(channel)
        roi_vals = avg_roi_by_channel[channel]
        ec_vals = ec.sel(media_channel=channel).values.flatten()
        s_vals = slope_var.sel(media_channel=channel).values.flatten()
        a_vals = alpha.sel(media_channel=channel).values.flatten()
        b_vals = beta.sel(media_channel=channel)
        if "geo" in b_vals.dims:
            b_vals = b_vals.isel(geo=0)
        b_vals = b_vals.values.flatten()

        avg_spend = float(np.mean(spend_np[..., ch_idx]))
        x = avg_spend / (sf[ch_idx] + 1e-10)
        x_ad = x / (1 - a_vals + 1e-10)

        denom = ec_vals**s_vals + x_ad**s_vals + 1e-10
        d_hill = s_vals * ec_vals**s_vals * x_ad**(s_vals - 1) / (denom**2 + 1e-10)

        mroi_vals = b_vals * d_hill / ((1 - a_vals + 1e-10) * (sf[ch_idx] + 1e-10)) * kpi_stdev

        mroi_med = float(np.median(mroi_vals))
        roi_med = float(np.median(roi_vals))

        if mroi_med > roi_med:
            headroom = "high"
        elif mroi_med > 1.0:
            headroom = "moderate"
        else:
            headroom = "low"

        results.append({
            "name": channel,
            "mroi_median": round(mroi_med, 4),
            "mroi_mean": round(float(np.mean(mroi_vals)), 4),
            "mroi_ci_lower": round(float(np.percentile(mroi_vals, (1 - ci) / 2 * 100)), 4),
            "mroi_ci_upper": round(float(np.percentile(mroi_vals, (1 + ci) / 2 * 100)), 4),
            "headroom_signal": headroom,
            "vs_average_roi": round(mroi_med - roi_med, 4),
        })

    return {
        "channels": results,
        "credible_interval_width": ci,
        "roi_source": roi_source,
        "interpretation": (
            "Channels with mROI > average ROI have headroom for increased spend. "
            "Channels with mROI < 1.0 are destroying value at the margin."
        ),
    }


# ---------------------------------------------------------------------------
# Tool 4: get_contribution_breakdown
# ---------------------------------------------------------------------------


@mcp.tool()
def get_contribution_breakdown(
    channels: list[str] | None = None,
    as_percentage: bool = True,
) -> dict:
    """Returns the modeled revenue contribution by channel — how much of total
    revenue each channel is responsible for according to the MMM. Computed from
    Hill saturation curves and beta coefficients at actual spend levels."""
    posterior = _get_posterior()
    model = _get_model()
    filtered = _filter_channels(channels)
    all_channels = _get_channels()

    ec = posterior["ec_m"]
    slope_var = posterior["slope_m"]
    alpha = posterior["alpha_m"]
    beta = posterior["beta_gm"]

    spend_np = model.media_tensors.media_spend.numpy()  # (1, 104, 10)
    sf = np.squeeze(model.media_tensors.media_transformer._scale_factors_gm.numpy())

    # Scale from normalized KPI space to revenue
    kpi_stdev = float(model.kpi_transformer._population_scaled_stdev.numpy())

    # Compute contribution for all channels (needed for total and percentages)
    contrib_by_channel: dict[str, np.ndarray] = {}
    for ch_idx, ch in enumerate(all_channels):
        ec_vals = ec.sel(media_channel=ch).values.flatten()
        s_vals = slope_var.sel(media_channel=ch).values.flatten()
        a_vals = alpha.sel(media_channel=ch).values.flatten()
        b_vals = beta.sel(media_channel=ch)
        if "geo" in b_vals.dims:
            b_vals = b_vals.isel(geo=0)
        b_vals = b_vals.values.flatten()

        avg_spend = float(np.mean(spend_np[..., ch_idx]))
        x = avg_spend / (sf[ch_idx] + 1e-10)
        x_ad = x / (1 - a_vals + 1e-10)
        hill_response = x_ad**s_vals / (ec_vals**s_vals + x_ad**s_vals + 1e-10)
        contrib_by_channel[ch] = b_vals * hill_response * kpi_stdev

    # Total contribution across all channels per draw
    total_per_draw = sum(contrib_by_channel.values())

    results = []
    for channel in filtered:
        vals = contrib_by_channel[channel]
        pct_per_draw = vals / (total_per_draw + 1e-10) * 100

        entry: dict[str, Any] = {
            "name": channel,
            "contribution_median": round(float(np.median(vals)), 2),
            "contribution_ci_lower": round(float(np.percentile(vals, 5)), 2),
            "contribution_ci_upper": round(float(np.percentile(vals, 95)), 2),
        }
        if as_percentage:
            entry["contribution_pct"] = round(float(np.median(pct_per_draw)), 2)

        results.append(entry)

    return {
        "channels": results,
        "total_modeled_contribution_median": round(float(np.median(total_per_draw)), 2),
    }


# ---------------------------------------------------------------------------
# Tool 5: get_saturation_curves
# ---------------------------------------------------------------------------


@mcp.tool()
def get_saturation_curves(
    channels: list[str] | None = None,
) -> dict:
    """Returns the Hill saturation curve parameters for each channel and
    computes current saturation percentage. Answers: how much runway does each
    channel have before diminishing returns dominate?"""
    posterior = _get_posterior()
    model = _get_model()
    filtered = _filter_channels(channels)

    # EC and slope should be in the posterior as sampled parameters
    if "ec_m" not in posterior or "slope_m" not in posterior:
        return {
            "error": (
                "Could not locate Hill curve parameters (ec_m / slope_m) in the "
                "posterior. Available vars: " + str(list(posterior.data_vars))
            ),
        }

    ec = posterior["ec_m"]
    slope = posterior["slope_m"]

    saturation_labels = {
        "low": "< 40% — significant headroom",
        "moderate": "40-65% — healthy operating range",
        "high": "65-85% — diminishing returns",
        "very_high": "> 85% — near-ceiling, reallocate",
    }

    results = []
    for channel in filtered:
        try:
            ec_vals = ec.sel(media_channel=channel).values.flatten()
            ec_med = float(np.median(ec_vals))

            slope_vals = slope.sel(media_channel=channel).values.flatten()
            slope_med = float(np.median(slope_vals))

            # Compute saturation at current spend level
            # Pipeline: raw spend → adstock → Hill (hill_before_adstock=False)
            current_sat_pct = None
            try:
                media_tensors = model.media_tensors
                spend = media_tensors.media_spend
                if hasattr(spend, "numpy"):
                    spend = spend.numpy()
                elif hasattr(spend, "values"):
                    spend = spend.values

                ch_idx = list(posterior.coords["media_channel"].values).index(channel)
                # Average weekly spend for this channel
                current_spend = float(np.mean(spend[..., ch_idx]))

                # Scale spend using media transformer
                transformer = media_tensors.media_transformer
                if hasattr(transformer, "_scale_factors_gm"):
                    scale_factors = transformer._scale_factors_gm
                    if hasattr(scale_factors, "numpy"):
                        scale_factors = scale_factors.numpy()
                    scaled_spend = current_spend / float(np.squeeze(scale_factors)[ch_idx])
                else:
                    scaled_spend = current_spend

                # Apply geometric adstock (steady-state approximation)
                alpha_vals = posterior["alpha_m"].sel(media_channel=channel).values.flatten()
                alpha_med = float(np.median(alpha_vals))
                adstocked = scaled_spend / (1 - alpha_med + 1e-10)

                # Hill saturation: sat = x^s / (ec^s + x^s)
                sat = adstocked**slope_med / (ec_med**slope_med + adstocked**slope_med + 1e-10)
                current_sat_pct = round(sat * 100, 1)
            except Exception as e:
                logger.debug(f"Could not compute saturation for {channel}: {e}")

            if current_sat_pct is not None:
                if current_sat_pct < 40:
                    label = "low"
                elif current_sat_pct < 65:
                    label = "moderate"
                elif current_sat_pct < 85:
                    label = "high"
                else:
                    label = "very_high"
            else:
                label = "unknown"

            interpretation = None
            if current_sat_pct is not None:
                remaining = round(100 - current_sat_pct, 0)
                interpretation = (
                    f"At {current_sat_pct}% saturation — incremental spend yields "
                    f"~{remaining}% of peak efficiency."
                )

            results.append({
                "name": channel,
                "ec_median": round(ec_med, 4),
                "slope_median": round(slope_med, 4),
                "current_saturation_pct": current_sat_pct,
                "saturation_label": label,
                "interpretation": interpretation,
            })
        except Exception as e:
            results.append({"name": channel, "error": str(e)})

    return {
        "channels": results,
        "saturation_labels": saturation_labels,
    }


# ---------------------------------------------------------------------------
# Tool 6: get_adstock_parameters
# ---------------------------------------------------------------------------


@mcp.tool()
def get_adstock_parameters(
    channels: list[str] | None = None,
) -> dict:
    """Returns the adstock (carry-over) decay rates for each channel. Answers:
    how long does each channel's advertising effect last?"""
    posterior = _get_posterior()
    filtered = _filter_channels(channels)

    alpha = posterior["alpha_m"]
    results = []

    for channel in filtered:
        vals = alpha.sel(media_channel=channel).values.flatten()
        decay_median = float(np.median(vals))

        if 0 < decay_median < 1:
            half_life = float(np.log(0.5) / np.log(decay_median))
            duration = float(np.log(0.1) / np.log(decay_median))
        else:
            half_life = 0.0
            duration = 0.0

        if half_life < 1:
            interpretation = (
                f"Short carry-over — 50% of effect gone within {half_life:.1f} weeks. "
                "Needs continuous spend."
            )
        elif half_life < 3:
            interpretation = (
                f"Moderate carry-over — effect lingers for ~{duration:.0f} weeks. "
                "Can tolerate some gaps in spend."
            )
        else:
            interpretation = (
                f"Long carry-over — effect lingers for ~{duration:.0f} weeks. "
                "Can pulse spend and still see returns."
            )

        results.append({
            "name": channel,
            "decay_rate_median": round(decay_median, 4),
            "decay_rate_ci_lower": round(float(np.percentile(vals, 5)), 4),
            "decay_rate_ci_upper": round(float(np.percentile(vals, 95)), 4),
            "half_life_weeks": round(half_life, 2),
            "effect_duration_weeks": round(duration, 2),
            "interpretation": interpretation,
        })

    return {
        "channels": results,
        "adstock_type": "geometric",
    }


# ---------------------------------------------------------------------------
# Tool 7: simulate_budget_reallocation
# ---------------------------------------------------------------------------


def _get_media_spend_baseline(model: Any, all_channels: list[str]) -> dict[str, float]:
    """Extract current average weekly spend per channel from model tensors."""
    media_tensors = model.media_tensors
    spend = media_tensors.media_spend
    if hasattr(spend, "numpy"):
        spend = spend.numpy()
    elif hasattr(spend, "values"):
        spend = spend.values

    current_spends: dict[str, float] = {}
    for i, ch in enumerate(all_channels):
        current_spends[ch] = float(np.mean(spend[..., i]))
    return current_spends


def _get_media_scale_factors(model: Any, all_channels: list[str]) -> np.ndarray:
    """Extract per-channel scale factors from media transformer."""
    transformer = model.media_tensors.media_transformer
    if hasattr(transformer, "_scale_factors_gm"):
        sf = transformer._scale_factors_gm
        if hasattr(sf, "numpy"):
            sf = sf.numpy()
        return np.squeeze(sf)
    # Fallback: no scaling
    return np.ones(len(all_channels))


@mcp.tool()
def simulate_budget_reallocation(
    proposed_budgets: dict[str, float],
    n_weeks: int = 4,
    credible_interval: float = 0.9,
) -> dict:
    """Takes a proposed budget reallocation and projects the impact on revenue
    using the posterior's Hill + adstock curves. Powers the Budget Optimizer
    skill's closed-loop workflow.

    Args:
        proposed_budgets: Channel name → new weekly spend amount.
        n_weeks: Projection horizon in weeks (default: 4).
        credible_interval: Width of credible interval (default: 0.9).
    """
    posterior = _get_posterior()
    model = _get_model()
    all_channels = _get_channels()
    ci = credible_interval

    # Validate proposed channels
    invalid = [c for c in proposed_budgets if c not in all_channels]
    if invalid:
        return {"error": f"Unknown channels: {invalid}. Available: {all_channels}"}

    # Check required posterior vars
    for var in ["ec_m", "slope_m", "alpha_m", "beta_gm"]:
        if var not in posterior:
            return {"error": f"Required posterior variable '{var}' not found."}

    ec = posterior["ec_m"]
    slope = posterior["slope_m"]
    alpha = posterior["alpha_m"]
    beta = posterior["beta_gm"]

    # Get current spend baseline and scale factors
    try:
        current_spends = _get_media_spend_baseline(model, all_channels)
        scale_factors = _get_media_scale_factors(model, all_channels)
    except Exception as e:
        return {"error": f"Cannot determine current spend levels: {e}"}

    # KPI scaling: convert from normalized to revenue units
    kpi_stdev = float(model.kpi_transformer._population_scaled_stdev.numpy())

    # Fill in channels not in proposed_budgets with current spend
    full_proposed = {
        ch: proposed_budgets.get(ch, current_spends.get(ch, 0))
        for ch in all_channels
    }

    n_draws = int(posterior.sizes["draw"])
    n_chains = int(posterior.sizes["chain"])
    n_samples = n_draws * n_chains

    current_revenue_samples = np.zeros(n_samples)
    proposed_revenue_samples = np.zeros(n_samples)
    channel_current_contrib: dict[str, np.ndarray] = {}
    channel_proposed_contrib: dict[str, np.ndarray] = {}

    for idx, ch in enumerate(all_channels):
        ec_vals = ec.sel(media_channel=ch).values.flatten()
        slope_vals = slope.sel(media_channel=ch).values.flatten()
        alpha_vals = alpha.sel(media_channel=ch).values.flatten()

        beta_vals = beta.sel(media_channel=ch)
        if "geo" in beta_vals.dims:
            beta_vals = beta_vals.isel(geo=0)
        beta_vals = beta_vals.values.flatten()

        sf = float(scale_factors[idx]) if scale_factors.ndim >= 1 else float(scale_factors)

        cur_spend = current_spends.get(ch, 0)
        new_spend = full_proposed[ch]

        # Pipeline: scale → adstock (geometric steady-state) → Hill
        cur_scaled = cur_spend / (sf + 1e-10)
        new_scaled = new_spend / (sf + 1e-10)

        cur_adstocked = cur_scaled / (1 - alpha_vals + 1e-10)
        new_adstocked = new_scaled / (1 - alpha_vals + 1e-10)

        def hill(x: np.ndarray, ec_v: np.ndarray, s_v: np.ndarray) -> np.ndarray:
            return x**s_v / (ec_v**s_v + x**s_v + 1e-10)

        cur_response = beta_vals * hill(cur_adstocked, ec_vals, slope_vals) * kpi_stdev
        new_response = beta_vals * hill(new_adstocked, ec_vals, slope_vals) * kpi_stdev

        channel_current_contrib[ch] = cur_response * n_weeks
        channel_proposed_contrib[ch] = new_response * n_weeks

        current_revenue_samples += cur_response * n_weeks
        proposed_revenue_samples += new_response * n_weeks

    # Channel-level impact details
    channel_impacts = []
    optimization_notes = []

    for ch in all_channels:
        cur_spend = current_spends.get(ch, 0)
        new_spend = full_proposed[ch]

        if cur_spend == new_spend and ch not in proposed_budgets:
            continue

        cur_contrib_med = float(np.median(channel_current_contrib[ch]))
        new_contrib_med = float(np.median(channel_proposed_contrib[ch]))
        spend_change_pct = ((new_spend - cur_spend) / cur_spend * 100) if cur_spend > 0 else 0

        channel_impacts.append({
            "name": ch,
            "current_spend": round(cur_spend, 2),
            "proposed_spend": round(new_spend, 2),
            "spend_change_pct": round(spend_change_pct, 1),
            "current_contribution_median": round(cur_contrib_med, 2),
            "projected_contribution_median": round(new_contrib_med, 2),
            "contribution_change": round(new_contrib_med - cur_contrib_med, 2),
        })

        nearly_same = new_contrib_med > cur_contrib_med * 0.9
        if new_spend < cur_spend and cur_contrib_med > 0 and nearly_same:
            optimization_notes.append(
                f"Reducing {ch} spend by {abs(spend_change_pct):.0f}% loses only "
                f"{(1 - new_contrib_med / cur_contrib_med) * 100:.0f}% contribution — "
                f"was in saturation zone."
            )
        elif new_spend > cur_spend and cur_contrib_med > 0 and new_contrib_med > cur_contrib_med:
            gain_pct = (new_contrib_med / cur_contrib_med - 1) * 100
            optimization_notes.append(
                f"Increasing {ch} captures additional {gain_pct:.0f}% contribution."
            )

    current_total_spend = sum(current_spends.get(ch, 0) for ch in all_channels)
    proposed_total_spend = sum(full_proposed[ch] for ch in all_channels)
    change_samples = proposed_revenue_samples - current_revenue_samples

    return {
        "current_total_weekly_spend": round(current_total_spend, 2),
        "proposed_total_weekly_spend": round(proposed_total_spend, 2),
        "spend_change_pct": round(
            (proposed_total_spend - current_total_spend) / current_total_spend * 100
            if current_total_spend > 0 else 0,
            1,
        ),
        "n_weeks": n_weeks,
        "projected_revenue": {
            "current_median": round(float(np.median(current_revenue_samples)), 2),
            "proposed_median": round(float(np.median(proposed_revenue_samples)), 2),
            "change_median": round(float(np.median(change_samples)), 2),
            "change_pct": round(
                float(np.median(change_samples / (current_revenue_samples + 1e-10) * 100)), 1
            ),
            "proposed_ci_lower": round(
                float(np.percentile(proposed_revenue_samples, (1 - ci) / 2 * 100)), 2
            ),
            "proposed_ci_upper": round(
                float(np.percentile(proposed_revenue_samples, (1 + ci) / 2 * 100)), 2
            ),
        },
        "channel_impacts": channel_impacts,
        "optimization_notes": optimization_notes,
    }


# ---------------------------------------------------------------------------
# Tool 8: get_weekly_contributions
# ---------------------------------------------------------------------------


@mcp.tool()
def get_weekly_contributions(
    channels: list[str] | None = None,
    last_n_weeks: int | None = None,
    start_week: str | None = None,
    end_week: str | None = None,
) -> dict:
    """Returns weekly data by channel over time — incremental contribution,
    spend, impressions, contribution %, and baseline revenue per week.

    Args:
        channels: Filter to specific channels. All if omitted.
        last_n_weeks: Return only the most recent N weeks. Overrides start/end.
        start_week: Start date (inclusive, ISO format e.g. '2025-06-01').
        end_week: End date (inclusive, ISO format e.g. '2025-10-27').
    """
    posterior = _get_posterior()
    model = _get_model()
    filtered = _filter_channels(channels)
    all_channels = _get_channels()

    ec = posterior["ec_m"]
    slope_var = posterior["slope_m"]
    alpha = posterior["alpha_m"]
    beta = posterior["beta_gm"]

    spend_np = model.media_tensors.media_spend.numpy()  # (1, n_time, n_channels)
    sf = np.squeeze(model.media_tensors.media_transformer._scale_factors_gm.numpy())
    kpi_stdev = float(model.kpi_transformer._population_scaled_stdev.numpy())

    # Extract impressions tensor if available
    impressions_np = None
    try:
        media_raw = model.media_tensors.media
        if hasattr(media_raw, "numpy"):
            impressions_np = media_raw.numpy()  # (1, n_time, n_channels)
        elif hasattr(media_raw, "values"):
            impressions_np = media_raw.values
    except (AttributeError, Exception):
        pass

    # Extract baseline (intercept) from posterior
    # Meridian stores the intercept as tau_g_excl (geo-level) or knot_values
    try:
        kpi_mean = float(model.kpi_transformer._population_scaled_mean.numpy())
        if "tau_g_excl" in posterior:
            tau = posterior["tau_g_excl"]
            if "geo" in tau.dims:
                tau = tau.isel(geo=0)
            baseline_per_draw = tau.values.flatten() * kpi_stdev + kpi_mean
            baseline_median = float(np.median(baseline_per_draw))
        elif "knot_values" in posterior:
            knot_vals = posterior["knot_values"]
            if "geo" in knot_vals.dims:
                knot_vals = knot_vals.isel(geo=0)
            baseline_per_draw = knot_vals.values.mean(axis=-1).flatten() * kpi_stdev + kpi_mean
            baseline_median = float(np.median(baseline_per_draw))
        else:
            baseline_median = kpi_mean
    except (AttributeError, Exception):
        baseline_median = None

    time_coords = list(posterior.coords["time"].values)
    time_labels = [str(t) for t in time_coords]
    n_time = len(time_labels)

    # Resolve time window
    if last_n_weeks is not None:
        t_start = max(0, n_time - last_n_weeks)
        t_end = n_time
    elif start_week is not None or end_week is not None:
        t_start = 0
        t_end = n_time
        if start_week is not None:
            for i, t in enumerate(time_labels):
                if t >= start_week:
                    t_start = i
                    break
        if end_week is not None:
            for i, t in enumerate(time_labels):
                if t > end_week:
                    t_end = i
                    break
    else:
        t_start = 0
        t_end = n_time

    week_indices = list(range(t_start, t_end))
    selected_labels = [time_labels[i] for i in week_indices]

    # Compute weekly contribution for ALL channels (needed for contribution %)
    all_contrib_weekly: dict[str, list[np.ndarray]] = {}
    for ch in all_channels:
        ch_idx = all_channels.index(ch)
        ec_vals = ec.sel(media_channel=ch).values.flatten()
        s_vals = slope_var.sel(media_channel=ch).values.flatten()
        a_vals = alpha.sel(media_channel=ch).values.flatten()
        b_vals = beta.sel(media_channel=ch)
        if "geo" in b_vals.dims:
            b_vals = b_vals.isel(geo=0)
        b_vals = b_vals.values.flatten()

        weekly_draws = []
        for wi in week_indices:
            week_spend = float(spend_np[0, wi, ch_idx])
            x = week_spend / (sf[ch_idx] + 1e-10)
            x_ad = x / (1 - a_vals + 1e-10)
            hill_response = x_ad**s_vals / (ec_vals**s_vals + x_ad**s_vals + 1e-10)
            weekly_draws.append(b_vals * hill_response * kpi_stdev)
        all_contrib_weekly[ch] = weekly_draws

    # Total incremental contribution across all channels per week per draw
    total_weekly_draws = [
        sum(all_contrib_weekly[ch][wi_idx] for ch in all_channels)
        for wi_idx in range(len(week_indices))
    ]

    # Build per-channel series for filtered channels
    channel_series: list[dict] = []
    for channel in filtered:
        ch_idx = all_channels.index(channel)
        weekly_draws = all_contrib_weekly[channel]

        weekly_medians = []
        weekly_ci_lower = []
        weekly_ci_upper = []
        weekly_spend = []
        weekly_impressions = []
        weekly_contrib_pct = []

        for wi_idx, wi in enumerate(week_indices):
            contrib = weekly_draws[wi_idx]
            weekly_medians.append(round(float(np.median(contrib)), 2))
            weekly_ci_lower.append(round(float(np.percentile(contrib, 5)), 2))
            weekly_ci_upper.append(round(float(np.percentile(contrib, 95)), 2))

            weekly_spend.append(round(float(spend_np[0, wi, ch_idx]), 2))

            if impressions_np is not None:
                weekly_impressions.append(round(float(impressions_np[0, wi, ch_idx]), 2))

            # Contribution % relative to total incremental
            pct_draws = contrib / (total_weekly_draws[wi_idx] + 1e-10) * 100
            weekly_contrib_pct.append(round(float(np.median(pct_draws)), 2))

        entry: dict[str, Any] = {
            "name": channel,
            "weeks": selected_labels,
            "contribution_median": weekly_medians,
            "contribution_ci_lower": weekly_ci_lower,
            "contribution_ci_upper": weekly_ci_upper,
            "contribution_pct": weekly_contrib_pct,
            "spend": weekly_spend,
            "total_median": round(sum(weekly_medians), 2),
            "total_spend": round(sum(weekly_spend), 2),
        }
        if weekly_impressions:
            entry["impressions"] = weekly_impressions
            entry["total_impressions"] = round(sum(weekly_impressions), 2)

        channel_series.append(entry)

    # Baseline series
    baseline_series = None
    if baseline_median is not None:
        baseline_series = {
            "weekly_baseline": round(baseline_median, 2),
            "total_baseline": round(baseline_median * len(week_indices), 2),
        }

    return {
        "channels": channel_series,
        "baseline": baseline_series,
        "n_weeks": len(week_indices),
        "time_range": {
            "start": selected_labels[0] if selected_labels else None,
            "end": selected_labels[-1] if selected_labels else None,
        },
    }


# ---------------------------------------------------------------------------
# Tool 9: get_channel_priors
# ---------------------------------------------------------------------------


def _extract_dist_params(dist: Any) -> dict[str, Any]:
    """Extract human-readable parameters from a TFP distribution."""
    dist_type = type(dist).__name__
    result: dict[str, Any] = {"distribution": dist_type}

    param_names: dict[str, list[str]] = {
        "LogNormal": ["loc", "scale"],
        "Normal": ["loc", "scale"],
        "HalfNormal": ["scale"],
        "TruncatedNormal": ["loc", "scale", "low", "high"],
        "Uniform": ["low", "high"],
        "Deterministic": ["loc"],
        "Beta": ["concentration1", "concentration0"],
        "Gamma": ["concentration", "rate"],
    }

    for param in param_names.get(dist_type, []):
        val = getattr(dist, param, None)
        if val is not None:
            try:
                import numpy as np

                arr = np.array(val)
                if arr.ndim == 0:
                    result[param] = round(float(arr), 6)
                else:
                    result[param] = [round(float(x), 6) for x in arr.flatten()]
            except Exception:
                result[param] = str(val)

    return result


@mcp.tool()
def get_channel_priors(
    channels: list[str] | None = None,
) -> dict:
    """Returns the prior distributions used for each channel in the model —
    ROI priors, adstock (alpha), half-saturation (EC), Hill slope, and beta
    effectiveness priors. Use this to verify modeling assumptions, check if
    informative priors were used, and answer questions about what the model
    assumed before seeing the data.

    Args:
        channels: Filter to specific channels. All if omitted.
    """
    model = _get_model()
    filtered = _filter_channels(channels)

    try:
        model_spec = model._model_context._model_spec
        prior = model_spec.prior
    except AttributeError:
        return {"error": "Prior distributions not available on this model."}

    media_prior_type = getattr(model_spec, "media_prior_type", None)
    media_effects_dist = getattr(model_spec, "media_effects_dist", None)

    # Extract prior distributions
    prior_attrs = {
        "roi_m": "ROI prior (used when media_prior_type='roi')",
        "alpha_m": "Adstock decay rate prior",
        "ec_m": "Half-saturation (EC) prior",
        "slope_m": "Hill slope prior",
        "beta_m": "Media effectiveness coefficient prior",
        "mroi_m": "Marginal ROI prior (used when media_prior_type='mroi')",
    }

    global_priors: dict[str, Any] = {}
    for attr, description in prior_attrs.items():
        dist = getattr(prior, attr, None)
        if dist is not None:
            params = _extract_dist_params(dist)
            params["description"] = description
            global_priors[attr] = params

    # Check if priors are per-channel (batch_shape matches n_channels)
    all_channels = _get_channels()
    n_channels = len(all_channels)
    per_channel = False

    # Build per-channel view
    channel_priors = []
    for channel in filtered:
        ch_idx = all_channels.index(channel)
        entry: dict[str, Any] = {"name": channel}

        for attr in prior_attrs:
            dist = getattr(prior, attr, None)
            if dist is None:
                continue
            try:
                import numpy as np

                batch = dist.batch_shape
                if len(batch) > 0 and batch[0] == n_channels:
                    per_channel = True
                    params = dist.parameters
                    ch_params: dict[str, Any] = {
                        "distribution": type(dist).__name__,
                    }
                    for k, v in params.items():
                        if k in (
                            "name", "validate_args",
                            "allow_nan_stats", "dtype",
                        ):
                            continue
                        arr = np.array(v)
                        if arr.ndim > 0 and arr.shape[0] == n_channels:
                            ch_params[k] = round(float(arr[ch_idx]), 6)
                        elif arr.ndim == 0:
                            ch_params[k] = round(float(arr), 6)
                    entry[attr] = ch_params
            except Exception:
                pass

        channel_priors.append(entry)

    result: dict[str, Any] = {
        "media_prior_type": str(media_prior_type),
        "media_effects_dist": str(media_effects_dist),
        "priors_are_per_channel": per_channel,
    }

    if per_channel:
        result["channels"] = channel_priors
    else:
        result["global_priors"] = global_priors
        result["applies_to_channels"] = filtered

    return result


# ---------------------------------------------------------------------------
# Tool 10: get_prior_posterior_comparison
# ---------------------------------------------------------------------------


def _prior_samples(dist: Any, n: int = 10000) -> np.ndarray | None:
    """Draw samples from a TFP-like prior distribution."""
    dist_type = type(dist).__name__
    rng = np.random.default_rng(42)

    try:
        params = dist.parameters
        if dist_type == "LogNormal":
            loc = float(np.array(params["loc"]))
            scale = float(np.array(params["scale"]))
            return rng.lognormal(loc, scale, n).astype(np.float64)
        elif dist_type == "Normal":
            loc = float(np.array(params["loc"]))
            scale = float(np.array(params["scale"]))
            return rng.normal(loc, scale, n).astype(np.float64)
        elif dist_type == "HalfNormal":
            scale = float(np.array(params["scale"]))
            return np.abs(rng.normal(0, scale, n)).astype(np.float64)
        elif dist_type == "Uniform":
            low = float(np.array(params["low"]))
            high = float(np.array(params["high"]))
            return rng.uniform(low, high, n).astype(np.float64)
        elif dist_type == "TruncatedNormal":
            loc = float(np.array(params["loc"]))
            scale = float(np.array(params["scale"]))
            low = float(np.array(params["low"]))
            high = float(np.array(params["high"]))
            samples = rng.normal(loc, scale, n * 5)
            samples = samples[(samples >= low) & (samples <= high)][:n]
            return samples.astype(np.float64)
        elif dist_type == "Deterministic":
            loc = float(np.array(params["loc"]))
            return np.full(n, loc, dtype=np.float64)
        elif dist_type == "Beta":
            a = float(np.array(params["concentration1"]))
            b = float(np.array(params["concentration0"]))
            return rng.beta(a, b, n).astype(np.float64)
    except Exception:
        pass
    return None


def _overlap_coefficient(
    prior_samples: np.ndarray, posterior_samples: np.ndarray, n_bins: int = 100,
) -> float:
    """Compute overlap coefficient between two distributions via histogram."""
    lo = min(prior_samples.min(), posterior_samples.min())
    hi = max(prior_samples.max(), posterior_samples.max())
    if hi - lo < 1e-10:
        return 1.0
    bins = np.linspace(lo, hi, n_bins + 1)
    prior_hist, _ = np.histogram(prior_samples, bins=bins, density=True)
    post_hist, _ = np.histogram(posterior_samples, bins=bins, density=True)
    bin_width = (hi - lo) / n_bins
    return float(np.minimum(prior_hist, post_hist).sum() * bin_width)


@mcp.tool()
def get_prior_posterior_comparison(
    channels: list[str] | None = None,
    credible_interval: float = 0.9,
) -> dict:
    """Compares prior assumptions to what the model actually learned from
    the data (posterior) for each channel. For each parameter, shows:

    - Prior and posterior summary statistics (median, CI)
    - Overlap coefficient (0–1): high overlap means the data didn't move
      the estimate much — the prior is doing the heavy lifting
    - A data_influence score: how much the posterior shifted from the prior
    - Flags channels where results may be prior-dominated

    Use this to answer: "Should I trust this estimate, or is it just
    reflecting what we assumed?"

    Args:
        channels: Filter to specific channels. All if omitted.
        credible_interval: Width of credible interval (default: 0.9).
    """
    model = _get_model()
    posterior = _get_posterior()
    filtered = _filter_channels(channels)
    ci = credible_interval

    try:
        prior = model._model_context._model_spec.prior
    except AttributeError:
        return {"error": "Prior distributions not available on this model."}

    # Parameters to compare: (prior_attr, posterior_var, label)
    comparisons = [
        ("roi_m", "roi_m", "ROI"),
        ("alpha_m", "alpha_m", "Adstock decay"),
        ("ec_m", "ec_m", "Half-saturation (EC)"),
        ("slope_m", "slope_m", "Hill slope"),
    ]

    channel_results = []
    for channel in filtered:
        params = []
        flags = []

        for prior_attr, post_var, label in comparisons:
            prior_dist = getattr(prior, prior_attr, None)
            if prior_dist is None or post_var not in posterior:
                continue

            post_vals = posterior[post_var].sel(
                media_channel=channel,
            ).values.flatten().astype(np.float64)
            prior_samp = _prior_samples(prior_dist)

            if prior_samp is None or len(prior_samp) < 100:
                continue

            # Posterior stats
            post_median = float(np.median(post_vals))
            post_mean = float(np.mean(post_vals))
            post_ci_lo = float(
                np.percentile(post_vals, (1 - ci) / 2 * 100),
            )
            post_ci_hi = float(
                np.percentile(post_vals, (1 + ci) / 2 * 100),
            )

            # Prior stats
            prior_median = float(np.median(prior_samp))
            prior_mean = float(np.mean(prior_samp))
            prior_ci_lo = float(
                np.percentile(prior_samp, (1 - ci) / 2 * 100),
            )
            prior_ci_hi = float(
                np.percentile(prior_samp, (1 + ci) / 2 * 100),
            )

            # Overlap and shift
            overlap = _overlap_coefficient(prior_samp, post_vals)
            prior_std = float(np.std(prior_samp))
            shift = (
                abs(post_median - prior_median) / (prior_std + 1e-10)
            )

            # Concentration: posterior narrower than prior = data informed
            prior_width = prior_ci_hi - prior_ci_lo
            post_width = post_ci_hi - post_ci_lo
            concentration = (
                1 - post_width / (prior_width + 1e-10)
                if prior_width > 0 else 0.0
            )

            # Data influence score: combines shift + concentration
            data_influence = min(1.0, (shift * 0.5 + max(0, concentration) * 0.5))

            if overlap > 0.7 and concentration < 0.3:
                influence_label = "prior_dominated"
                flags.append(
                    f"{label}: posterior closely matches the prior "
                    f"(overlap={overlap:.0%}). The data provided little "
                    f"information for this parameter."
                )
            elif overlap > 0.4:
                influence_label = "weakly_informed"
            elif concentration > 0.5:
                influence_label = "data_driven"
            else:
                influence_label = "moderately_informed"

            params.append({
                "parameter": prior_attr,
                "label": label,
                "prior": {
                    "distribution": type(prior_dist).__name__,
                    "median": round(prior_median, 4),
                    "mean": round(prior_mean, 4),
                    "ci_lower": round(prior_ci_lo, 4),
                    "ci_upper": round(prior_ci_hi, 4),
                },
                "posterior": {
                    "median": round(post_median, 4),
                    "mean": round(post_mean, 4),
                    "ci_lower": round(post_ci_lo, 4),
                    "ci_upper": round(post_ci_hi, 4),
                },
                "overlap_coefficient": round(overlap, 3),
                "prior_posterior_shift": round(shift, 3),
                "concentration": round(concentration, 3),
                "data_influence": round(data_influence, 3),
                "influence_label": influence_label,
            })

        channel_results.append({
            "name": channel,
            "parameters": params,
            "flags": flags if flags else None,
        })

    return {
        "channels": channel_results,
        "credible_interval_width": ci,
        "interpretation": {
            "overlap_coefficient": (
                "0–1 scale. High (>0.7) means prior and posterior are "
                "similar — data didn't override the assumption."
            ),
            "data_influence": (
                "0–1 scale. Low (<0.3) = prior-dominated, "
                "High (>0.7) = data-driven."
            ),
            "influence_labels": {
                "prior_dominated": "Posterior ≈ prior. Estimate reflects assumptions, not data.",
                "weakly_informed": "Some learning, but prior still has strong influence.",
                "moderately_informed": "Data moved the estimate meaningfully from the prior.",
                "data_driven": "Posterior is substantially narrower/shifted. Data is informative.",
            },
        },
    }


# ---------------------------------------------------------------------------
# Tool 11: list_models (S3)
# ---------------------------------------------------------------------------


@mcp.tool()
def list_models() -> dict:
    """Lists available Meridian model files from the configured S3 bucket.
    Returns model name, size, and last modified date for each .pkl file found.
    Also indicates which model is currently loaded (if any).
    Requires MERIDIAN_S3_BUCKET to be set."""
    bucket = os.environ.get("MERIDIAN_S3_BUCKET")
    if not bucket:
        return {
            "error": (
                "S3 is not configured. Set the MERIDIAN_S3_BUCKET environment variable "
                "to enable multi-model support."
            )
        }

    import boto3

    prefix = os.environ.get("MERIDIAN_S3_PREFIX", "")
    s3 = boto3.client("s3")

    paginator = s3.get_paginator("list_objects_v2")
    pages = paginator.paginate(Bucket=bucket, Prefix=prefix)

    models = []
    for page in pages:
        for obj in page.get("Contents", []):
            key = obj["Key"]
            if not key.endswith(".pkl"):
                continue
            name = os.path.basename(key)
            models.append({
                "name": name,
                "s3_key": key,
                "size_mb": round(obj["Size"] / (1024 * 1024), 2),
                "last_modified": obj["LastModified"].isoformat(),
                "is_loaded": name == _loaded_model_name,
            })

    return {
        "bucket": bucket,
        "prefix": prefix,
        "models": models,
        "count": len(models),
        "currently_loaded": _loaded_model_name,
    }


# ---------------------------------------------------------------------------
# Tool 12: load_model (S3)
# ---------------------------------------------------------------------------


@mcp.tool()
def load_model(name: str) -> dict:
    """Loads a Meridian model by name. If MERIDIAN_S3_BUCKET is configured, downloads
    the model from S3 to a local cache (skipping download if already cached), then
    loads it. If S3 is not configured, treats the name as a local file path.
    Replaces the currently loaded model and returns the model summary.

    Args:
        name: Model filename (e.g. 'client_20260301.pkl') or local path.
    """
    bucket = os.environ.get("MERIDIAN_S3_BUCKET")

    if bucket:
        prefix = os.environ.get("MERIDIAN_S3_PREFIX", "")
        s3_key = f"{prefix}{name}" if prefix and not name.startswith(prefix) else name
        filename = os.path.basename(s3_key)

        cache_dir = _get_cache_dir()
        local_path = cache_dir / filename

        if not local_path.exists():
            try:
                _download_from_s3(bucket, s3_key, local_path)
            except Exception as e:
                return {"error": f"Failed to download s3://{bucket}/{s3_key}: {e}"}

        _load_model(path=str(local_path))
    else:
        # Treat name as a local path
        _load_model(path=name)

    return get_model_summary()


# ---------------------------------------------------------------------------
# Tool 13: get_model_settings
# ---------------------------------------------------------------------------


@mcp.tool()
def get_model_settings() -> dict:
    """Returns detailed model training settings — MCMC configuration (chains,
    draws, warmup/burn-in), holdout/test set strategy, prior parameterization
    details, and sampler tuning parameters.

    Some settings (warmup, holdout, sampler config) are only available if they
    were stored as posterior attrs at training time. Missing settings are
    reported as "not_stored" with guidance on how to persist them.
    """
    trace = _get_trace()
    posterior = _get_posterior()
    model = _get_model()
    attrs = dict(posterior.attrs)

    # --- MCMC configuration ---
    mcmc: dict[str, Any] = {
        "n_chains": int(posterior.sizes["chain"]),
        "n_draws": int(posterior.sizes["draw"]),
        "n_warmup": attrs.get("n_warmup", "not_stored"),
        "target_accept_rate": attrs.get("target_accept_rate", "not_stored"),
    }

    # Realized sampler diagnostics from sample_stats
    try:
        sample_stats = trace.sample_stats
        mcmc["realized_accept_ratio"] = float(
            sample_stats["accept_ratio"].values.mean()
        )
        diverging = sample_stats["diverging"].values
        mcmc["n_divergences"] = int(diverging.sum())
    except (KeyError, AttributeError):
        pass

    # --- Holdout / test set ---
    holdout: dict[str, Any] = {
        "strategy": attrs.get("holdout_strategy", "not_stored"),
        "n_holdout_points": attrs.get("n_holdout_points", "not_stored"),
        "random_seed": attrs.get("holdout_seed", "not_stored"),
    }

    # --- Prior parameterization ---
    model_spec = None
    try:
        model_spec = model._model_context._model_spec
    except AttributeError:
        pass

    priors_info: dict[str, Any] = {}
    if model_spec is not None:
        priors_info["media_prior_type"] = getattr(model_spec, "media_prior_type", "unknown")
        priors_info["media_effects_dist"] = getattr(model_spec, "media_effects_dist", "unknown")

        prior = getattr(model_spec, "prior", None)
        if prior is not None:
            prior_key = "roi_m" if priors_info["media_prior_type"] == "roi" else "mroi_m"
            dist = getattr(prior, prior_key, None)
            if dist is not None:
                params = getattr(dist, "parameters", {})
                dist_name = type(dist).__name__
                priors_info["roi_prior_example"] = {
                    "distribution": dist_name,
                    "parameters": {
                        k: float(v) for k, v in params.items()
                        if k not in ("name", "validate_args", "allow_nan_stats", "dtype")
                        and hasattr(v, "__float__")
                    },
                    "note": (
                        "LogNormal(loc=μ, scale=σ) means log(X) ~ Normal(μ, σ). "
                        "No additional transformations are applied to loc/scale."
                        if dist_name == "LogNormal"
                        else f"Using {dist_name} distribution with native TFP parameterization."
                    ),
                }

    # --- Model structure ---
    structure: dict[str, Any] = {}
    if model_spec is not None:
        structure["max_lag"] = getattr(model_spec, "max_lag", "unknown")
        structure["hill_before_adstock"] = getattr(model_spec, "hill_before_adstock", "unknown")
        structure["adstock_type"] = "geometric"
        structure["saturation_type"] = "hill"

    # --- Guidance for missing settings ---
    not_stored_fields = []
    for section_name, section in [("mcmc", mcmc), ("holdout", holdout)]:
        for k, v in section.items():
            if v == "not_stored":
                not_stored_fields.append(f"{section_name}.{k}")

    guidance = None
    if not_stored_fields:
        guidance = (
            f"Fields {not_stored_fields} were not found in posterior.attrs. "
            "To persist these, add them at training time before pickling: "
            "trace.posterior.attrs['n_warmup'] = 500, "
            "trace.posterior.attrs['holdout_strategy'] = 'last_8_weeks', etc."
        )

    return {
        "mcmc": mcmc,
        "holdout": holdout,
        "prior_parameterization": priors_info,
        "model_structure": structure,
        "guidance_for_missing": guidance,
    }


# ---------------------------------------------------------------------------
# Entry point
# ---------------------------------------------------------------------------


def main():
    mcp.run()
