"""
Plotting helpers for the synthetic-data post.

Conventions used throughout (kept deliberately uniform across all figures):
- promotional weeks are shaded light grey in every time-series panel;
- channels keep fixed colours: PLA blue, Meta green, TV yellow
  (the Okabe-Ito palette, readable for colour-blind readers);
- series are labelled directly at their right-hand ends instead of legend boxes;
- axes carry only the ink they need: no top/right spines, at most a light
  horizontal grid, no dual axes.
"""

import numpy as np
import pandas as pd
import matplotlib.pyplot as plt

# Okabe-Ito palette
COLORS = {
    'pla': '#0072B2',        # blue
    'meta': '#009E73',       # green
    'tv': '#D55E00',         # yellow
    'sales': '#000000',
    'true': '#009E73',
    'observed': '#E69F00',   # orange
    'error': '#D55E00',
    'muted': '#808080',
}


# ---------------------------------------------------------------------------
# Shared styling utilities
# ---------------------------------------------------------------------------

def _tufte(ax, grid='y'):
    """Strip an axis down to the ink it needs."""
    for side in ('top', 'right'):
        ax.spines[side].set_visible(False)
    if grid == 'y':
        ax.grid(axis='y', color='0.92', linewidth=0.6)
        ax.set_axisbelow(True)
    else:
        ax.grid(False)
    ax.tick_params(length=3, labelsize=9)


def _shade_promotions(ax, promo_weeks):
    """Light grey bands for promotional weeks - the same encoding everywhere."""
    for w in promo_weeks:
        ax.axvspan(w - 0.5, w + 0.5, color='0.5', alpha=0.12, linewidth=0)


def _label_end(ax, x, y, text, color, fontsize=9, dx=2.0):
    """Direct label at the right-hand end of a series."""
    ax.annotate(text, xy=(x, y), xytext=(x + dx, y),
                va='center', ha='left', color=color, fontsize=fontsize,
                annotation_clip=False)


def _spread_labels(ax, x, entries, dx=2.0, fontsize=9):
    """Direct end labels with a minimum vertical gap, computed in display
    space so labels never collide regardless of panel height.
    entries: list of (y, text, color), any order."""
    ax.figure.canvas.draw()
    height_px = fontsize * 1.45 * ax.figure.dpi / 72.0
    inv = ax.transData.inverted()
    min_gap = abs(inv.transform((0, height_px))[1] - inv.transform((0, 0))[1])
    entries = sorted(entries, key=lambda e: e[0])
    ys = [e[0] for e in entries]
    for i in range(1, len(ys)):
        if ys[i] - ys[i - 1] < min_gap:
            ys[i] = ys[i - 1] + min_gap
    shift = (ys[-1] - entries[-1][0]) / 2 if len(ys) > 1 else 0.0
    ys = [y - shift for y in ys]
    for (y_orig, text, color), y in zip(entries, ys):
        ax.annotate(text, xy=(x, y_orig), xytext=(x + dx, y),
                    va='center', ha='left', color=color, fontsize=fontsize,
                    annotation_clip=False)


def _title(ax, text):
    """Left-aligned, regular-weight sentence titles."""
    ax.set_title(text, loc='left', fontsize=10.5, fontweight='normal', pad=8)


# ---------------------------------------------------------------------------
# Step 1: base sales components
# ---------------------------------------------------------------------------

