"""
Plotting helpers for the estimation post.

Same conventions as the synthetic-data post, so the figures of the two
posts read as one series:
- channels keep fixed colours: PLA blue, Meta green, TV yellow
  (the Okabe-Ito palette, readable for colour-blind readers);
- in addition, three fixed roles: truth green, estimate yellow, fit blue;
- series are labelled directly instead of legend boxes; where direct labels
  would clutter, the encoding is stated in the title;
- axes carry only the ink they need: no top/right spines, at most a light
  horizontal grid, no dual axes.

The helpers own no model quantities: parameters, traces, and the saturation
function are passed in from the notebook, which remains the single source
of truth.
"""

import numpy as np
import pandas as pd
import matplotlib.pyplot as plt
import arviz as az
import matplotlib.transforms as mtransforms

# Okabe-Ito palette, plus the three roles of this post
COLORS = {
    'pla': '#0072B2',        # blue
    'meta': '#009E73',       # green
    'tv': '#D55E00',         # yellow
    'sales': '#000000',
    'truth': '#009E73',
    'estimate': '#D55E00',
    'fit': '#0072B2',
    'muted': '#808080',
}


# ---------------------------------------------------------------------------
# Shared styling utilities (identical to the synthetic-data helpers)
# ---------------------------------------------------------------------------

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


def _clean(channel):
    """'pla_spend' -> 'PLA'."""
    return channel.replace('_spend', '').upper()


# ---------------------------------------------------------------------------
# Step 3: ROAS estimates against the truth
# ---------------------------------------------------------------------------

def visualize_roas_comparison(roas, true_roas, media_columns):
    """Estimated vs true ROAS: dot-and-interval against grey truth bars,
    and the full posterior distributions. `roas` is the (chain, draw,
    channels) DataArray from roas_by_channel."""
    fig, axes = plt.subplots(1, 2, figsize=(9.5, 4.2),
                             gridspec_kw={'wspace': 0.28})

    # Panel 1: estimates vs truth
    ax = axes[0]
    x_pos = np.arange(len(media_columns))
    ax.bar(x_pos - 0.18, [true_roas[ch] for ch in media_columns], 0.36,
           color='0.85', linewidth=0)
    for x, ch in zip(x_pos, media_columns):
        key = ch.replace('_spend', '')
        d = roas.sel(channels=key).values.flatten()
        lo, hi = az.hdi(d, hdi_prob=0.90)
        m = d.mean()
        ax.errorbar(x + 0.18, m, yerr=[[m - lo], [hi - m]],
                    fmt='o', capsize=3, markersize=6, lw=1.2,
                    color=COLORS[key])
        ax.annotate(f"{true_roas[ch]:.1f}x", xy=(x - 0.18, true_roas[ch]),
                    xytext=(0, 3), textcoords='offset points',
                    ha='center', va='bottom', fontsize=8.5, color='0.35')
    ax.axhline(0, color='0.7', lw=0.8)
    ax.set_xticks(x_pos)
    ax.set_xticklabels([_clean(ch) for ch in media_columns], fontsize=9)
    ax.set_ylabel('ROAS', fontsize=9.5)
    _tufte(ax)
    _title(ax, 'True ROAS (grey bars) vs estimated\n(dots, 90% HDIs)')

    # Panel 2: full posterior distributions, channel names at the modes
    ax = axes[1]
    for ch in media_columns:
        key = ch.replace('_spend', '')
        dist = roas.sel(channels=key).values.flatten()
        counts, edges, _ = ax.hist(dist, bins=30, density=True,
                                   alpha=0.55, color=COLORS[key], linewidth=0)
        peak = np.argmax(counts)
        ax.annotate(_clean(ch),
                    xy=((edges[peak] + edges[peak + 1]) / 2, counts[peak]),
                    xytext=(0, 4), textcoords='offset points',
                    ha='center', va='bottom', fontsize=9, color=COLORS[key])
        ax.axvline(true_roas[ch], color=COLORS[key], linestyle='--', lw=1.4)
    ax.set_xlabel('ROAS', fontsize=9.5)
    ax.set_ylabel('Density', fontsize=9.5)
    ax.margins(y=0.12)
    _tufte(ax)
    _title(ax, 'Posterior ROAS distributions;\ndashed lines mark the truths')

    plt.show()


# ---------------------------------------------------------------------------
# Step 4: fit and decomposition
# ---------------------------------------------------------------------------

