import logging
import os
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


def _load_model() -> None:
    global _model, _trace, _posterior
    if _model is not None:
        return

    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."
        )
    if not os.path.exists(model_path):
        raise RuntimeError(f"Model file not found: {model_path}")

    from meridian.model import model as meridian_model

    _model = meridian_model.load_mmm(model_path)
    _trace = _model._inference_data
    _posterior = _trace.posterior


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(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."""
    posterior = _get_posterior()
    filtered = _filter_channels(channels)
    ci = credible_interval

    roi = posterior["roi_m"]
    results = []

    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),
        })

    return {
        "channels": results,
        "credible_interval_width": ci,
        "n_samples": int(roi.sizes["draw"] * roi.sizes["chain"]),
    }


# ---------------------------------------------------------------------------
# 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."""
    posterior = _get_posterior()
    model = _get_model()
    filtered = _filter_channels(channels)
    all_channels = _get_channels()
    ci = credible_interval

    roi = posterior["roi_m"]
    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()  # (1, 104, 10)
    sf = model.media_tensors.media_transformer._scale_factors_gm.numpy()  # (1, 10)
    sf = np.squeeze(sf)

    results = []
    for channel in filtered:
        ch_idx = all_channels.index(channel)
        roi_vals = roi.sel(media_channel=channel).values.flatten()
        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)
        # Adstock steady-state
        x_ad = x / (1 - a_vals + 1e-10)

        # Hill derivative: d/dx [x^s / (ec^s + x^s)] = s * ec^s * x^(s-1) / (ec^s + x^s)^2
        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)

        # Chain rule: d(contribution)/d(spend) = beta * d_hill * (1/(1-alpha)) * (1/scale)
        # Scale from normalized KPI space to revenue using kpi_stdev
        kpi_stdev = float(model.kpi_transformer._population_scaled_stdev.numpy())
        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,
        "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),
        })

        if new_spend < cur_spend and cur_contrib_med > 0 and new_contrib_med > cur_contrib_med * 0.9:
            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 contribution by channel over time — a time series of how
    much each channel contributed to the KPI each week. Use this to see trends,
    recent performance, or compare specific time windows.

    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())

    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]

    # For each channel, compute weekly contribution (median across posterior)
    channel_series: list[dict] = []
    for channel in filtered:
        ch_idx = all_channels.index(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()

        weekly_medians = []
        weekly_ci_lower = []
        weekly_ci_upper = []

        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)
            contrib = b_vals * hill_response * kpi_stdev

            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))

        channel_series.append({
            "name": channel,
            "weeks": selected_labels,
            "contribution_median": weekly_medians,
            "contribution_ci_lower": weekly_ci_lower,
            "contribution_ci_upper": weekly_ci_upper,
            "total_median": round(sum(weekly_medians), 2),
        })

    return {
        "channels": channel_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,
        },
    }


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


def main():
    mcp.run()