def visualize_base_sales_components(base_data, show_summary=True):
    """
    Small multiples of the base-sales components on a shared time axis.
    """
    weeks = base_data['week'].values
    n_weeks = len(weeks)

    fig, axes = plt.subplots(5, 1, figsize=(9.5, 10.5), sharex=True,
                             gridspec_kw={'hspace': 0.55})

    # Panel 1: the composite series
    ax = axes[0]
    ax.plot(weeks, base_data['base_sales'] / 1000, color=COLORS['sales'], lw=1.6)
    _tufte(ax)
    _title(ax, 'Total base sales - the counterfactual without marketing (EUR millions)')
    peak_idx = base_data['base_sales'].idxmax()
    ax.annotate(f"seasonal peak, EUR {base_data['base_sales'].max()/1000:.1f}M",
                xy=(base_data.loc[peak_idx, 'week'], base_data['base_sales'].max() / 1000),
                xytext=(base_data.loc[peak_idx, 'week'] + 12,
                        base_data['base_sales'].max() / 1000),
                fontsize=8.5, color='0.35', va='center',
                arrowprops=dict(arrowstyle='-', color='0.6', lw=0.7))

    # Panel 2: seasonality
    ax = axes[1]
    ax.plot(weeks, base_data['seasonal_component'] * 100, color=COLORS['meta'], lw=1.3)
    ax.axhline(0, color='0.8', lw=0.7)
    _tufte(ax)
    _title(ax, 'Seasonality component - three harmonics, a function of the calendar alone (%)')
    for year_start in range(52, n_weeks, 52):
        ax.axvline(weeks[0] + year_start, color='0.85', lw=0.7, linestyle=':')

    # Panel 3: quality
    ax = axes[2]
    ax.plot(weeks, base_data['quality_factor'] * 100, color='#CC79A7', lw=1.3)
    ax.axhline(0, color='0.8', lw=0.7)
    _tufte(ax)
    _title(ax, 'Product quality component - slow mean-reverting drift, unobserved, multiplicative (%)')

    # Panel 4: competitive price positioning
    ax = axes[3]
    ax.plot(weeks, base_data['base_price_level'], color=COLORS['observed'], lw=1.3)
    ax.axhline(1.0, color='0.8', lw=0.7)
    ax.annotate('1.0 = market level', xy=(weeks[-1], 1.0), xytext=(weeks[-1] + 2, 1.0),
                va='center', fontsize=8.5, color='0.5', annotation_clip=False)
    _tufte(ax)
    _title(ax, 'Competitive price positioning component - partially observed later, base for promotions')

    # Panel 5: sentiment
    ax = axes[4]
    ax.plot(weeks, base_data['market_sentiment'] * 100, color='0.45', lw=1.1)
    ax.axhline(0, color='0.8', lw=0.7)
    _tufte(ax)
    _title(ax, 'Market sentiment component - faster-moving, unobserved, additive (%)')

    ax.set_xlabel('Week', fontsize=9.5)
    ax.set_xlim(weeks[0], weeks[-1] * 1.06)

    plt.show()

    if show_summary:
        print("BASE SALES SUMMARY")
        print("=" * 50)
        seasonal_range = base_data['seasonal_component'].max() - base_data['seasonal_component'].min()
        quality_range = base_data['quality_factor'].max() - base_data['quality_factor'].min()
        price_range = base_data['base_price_level'].max() - base_data['base_price_level'].min()
        print(f"Seasonal variation:     {seasonal_range:.1%} peak-to-trough")
        print(f"Quality evolution:      {quality_range:.1%} range")
        print(f"Price positioning:      {price_range:.1%} range")
        cv = base_data['base_sales'].std() / base_data['base_sales'].mean()
        print(f"Base sales:             EUR {base_data['base_sales'].mean():,.0f}K per week on average, "
              f"{cv:.1%} coefficient of variation")


# ---------------------------------------------------------------------------
# Step 2: promotional effects and measurement error
# ---------------------------------------------------------------------------