def visualize_fit_and_decomposition(trace, dates, observed_sales, true_uplift,
                                    y_max, show_summary=True):
    """
    Two panels sharing a time axis: observed vs fitted sales, and the
    estimated marketing contribution against the true uplift. The point of
    the pairing: a model can win the first panel and lose the second.
    """
    dates = pd.Series(pd.to_datetime(np.asarray(dates)))
    observed_sales = np.asarray(observed_sales)
    true_uplift = np.asarray(true_uplift)

    mu_post = trace.posterior['mu'].values
    mu_post = mu_post.reshape(-1, mu_post.shape[-1])
    mu_mean = mu_post.mean(axis=0) * y_max
    mu_lo = np.percentile(mu_post, 5, axis=0) * y_max
    mu_hi = np.percentile(mu_post, 95, axis=0) * y_max

    media_post = trace.posterior['media_contribution'].values
    media_post = media_post.reshape(-1, media_post.shape[-1])
    media_mean = media_post.mean(axis=0) * y_max

    label_dx = pd.Timedelta(weeks=2)
    label_room = pd.Timedelta(weeks=20)

    fig, axes = plt.subplots(2, 1, figsize=(9.5, 7.0), sharex=True,
                             gridspec_kw={'hspace': 0.45})

    # Panel 1: observed vs fitted sales
    ax = axes[0]
    ax.fill_between(dates, mu_lo, mu_hi, color=COLORS['fit'], alpha=0.18,
                    linewidth=0)
    ax.plot(dates, observed_sales, color=COLORS['sales'], lw=1.0)
    ax.plot(dates, mu_mean, color=COLORS['fit'], lw=1.5)
    _spread_labels(ax, dates.iloc[-1],
                   [(observed_sales[-1], 'observed', COLORS['sales']),
                    (mu_mean[-1], 'fitted', COLORS['fit'])],
                   dx=label_dx)
    ax.set_ylabel('Sales (EUR K)', fontsize=9.5)
    _tufte(ax)
    _title(ax, 'The model fits the sales series well '
               '(posterior mean and 90% interval)...')
    ax.set_xlim(dates.iloc[0], dates.iloc[-1] + label_room)

    # Panel 2: estimated vs true marketing contribution
    ax = axes[1]
    ax.plot(dates, true_uplift, color=COLORS['truth'], lw=1.5)
    ax.plot(dates, media_mean, color=COLORS['estimate'], lw=1.5)
    _spread_labels(ax, dates.iloc[-1],
                   [(true_uplift[-1], 'true uplift', COLORS['truth']),
                    (media_mean[-1], 'estimated', COLORS['estimate'])],
                   dx=label_dx)
    ax.set_ylabel('Incremental sales (EUR K)', fontsize=9.5)
    ax.set_xlabel('Date', fontsize=9.5)
    _tufte(ax)
    _title(ax, '...while attributing too much of it to marketing')
    ax.set_xlim(dates.iloc[0], dates.iloc[-1] + label_room)

    plt.show()

    if show_summary:
        total = observed_sales.sum()
        print(f"True marketing share of sales:      {true_uplift.sum() / total:.1%}")
        print(f"Estimated marketing share of sales: {media_mean.sum() / total:.1%}")


# ---------------------------------------------------------------------------
# Step 4: response curves
# ---------------------------------------------------------------------------