def visualize_promotional_effects(promotional_data_df, show_summary=True):
    """
    Three panels: sales with and without promotional uplift; true vs observed
    price level; and promotion detection, with disagreements marked.
    """
    df = promotional_data_df
    weeks = df['week'].values
    promo_weeks = df.loc[df['true_promotion_dummy'] == 1, 'week'].values

    fig, axes = plt.subplots(3, 1, figsize=(9.5, 8.5),
                             gridspec_kw={'hspace': 0.55, 'height_ratios': [1.4, 1.1, 0.7]})

    # Panel 1: sales impact
    ax = axes[0]
    _shade_promotions(ax, promo_weeks)
    ax.plot(weeks, df['base_sales'] / 1000, color='0.65', lw=1.1)
    ax.plot(weeks, df['true_sales_with_promotions'] / 1000, color=COLORS['sales'], lw=1.5)
    _spread_labels(ax, weeks[-1],
                   [(df['base_sales'].iloc[-1] / 1000, 'base', '0.5'),
                    (df['true_sales_with_promotions'].iloc[-1] / 1000,
                     'with promotions', COLORS['sales'])])
    _tufte(ax)
    _title(ax, 'Promotional uplift on top of base sales (EUR millions); promotion weeks shaded')
    ax.set_xlim(weeks[0], weeks[-1] * 1.14)

    # Panel 2: price levels, true vs observed
    ax = axes[1]
    _shade_promotions(ax, promo_weeks)
    ax.plot(weeks, df['true_promotional_price_level'], color=COLORS['true'], lw=1.4)
    ax.plot(weeks, df['observed_promotional_price_level'], color=COLORS['observed'],
            lw=0.9, alpha=0.85)
    _spread_labels(ax, weeks[-1],
                   [(df['true_promotional_price_level'].iloc[-1], 'true', COLORS['true']),
                    (df['observed_promotional_price_level'].iloc[-1],
                     'observed', COLORS['observed'])])
    _tufte(ax)
    _title(ax, 'Price level: true vs observed, with autocorrelated measurement error')
    ax.set_xlim(weeks[0], weeks[-1] * 1.14)

    # Panel 3: promotion detection, true vs observed
    # (event markers only - the ~130 zero-weeks carry no information worth ink)
    ax = axes[2]
    true_w = df.loc[df['true_promotion_dummy'] == 1, 'week'].values
    obs_w = df.loc[df['observed_promotion_dummy'] == 1, 'week'].values
    ax.scatter(true_w, np.full(len(true_w), 1.0), s=22, color='0.3', zorder=3)
    ax.scatter(obs_w, np.full(len(obs_w), 1.07), s=26, facecolors='none',
               edgecolors=COLORS['observed'], marker='^', linewidths=1.2, zorder=3)
    ax.axhline(0, color='0.88', lw=0.8)
    ax.set_ylim(-0.18, 1.30)
    ax.set_xlim(weeks[0], weeks[-1] * 1.14)
    _spread_labels(ax, weeks[-1], [(1.0, 'true', '0.35'),
                                   (1.07, 'observed', COLORS['observed'])])
    ax.set_yticks([0, 1])
    ax.set_yticklabels(['0', '1'], fontsize=9)
    ax.set_ylabel('Promotion flag', fontsize=9)
    _tufte(ax, grid=None)
    _title(ax, 'Promotion detection: true vs observed '
               '(70% of true weeks captured, 15% false-positive rate)')
    ax.set_xlabel('Week', fontsize=9.5)

    plt.show()

    if show_summary:
        print("PROMOTIONAL EFFECTS SUMMARY")
        print("=" * 50)
        uplift_max = df['true_promotional_uplift'].max()
        print(f"Largest weekly promotional uplift: EUR {uplift_max:,.0f}K")
        cyber = df[df['promotion_type'] == 'cyber_week']
        if len(cyber) > 0:
            mult = (cyber['true_sales_with_promotions'] / cyber['base_sales']).iloc[0]
            print(f"Cyber Week sales multiplier:       {mult:.2f}x base sales")
        accuracy = (df['observed_promotion_dummy'] == df['true_promotion_dummy']).mean()
        tp = ((df['observed_promotion_dummy'] == 1) & (df['true_promotion_dummy'] == 1)).sum()
        fp = ((df['observed_promotion_dummy'] == 1) & (df['true_promotion_dummy'] == 0)).sum()
        fn = ((df['observed_promotion_dummy'] == 0) & (df['true_promotion_dummy'] == 1)).sum()
        print(f"Weekly detection accuracy:         {accuracy:.1%}")
        if tp + fn > 0:
            print(f"True promotions captured:          {tp / (tp + fn):.0%} ({fn} missed)")
        print(f"False promotion flags:             {fp} weeks")


# ---------------------------------------------------------------------------
# Step 3: marketing spend patterns and endogeneity
# ---------------------------------------------------------------------------