def visualize_response_curves(trace, media_columns, spend_data, scale_factors,
                              true_configs, saturation_fn, n_lags=6,
                              n_draws=200, seed=1):
    """
    Steady-state response curves per channel: posterior draws (spaghetti),
    the posterior mean, and the true curve. Along the x-axis, a rug of the
    channel's actual adstocked spends (computed at the posterior-mean carryover
    rate) shows where the channel operated — and hence which part of the curve
    the data could constrain. `saturation_fn` is the notebook's own saturation
    function, passed in so the definition lives in exactly one place.
    """
    y_max = scale_factors['total_sales']
    post = trace.posterior
    total_draws = post.sizes['chain'] * post.sizes['draw']
    n_channels = len(media_columns)
    idx = np.random.default_rng(seed).choice(
        total_draws, min(n_draws, total_draws), replace=False)
    beta_s = post['beta_media'].values.reshape(-1, n_channels)
    lam_s = post['saturation_lambda'].values.reshape(-1, n_channels)
    alpha_mean = post['adstock_alpha'].values.reshape(-1, n_channels).mean(axis=0)

    fig, axes = plt.subplots(1, n_channels, figsize=(9.5, 3.5),
                             gridspec_kw={'wspace': 0.25})

    for ch_idx, (channel, ax) in enumerate(zip(media_columns, np.atleast_1d(axes))):
        cfg = true_configs[channel]
        x_max = scale_factors[channel]
        spend = np.asarray(spend_data[channel])

        # actual adstocked spend this channel saw, at posterior-mean alpha (EUR K)
        w = alpha_mean[ch_idx] ** np.arange(n_lags)
        w = w / w.sum()
        adstocked = np.convolve(spend, w)[:len(spend)]

        x_grid = np.linspace(0, adstocked.max() * 1.1, 200)     # EUR K, adstock scale
        x_scaled = x_grid / x_max

        draws = np.array([beta_s[i, ch_idx]
                          * saturation_fn(x_scaled, lam_s[i, ch_idx]) * y_max
                          for i in idx])
        for resp in draws:
            ax.plot(x_grid, resp, color=COLORS['fit'], alpha=0.03, lw=0.8)
        resp_mean = draws.mean(axis=0)
        ax.plot(x_grid, resp_mean, color=COLORS['fit'], lw=1.8)

        true_resp = cfg['beta'] * saturation_fn(x_grid, cfg['lambda'])
        ax.plot(x_grid, true_resp, color=COLORS['truth'], lw=1.6, linestyle='--')

        # rug of actual adstocked spends, pinned to the panel floor (x data, y axes)
        rug_tf = mtransforms.blended_transform_factory(ax.transData, ax.transAxes)
        ax.plot(adstocked, np.zeros_like(adstocked), marker='|', linestyle='none',
                color='0.55', markersize=6, markeredgewidth=0.8, alpha=0.5,
                transform=rug_tf, clip_on=False)

        if ch_idx == 0:
            _spread_labels(ax, x_grid[-1],
                           [(resp_mean[-1], 'estimated', COLORS['fit']),
                            (true_resp[-1], 'true', COLORS['truth'])],
                           dx=x_grid[-1] * 0.02, fontsize=8.5)
            ax.set_xlim(0, x_grid[-1] * 1.30)
            ax.annotate('adstocked spend', xy=(adstocked.mean(), 0),
                        xycoords=rug_tf, xytext=(4, 3), textcoords='offset points',
                        fontsize=7.5, color='0.5', va='bottom')
            ax.set_ylabel('Weekly incremental sales (EUR K)', fontsize=9.5)
        else:
            ax.set_xlim(0, x_grid[-1] * 1.10)
        _tufte(ax)
        _title(ax, _clean(channel))
        ax.set_xlabel('Adstocked weekly spend (EUR K)', fontsize=9.5)

    plt.show()


# ---------------------------------------------------------------------------
# Step 6: the three specifications side by side
# ---------------------------------------------------------------------------

def visualize_specification_ladder(true_roas, results_by_spec, media_columns):
    """
    One panel per channel: the estimated ROAS (dot, 90% HDI) under each
    specification, against the true value (dashed line). `results_by_spec`
    is an ordered dict mapping specification labels to the DataArrays
    returned by roas_by_channel.
    """
    spec_names = list(results_by_spec.keys())
    x_pos = np.arange(len(spec_names))
    n_channels = len(media_columns)

    fig, axes = plt.subplots(1, n_channels, figsize=(9.5, 3.5),
                             gridspec_kw={'wspace': 0.12},
                             sharey=True)
    
    for channel, ax in zip(media_columns, np.atleast_1d(axes)):
        key = channel.replace('_spend', '')
        color = COLORS.get(key, '0.3')
        for x, s in zip(x_pos, spec_names):
            d = results_by_spec[s].sel(channels=key).values.flatten()
            lo, hi = az.hdi(d, hdi_prob=0.90)
            m = d.mean()
            ax.errorbar(x, m, yerr=[[m - lo], [hi - m]],
                        fmt='o', capsize=3, markersize=6, lw=1.2, color=color)
        ax.axhline(true_roas[channel], color='0.4', lw=1.0, linestyle='--')
        ax.annotate(f"true {true_roas[channel]:.1f}x",
                    xy=(x_pos[-1], true_roas[channel]),
                    xytext=(0, 4), textcoords='offset points',
                    ha='right', va='bottom', fontsize=8.5, color='0.4')
        ax.axhline(0, color='0.85', lw=0.7)
        ax.set_xticks(x_pos)
        ax.set_xticklabels(spec_names, fontsize=8.5)
        ax.set_xlim(-0.5, len(spec_names) - 0.5)
        _tufte(ax)
        _title(ax, _clean(channel))

    np.atleast_1d(axes)[0].set_ylabel('ROAS (90% HDI)', fontsize=9.5)
    plt.show()