def visualize_marketing_spend_patterns(marketing_spend, planned_pla, promotion_types,
                                       base_data, promo_data, quarterly_adj,
                                       n_weeks=156, show_summary=True):
    """
    Four panels: spend by channel; spend vs sales timing (indexed to a common
    scale); lead-lag correlations; and planned vs performance-chased PLA.
    """
    weeks = np.arange(1, n_weeks + 1)
    promo_weeks = [i + 1 for i, p in enumerate(promotion_types) if p != 'none']
    sales = np.asarray(promo_data['true_sales_with_promotions'])

    fig, axes = plt.subplots(4, 1, figsize=(9.5, 13.5),
                             gridspec_kw={'hspace': 0.5})

    # Panel 1: spend by channel
    ax = axes[0]
    _shade_promotions(ax, promo_weeks)
    for ch in ('pla', 'meta', 'tv'):
        ax.plot(weeks, marketing_spend[ch], color=COLORS[ch], lw=1.2)
    # quarterly budget shifts as a sparse row of small labels along the top
    max_spend = max(np.max(marketing_spend[c]) for c in marketing_spend)
    ax.set_ylim(top=max_spend * 1.22)
    q_adj = np.asarray(quarterly_adj)
    for q_start in range(26, n_weeks, 13):
        if abs(q_adj[q_start] - q_adj[q_start - 1]) > 0.02:
            ax.annotate(f"x{q_adj[q_start]:.2f}",
                        xy=(q_start + 1, max_spend * 1.18), fontsize=7.5,
                        color='0.5', ha='center', va='top')
            ax.axvline(q_start + 0.5, color='0.9', lw=0.6, ymax=0.96)
    ax.annotate('budget reviews', xy=(1, max_spend * 1.18), fontsize=7.5,
                color='0.5', ha='left', va='top')
    _spread_labels(ax, weeks[-1],
                   [(marketing_spend[ch][-1], ch.upper(), COLORS[ch])
                    for ch in ('pla', 'meta', 'tv')])
    tv = marketing_spend['tv']
    xmas_week = int(np.argmax(tv[:52])) + 1
    ax.annotate('TV Christmas burst, 15x base',
                xy=(xmas_week + 0.5, tv[xmas_week - 1] * 0.9),
                xytext=(xmas_week + 10, tv[xmas_week - 1] * 0.88),
                fontsize=8.5, color='0.35', va='center',
                arrowprops=dict(arrowstyle='-', color='0.6', lw=0.7))
    _tufte(ax)
    _title(ax, 'Weekly spend by channel (EUR thousands); quarterly budget shifts noted along the top')
    ax.set_xlim(weeks[0], weeks[-1] * 1.10)

    # Panel 2: timing of spend vs sales, indexed to mean = 100
    ax = axes[1]
    _shade_promotions(ax, promo_weeks)
    sales_idx = sales / sales.mean() * 100
    pla_idx = marketing_spend['pla'] / marketing_spend['pla'].mean() * 100
    meta_idx = marketing_spend['meta'] / marketing_spend['meta'].mean() * 100
    ax.plot(weeks, sales_idx, color=COLORS['sales'], lw=1.8)
    ax.plot(weeks, pla_idx, color=COLORS['pla'], lw=1.0, alpha=0.9)
    ax.plot(weeks, meta_idx, color=COLORS['meta'], lw=1.0, alpha=0.9)
    top = max(sales_idx.max(), pla_idx.max(), meta_idx.max())
    ax.set_ylim(top=top * 1.20)
    ax.set_xlim(weeks[0], weeks[-1] * 1.10)
    _spread_labels(ax, weeks[-1],
                   [(sales_idx[-1], 'sales', COLORS['sales']),
                    (pla_idx[-1], 'PLA', COLORS['pla']),
                    (meta_idx[-1], 'Meta', COLORS['meta'])])
    cyber_idx = next((i for i, p in enumerate(promotion_types) if p == 'cyber_week'), None)
    if cyber_idx is not None and cyber_idx > 3:
        pre = cyber_idx - 2
        ax.annotate('strategic: spend rises before the event',
                    xy=(pre + 1, meta_idx[pre]), xytext=(2, top * 1.14),
                    fontsize=8.5, color='0.35', va='top',
                    arrowprops=dict(arrowstyle='-', color='0.6', lw=0.7))
        post = cyber_idx + 2
        ax.annotate('algorithmic: spend follows realised sales',
                    xy=(post + 1, pla_idx[post]), xytext=(post + 13, top * 0.99),
                    fontsize=8.5, color='0.35', va='top',
                    arrowprops=dict(arrowstyle='-', color='0.6', lw=0.7))
    _tufte(ax)
    _title(ax, 'Timing: sales, PLA and Meta on a common index (mean = 100); TV omitted for scale')

    # Panel 3: lead-lag correlations
    ax = axes[2]
    lag_range = range(-4, 5)
    correlation_results = {}
    for ch in ('pla', 'meta', 'tv'):
        spend = np.asarray(marketing_spend[ch])
        corrs = []
        for lag in lag_range:
            if lag < 0:      # spend leads sales
                corrs.append(np.corrcoef(spend[:lag], sales[-lag:])[0, 1])
            elif lag > 0:    # sales lead spend
                corrs.append(np.corrcoef(spend[lag:], sales[:-lag])[0, 1])
            else:
                corrs.append(np.corrcoef(spend, sales)[0, 1])
        correlation_results[ch] = corrs
        ax.plot(list(lag_range), corrs, color=COLORS[ch], lw=1.4,
                marker='o', markersize=3.5)
    _spread_labels(ax, 4, [(correlation_results[ch][-1], ch.upper(), COLORS[ch])
                           for ch in ('pla', 'meta', 'tv')], dx=0.25)
    ax.axhline(0, color='0.85', lw=0.8)
    ax.axvline(0, color='0.85', lw=0.8)
    _tufte(ax)
    _title(ax, 'Correlation of spend with sales at different lags')
    ax.set_xlabel('Lag in weeks    (left: spend leads sales - strategic; '
                  'right: sales lead spend - algorithmic)', fontsize=9)
    ax.set_xlim(-4.4, 5.2)

    # Panel 4: planned vs performance-chased PLA (identical random draws)
    ax = axes[3]
    _shade_promotions(ax, promo_weeks)
    ax.plot(weeks, planned_pla, color='0.6', lw=1.1, linestyle='--')
    ax.plot(weeks, marketing_spend['pla'], color=COLORS['pla'], lw=1.4)
    _spread_labels(ax, weeks[-1],
                   [(planned_pla[-1], 'planned only', '0.5'),
                    (marketing_spend['pla'][-1], 'with chasing', COLORS['pla'])])
    _tufte(ax)
    _title(ax, 'PLA spend with and without the performance-chasing multiplier, '
               'same random draws')
    ax.set_xlabel('Week', fontsize=9.5)
    ax.set_xlim(weeks[0], weeks[-1] * 1.10)

    plt.show()

    if show_summary:
        planned_cv = np.std(planned_pla) / np.mean(planned_pla)
        final_cv = np.std(marketing_spend['pla']) / np.mean(marketing_spend['pla'])
        print("MARKETING SPEND SUMMARY")
        print("=" * 50)
        print(f"PLA planned only:   mean EUR {np.mean(planned_pla):,.0f}K, CV {planned_cv:.2f}")
        print(f"PLA with chasing:   mean EUR {np.mean(marketing_spend['pla']):,.0f}K, CV {final_cv:.2f}")
        print("\nLead-lag correlation peaks (spend vs sales):")
        lags = list(range(-4, 5))
        for ch in ('pla', 'meta', 'tv'):
            corrs = correlation_results[ch]
            neg_peak = max(range(0, 4), key=lambda i: corrs[i])
            pos_peak = max(range(5, 9), key=lambda i: corrs[i])
            print(f"  {ch.upper():5s} at negative lags (spend leads): {corrs[neg_peak]:+.2f} at {lags[neg_peak]:+d}; "
                  f"at positive lags (sales lead): {corrs[pos_peak]:+.2f} at {lags[pos_peak]:+d}")


# ---------------------------------------------------------------------------
# Step 4: transformations
# ---------------------------------------------------------------------------

def visualize_marketing_transformations(marketing_spend, adstocked_spend, saturated_spend,
                                        marketing_effects, transformation_details,
                                        CHANNEL_CONFIGS, show_details=True, n_weeks=156):
    """
    One panel per channel: raw spend, adstocked spend, and the resulting
    sales effect, all in EUR thousands on a single axis.
    """
    weeks = np.arange(1, n_weeks + 1)
    channels = list(marketing_effects.keys())

    fig, axes = plt.subplots(len(channels), 1, figsize=(9.5, 3.1 * len(channels)),
                             sharex=True, gridspec_kw={'hspace': 0.5})
    if len(channels) == 1:
        axes = [axes]

    for ax, ch in zip(axes, channels):
        color = COLORS.get(ch, '0.3')
        spend = np.asarray(marketing_spend[ch])
        adstocked = np.asarray(adstocked_spend[ch])
        effects = np.asarray(marketing_effects[ch])
        cfg = CHANNEL_CONFIGS[ch]

        ax.plot(weeks, spend, color=color, lw=0.8, linestyle=':', alpha=0.55)
        ax.plot(weeks, adstocked, color=color, lw=1.2, alpha=0.8)
        ax.plot(weeks, effects, color=color, lw=2.2)
        _spread_labels(ax, weeks[-1],
                       [(spend[-1], 'spend', color),
                        (adstocked[-1], 'adstocked', color),
                        (effects[-1], 'effect', color)], fontsize=8.5)

        roas = effects.mean() / spend.mean()
        avg_sat = np.asarray(saturated_spend[ch]).mean()
        ax.text(0.01, 0.97,
                f"carryover {cfg['adstock_rate']:.0%} - "
                f"avg saturation {avg_sat:.0%} - ROAS {roas:.1f}x",
                transform=ax.transAxes, va='top', ha='left', fontsize=8.5, color='0.35')
        _tufte(ax)
        _title(ax, f"{ch.upper()}: spend, adstocked spend, and sales effect (EUR thousands)")
        ax.set_xlim(weeks[0], weeks[-1] * 1.12)

    axes[-1].set_xlabel('Week', fontsize=9.5)
    plt.show()

    if show_details:
        print("TRANSFORMATION CHECKS")
        print("=" * 50)
        for ch, details in transformation_details.items():
            lam = CHANNEL_CONFIGS[ch]['saturation_lambda']
            spend = np.asarray(marketing_spend[ch])
            effects = np.asarray(marketing_effects[ch])
            roas = effects.mean() / spend.mean()
            print(f"\n{ch.upper()}:")
            print(f"  average spend           EUR {details['avg_spend']:,.0f}K per week")
            print(f"  half-saturation spend   EUR {np.log(3) / lam:,.0f}K "
                  f"(spend at which the curve reaches 50%)")
            print(f"  implied ROAS            {roas:.1f}x")
            curve = "  saturation at spend of "
            for s in (50, 100, 150, 200):
                sat = (1 - np.exp(-lam * s)) / (1 + np.exp(-lam * s))
                curve += f"EUR {s}K: {sat:.0%}   "
            print(curve)


# ---------------------------------------------------------------------------
# Step 5: the complete dataset
# ---------------------------------------------------------------------------

def visualize_complete_dataset(dataset, show_summary=True):
    """
    Four panels: the causal decomposition in levels; the same as shares of
    total sales; channel effects; and the measurement-error columns.
    """
    weeks = dataset['week'].values
    promo_weeks = dataset.loc[dataset['true_promotion_dummy'] == 1, 'week'].values

    fig, axes = plt.subplots(4, 1, figsize=(9.5, 13),
                             gridspec_kw={'hspace': 0.5,
                                          'height_ratios': [1.9, 0.8, 1.2, 1.1]})

    base = dataset['base_sales'].values / 1000
    promo = dataset['promotional_sales_uplift'].values / 1000
    marketing = dataset['marketing_sales_uplift'].values / 1000
    total = dataset['total_sales'].values / 1000

    # Panel 1: decomposition in levels
    ax = axes[0]
    _shade_promotions(ax, promo_weeks)
    ax.fill_between(weeks, 0, base, color='0.85', linewidth=0)
    ax.fill_between(weeks, base, base + promo, color=COLORS['observed'],
                    alpha=0.6, linewidth=0)
    ax.fill_between(weeks, base + promo, base + promo + marketing,
                    color=COLORS['meta'], alpha=0.55, linewidth=0)
    ax.plot(weeks, total, color=COLORS['sales'], lw=1.4)
    _spread_labels(ax, weeks[-1],
                   [(base[-1] / 2, 'base', '0.4'),
                    ((base + promo)[-1], 'promotions', COLORS['observed']),
                    (base[-1] + promo[-1] + marketing[-1] / 2, 'marketing', COLORS['meta']),
                    (total[-1], 'total', COLORS['sales'])])
    _tufte(ax)
    _title(ax, 'Sales decomposition: base + promotional uplift + marketing uplift '
               '(EUR millions)')
    ax.set_xlim(weeks[0], weeks[-1] * 1.13)

    # Panel 2: the same as shares
    ax = axes[1]
    share_base = base / total * 100
    share_promo = promo / total * 100
    share_mkt = marketing / total * 100
    ax.fill_between(weeks, 0, share_base, color='0.85', linewidth=0)
    ax.fill_between(weeks, share_base, share_base + share_promo,
                    color=COLORS['observed'], alpha=0.6, linewidth=0)
    ax.fill_between(weeks, share_base + share_promo, 100,
                    color=COLORS['meta'], alpha=0.55, linewidth=0)
    ax.set_ylim(0, 100)
    _tufte(ax, grid=None)
    _title(ax, 'The same decomposition as shares of total sales (%)')
    ax.set_xlim(weeks[0], weeks[-1] * 1.13)

    # Panel 3: channel effects
    ax = axes[2]
    _shade_promotions(ax, promo_weeks)
    end_vals = []
    for ch in ('pla', 'meta', 'tv'):
        eff = dataset[f'{ch}_effect'].values / 1000
        ax.plot(weeks, eff, color=COLORS[ch], lw=1.2)
        end_vals.append((eff[-1], ch.upper(), COLORS[ch]))
    _spread_labels(ax, weeks[-1], end_vals)
    _tufte(ax)
    _title(ax, 'True marketing effects by channel (EUR millions) - '
               'the ground truth the MMM will be graded on')
    ax.set_xlim(weeks[0], weeks[-1] * 1.13)

    # Panel 4: measurement error - noisy prices, plus true vs observed
    # promotion flags as event markers along the bottom
    ax = axes[3]
    _shade_promotions(ax, promo_weeks)
    ax.plot(weeks, dataset['true_promotional_price_level'], color=COLORS['true'], lw=1.3)
    ax.plot(weeks, dataset['observed_promotional_price_level'], color=COLORS['observed'],
            lw=0.9, alpha=0.85)
    true_w = dataset.loc[dataset['true_promotion_dummy'] == 1, 'week'].values
    obs_w = dataset.loc[dataset['observed_promotion_dummy'] == 1, 'week'].values
    ax.scatter(true_w, np.full(len(true_w), 0.36), s=16, color='0.35', zorder=3)
    ax.scatter(obs_w, np.full(len(obs_w), 0.44), s=20, facecolors='none',
               edgecolors=COLORS['observed'], marker='^', linewidths=1.1, zorder=3)
    ax.set_ylim(0.24, None)
    ax.set_xlim(weeks[0], weeks[-1] * 1.16)
    _spread_labels(ax, weeks[-1],
                   [(dataset['true_promotional_price_level'].iloc[-1], 'true price', COLORS['true']),
                    (dataset['observed_promotional_price_level'].iloc[-1], 'observed price',
                     COLORS['observed']),
                    (0.44, 'observed flag', COLORS['observed']),
                    (0.36, 'true flag', '0.35')])
    _tufte(ax)
    _title(ax, 'What the model will see: noisy prices and mislabelled promotion weeks')
    ax.set_xlabel('Week', fontsize=9.5)

    plt.show()

    if show_summary:
        print("COMPLETE DATASET SUMMARY")
        print("=" * 50)
        print(f"Shape: {dataset.shape[0]} weeks ({dataset.shape[0] / 52:.1f} years), "
              f"{dataset.shape[1]} columns")
        print(f"Total sales: EUR {dataset['total_sales'].min():,.0f}K to "
              f"EUR {dataset['total_sales'].max():,.0f}K per week")
        tot = dataset['total_sales'].sum()
        print("\nDecomposition of cumulative sales:")
        print(f"  base sales           {dataset['base_sales'].sum() / tot:.1%}")
        print(f"  promotional uplift   {dataset['promotional_sales_uplift'].sum() / tot:.1%}")
        print(f"  marketing uplift     {dataset['marketing_sales_uplift'].sum() / tot:.1%}")
        print("\nTrue ROAS by channel:")
        for ch in ('pla', 'meta', 'tv'):
            roas = dataset[f'{ch}_effect'].sum() / dataset[f'{ch}_spend'].sum()
            print(f"  {ch.upper():5s} {roas:.2f}x")
        acc = (dataset['true_promotion_dummy'] == dataset['observed_promotion_dummy']).mean()
        corr = dataset['true_promotional_price_level'].corr(
            dataset['observed_promotional_price_level'])
        print(f"\nPromotion detection accuracy: {acc:.1%}")
        print(f"Price measurement correlation: {corr:.3f}")


# ---------------------------------------------------------------------------
# Helper visualisations: adstock and saturation
# ---------------------------------------------------------------------------


def visualize_adstock_weights(rates, l_max=6, colors=None):
    """
    Normalised adstock weights for a set of retention rates: the share of a
    week's spend landing in the current and following weeks. The half-life
    ln(1/2) / ln(rate) marks where the weekly weight has halved.
    """
    colors = colors or [COLORS['pla'], COLORS['meta'], COLORS['tv']]
    fig, axes = plt.subplots(1, len(rates), figsize=(9.5, 3.0), sharey=True,
                             gridspec_kw={'wspace': 0.12})
    lags = np.arange(l_max)
    for ax, rate, color in zip(np.atleast_1d(axes), rates, colors):
        w = rate ** lags
        w = w / w.sum()
        ax.bar(lags, w, width=0.62, color=color)
        ax.annotate(f"{w[0]:.0%}", xy=(0, w[0]), xytext=(0, w[0] + 0.015),
                    ha='center', fontsize=8, color='0.35')
        half_life = np.log(0.5) / np.log(rate)
        ax.text(0.97, 0.95, f"half-life {half_life:.1f} wk",
                transform=ax.transAxes, va='top', ha='right',
                fontsize=8.5, color='0.35')
        _tufte(ax)
        _title(ax, f"retention {rate:.0%}")
        ax.set_xticks(lags)
        ax.set_xlabel('Weeks after spend', fontsize=9)
    axes[0].set_ylabel('Share of effect', fontsize=9.5)
    plt.show()


def visualize_saturation_curves(saturation_fn, lambdas, colors=None, max_spend=350):
    """
    Saturation curves for a set of lambda values, using the caller's own
    saturation function. Dots mark the half-saturation spend ln(3) / lambda:
    the (adstocked) spend delivering half of the maximum possible effect.
    """
    colors = colors or [COLORS['pla'], COLORS['meta'], COLORS['tv']]
    fig, ax = plt.subplots(figsize=(9.5, 4.6))
    x = np.linspace(0, max_spend, 500)

    entries = []
    for lam, color in zip(lambdas, colors):
        ax.plot(x, saturation_fn(x, lam), color=color, lw=1.5)
        x_half = np.log(3) / lam
        ax.vlines(x_half, 0, 0.5, color=color, lw=0.7, linestyle=':', alpha=0.7)
        ax.plot(x_half, 0.5, 'o', ms=4.5, color=color, zorder=3)
        ax.annotate(f"{x_half:.0f}", xy=(x_half, 0.015), ha='center',
                    va='bottom', fontsize=8, color=color)
        entries.append((saturation_fn(max_spend, lam), f"λ = {lam:g}", color))

    ax.axhline(0.5, color='0.85', lw=0.8, linestyle='--')
    ax.axhline(1.0, color='0.85', lw=0.6)
    ax.annotate('ceiling = 1', xy=(0, 1.015), fontsize=8, color='0.5', va='bottom')
    ax.set_ylim(0, 1.08)
    ax.set_xlim(0, max_spend * 1.10)
    _spread_labels(ax, max_spend, entries)
    _tufte(ax)
    _title(ax, 'Logistic saturation; dots mark the half-saturation spend '
               '$x_{1/2} = \\ln 3 / \\lambda$ (EUR thousands)')
    ax.set_xlabel('Adstocked spend (EUR thousands)', fontsize=9.5)
    ax.set_ylabel('Fraction of maximum effect', fontsize=9.5)
    plt.show()