def align_to(ref, other):
    """Permutation `perm` such that other[perm] best matches ref (rows are topics)."""
    cost = ((ref[:, None, :] - other[None, :, :]) ** 2).sum(-1)
    return linear_sum_assignment(cost)[1]


def align_chains(phi_draws):
    """phi_draws: (chain, draw, topic, word). Returns one permutation per chain (chain 0 = reference)."""
    means = phi_draws.mean(1)
    return np.array([align_to(means[0], means[c]) for c in range(len(means))])


phi_s = idata_sim.posterior["phi"].values
perms_sim = align_chains(phi_s)
print("topic permutation per chain:\n", perms_sim)
phi_aligned = np.stack([phi_s[c][:, perms_sim[c]] for c in range(phi_s.shape[0])])
theta_s = idata_sim.posterior["theta"].values
theta_aligned = np.stack([theta_s[c][:, :, perms_sim[c]] for c in range(theta_s.shape[0])])
rhat_aligned = az.rhat(np.asarray(phi_aligned).reshape(phi_aligned.shape[0], phi_aligned.shape[1], -1))
print("largest r_hat of phi after alignment:", round(float(np.max(rhat_aligned)), 3))

# match the fitted topics to the true ones for the plot
to_true = align_to(phi_true, phi_aligned.mean((0, 1)))
phi_fit = phi_aligned.mean((0, 1))[to_true]
theta_fit = theta_aligned.mean((0, 1))[:, to_true]

fig, axes = plt.subplots(1, 2, figsize=(11, 3.6))
for k in range(K_sim):
    axes[0].plot(phi_true[k], "o-", color=f"C{k}", lw=1, label=f"topic {k}, true")
    axes[0].plot(phi_fit[k], "x", color=f"C{k}", ms=8)
axes[0].set(xlabel="word", ylabel="probability within topic", title="Topics: true (o) and posterior mean (x)")
axes[0].legend(fontsize=8)
axes[1].scatter(theta_true.ravel(), theta_fit.ravel(), s=6, alpha=0.5)
axes[1].plot([0, 1], [0, 1], color="k", lw=0.8)
axes[1].set(xlabel="true topic share", ylabel="posterior mean share", title="Document mixtures (200 docs x 3 topics)");

data.describe("upworthy_exploratory")
uw = data.load("upworthy_exploratory")

STOP = set("""a about above after again against all am an and any are aren't as at be because been before
being below between both but by can can't cannot could couldn't did didn't do does doesn't doing don't down
during each few for from further had hadn't has hasn't have haven't having he he'd he'll he's her here here's
hers herself him himself his how how's i i'd i'll i'm i've if in into is isn't it it's its itself let's me
more most mustn't my myself no nor not of off on once only or other ought our ours ourselves out over own same
shan't she she'd she'll she's should shouldn't so some such than that that's the their theirs them themselves
then there there's these they they'd they'll they're they've this those through to too under until up very
was wasn't we we'd we'll we're we've were weren't what what's when when's where where's which while who who's
whom why why's with won't would wouldn't you you'd you'll you're you've your yours yourself yourselves just
like get gets got one two really thing things way ever even also make makes made know see go going goes say
says said will new something every much many want wants yes""".split())


def tokens(headline):
    words = re.findall(r"[a-z][a-z']+", headline.lower().replace("’", "'"))
    return [w for w in words if w not in STOP and len(w) > 2]


heads = uw.drop_duplicates(["test_id", "headline"])
docs = heads.groupby("test_id")["headline"].apply(lambda h: sum((tokens(x) for x in h), []))
doc_freq = Counter(w for d in docs for w in set(d))
fold = {w: w[:-1] for w in doc_freq if w.endswith("s") and doc_freq.get(w[:-1], 0) >= 20}
docs = docs.map(lambda d: [fold.get(w, w) for w in d])
doc_freq = Counter(w for d in docs for w in set(d))
vocab = np.array(sorted(w for w, c in doc_freq.items() if 20 <= c <= 0.1 * len(docs)))
w_index = {w: i for i, w in enumerate(vocab)}
X_all = np.zeros((len(docs), len(vocab)), dtype=np.int64)
for d, ws in enumerate(docs):
    for w in ws:
        if w in w_index:
            X_all[d, w_index[w]] += 1
keep = X_all.sum(1) >= 3
X_txt, test_ids = X_all[keep], docs.index.to_numpy()[keep]
print(f"{len(docs)} tests, vocabulary of {len(vocab)} words; {keep.sum()} tests keep >= 3 tokens")
print(f"tokens per test: median {np.median(X_txt.sum(1)):.0f}, total {X_txt.sum():,}")
print("most common kept words:", ", ".join(vocab[np.argsort(-X_txt.sum(0))[:25]]))

@numba.njit
def _gibbs_sweeps(doc, word, z, ndk, nkw, nk, alpha, beta, n_sweeps, seed):
    np.random.seed(seed)
    K, V = nkw.shape
    cum = np.empty(K)
    for _ in range(n_sweeps):
        for i in range(doc.shape[0]):
            d, w, k = doc[i], word[i], z[i]
            ndk[d, k] -= 1
            nkw[k, w] -= 1
            nk[k] -= 1
            total = 0.0
            for j in range(K):
                total += (ndk[d, j] + alpha) * (nkw[j, w] + beta) / (nk[j] + V * beta)
                cum[j] = total
            u = np.random.random() * total
            k = 0
            while cum[k] < u:
                k += 1
            z[i] = k
            ndk[d, k] += 1
            nkw[k, w] += 1
            nk[k] += 1


def lda_gibbs(X, K, alpha=0.1, beta=0.05, burn=2000, draws=40, thin=10, seed=0):
    """Collapsed Gibbs LDA on a document x word count matrix.

    Returns posterior draws of phi (draw, topic, word) and theta (draw, doc, topic), and the
    log-likelihood of the tokens at each saved sweep.
    """
    g = np.random.default_rng(seed)
    d_i, w_i = np.nonzero(X)
    n = X[d_i, w_i]
    doc, word = np.repeat(d_i, n), np.repeat(w_i, n)
    z = g.integers(K, size=doc.size)
    D, V = X.shape
    ndk, nkw = np.zeros((D, K), np.int64), np.zeros((K, V), np.int64)
    np.add.at(ndk, (doc, z), 1)
    np.add.at(nkw, (z, word), 1)
    nk = nkw.sum(1)
    _gibbs_sweeps(doc, word, z, ndk, nkw, nk, alpha, beta, burn, seed)
    phis, thetas, lls = [], [], []
    for s in range(draws):
        _gibbs_sweeps(doc, word, z, ndk, nkw, nk, alpha, beta, thin, seed + 1 + s)
        phi_s = g.gamma(nkw + beta)
        theta_s = g.gamma(ndk + alpha)
        phis.append(phi_s / phi_s.sum(1, keepdims=True))
        thetas.append(theta_s / theta_s.sum(1, keepdims=True))
        lls.append(np.sum(X * np.log(thetas[-1] @ phis[-1])))
    return np.array(phis), np.array(thetas), np.array(lls)


K_TXT = 6
t0 = time.perf_counter()
fits = [lda_gibbs(X_txt, K_TXT, seed=c) for c in range(4)]
print(f"4 chains x 2,400 sweeps: {time.perf_counter() - t0:.1f} s (includes Numba compilation)")
phi_txt = np.stack([f[0] for f in fits])       # chain, draw, topic, word
theta_txt = np.stack([f[1] for f in fits])     # chain, draw, doc, topic
ll_txt = np.stack([f[2] for f in fits])
print("token log-likelihood per chain (mean over saved sweeps):", np.round(ll_txt.mean(1)))

def headline_tokens(word):
    toks = tokens(word)
    return fold.get(toks[0], toks[0]) if toks else None


two_topic = np.flatnonzero(np.sort(theta_ref, 1)[:, -2] > 0.3)
rows = np.random.default_rng(8).choice(two_topic, 9, replace=False)
fig, ax = plt.subplots(figsize=(13, 0.55 * len(rows) + 1.2))
ax.set(xlim=(0, 1), ylim=(len(rows) - 0.4, -1.2))
ax.axis("off")
renderer = fig.canvas.get_renderer()
to_data = ax.transData.inverted()
for r, d in enumerate(rows):
    left = 0.0
    for k in range(K_TXT):                                        # the story's mixture
        ax.barh(r, theta_ref[d, k] * 0.12, left=left, height=0.5, color=TOPIC_COLOURS[k],
                edgecolor="white", lw=1)
        left += theta_ref[d, k] * 0.12
    x = 0.14
    for word in first_headline[test_ids[d]].split():
        t = ax.text(x, r, word, va="center", fontsize=10, color=INK)
        w = headline_tokens(word)
        if w in w_index:
            k_w = int(np.argmax(theta_ref[d] * phi_ref[:, w_index[w]]))
            t.set_bbox(dict(facecolor=TOPIC_COLOURS[k_w], alpha=0.35, edgecolor="none", pad=1.5))
        else:
            t.set_color(MUTED)
        x = to_data.transform((t.get_window_extent(renderer).x1, 0))[0] + 0.005
x = 0.14
for k in range(K_TXT):                                            # legend row, laid out by measured width
    t = ax.text(x, -0.95, f"topic {k}: {labels_rel[k]}", fontsize=8, color=INK, va="center",
                bbox=dict(facecolor=TOPIC_COLOURS[k], alpha=0.35, edgecolor="none", pad=1.5))
    x = to_data.transform((t.get_window_extent(renderer).x1, 0))[0] + 0.012
ax.text(0, -0.95, "topic mix", fontsize=8, color=MUTED, va="center");

@numba.njit
def _foldin_sweeps(doc, word, z, ndk, phi, alpha, n_sweeps, seed):
    np.random.seed(seed)
    K = phi.shape[0]
    cum = np.empty(K)
    for _ in range(n_sweeps):
        for i in range(doc.shape[0]):
            d, w, k = doc[i], word[i], z[i]
            ndk[d, k] -= 1
            total = 0.0
            for j in range(K):
                total += (ndk[d, j] + alpha) * phi[j, w]
                cum[j] = total
            u = np.random.random() * total
            k = 0
            while cum[k] < u:
                k += 1
            z[i] = k
            ndk[d, k] += 1


def completion_score(X_seen, X_hidden, phis, alpha=0.1, sweeps=50, seed=0):
    """Log-probability per hidden token, theta inferred from the seen half, averaged over phi draws."""
    g = np.random.default_rng(seed)
    d_i, w_i = np.nonzero(X_seen)
    n = X_seen[d_i, w_i]
    doc, word = np.repeat(d_i, n), np.repeat(w_i, n)
    K = phis.shape[1]
    prob = np.zeros(X_hidden.shape)
    for s, phi_s in enumerate(phis):
        z = g.integers(K, size=doc.size)
        ndk = np.zeros((X_seen.shape[0], K), np.int64)
        np.add.at(ndk, (doc, z), 1)
        _foldin_sweeps(doc, word, z, ndk, phi_s, alpha, sweeps, seed + s)
        prob += (ndk + alpha) / (ndk.sum(1, keepdims=True) + K * alpha) @ phi_s
    return np.sum(X_hidden * np.log(prob / len(phis))) / X_hidden.sum()


split_rng = np.random.default_rng(1)
is_test = split_rng.random(len(X_txt)) < 0.2
X_train, X_test = X_txt[~is_test], X_txt[is_test]
X_seen = np.zeros_like(X_test)
for d in range(len(X_test)):
    toks = np.repeat(np.arange(len(vocab)), X_test[d])
    split_rng.shuffle(toks)
    np.add.at(X_seen[d], toks[: len(toks) // 2], 1)
X_hidden = X_test - X_seen
unigram = (X_train.sum(0) + 0.5) / (X_train.sum() + 0.5 * len(vocab))
print(f"one-topic baseline: {np.sum(X_hidden * np.log(unigram)) / X_hidden.sum():.3f} nats per token")

tests = uw.groupby("test_id").agg(clicks=("clicks", "sum"), impressions=("impressions", "sum"),
                                  week=("test_week", "first")).loc[test_ids]
logit_ctr = np.log((tests["clicks"] + 0.5) / (tests["impressions"] - tests["clicks"] + 0.5)).to_numpy()
year, wk = tests["week"].to_numpy() // 100, tests["week"].to_numpy() % 100
month = (year - year.min()) * 12 + np.minimum((wk - 1) * 12 // 52, 11)
month_idx, month_codes = pd.factorize(month, sort=True)
y_rel = logit_ctr - pd.Series(logit_ctr).groupby(month_idx).transform("mean").to_numpy()
cv_fold = np.random.default_rng(2).integers(5, size=len(y_rel))


def cv_r2(F, lam=1.0):
    pred = np.zeros_like(y_rel)
    for f in range(5):
        tr = cv_fold != f
        A = np.c_[np.ones(tr.sum()), F[tr]]
        b = np.linalg.solve(A.T @ A + lam * np.eye(A.shape[1]), A.T @ y_rel[tr])
        pred[~tr] = np.c_[np.ones((~tr).sum()), F[~tr]] @ b
    return 1 - np.mean((y_rel - pred) ** 2) / np.var(y_rel)


rows = []
for K in [2, 4, 6, 10, 20, 40, 80]:
    phi_tr, _, _ = lda_gibbs(X_train, K, burn=500, draws=10, thin=20, seed=K)
    _, theta_all, _ = lda_gibbs(X_txt, K, burn=500, draws=10, thin=20, seed=100 + K)
    rows.append({"K": K, "held-out nats/token": completion_score(X_seen, X_hidden, phi_tr[::2], seed=K),
                 "CV R2 of click-through": cv_r2(theta_all.mean(0))})
choose = pd.DataFrame(rows).set_index("K")
print(choose.round(3))

fig, axes = plt.subplots(1, 2, figsize=(12, 3.6))
axes[0].plot(choose.index, choose.iloc[:, 0], "o-")
axes[0].set(xscale="log", xlabel="number of topics K", ylabel="nats per held-out token",
            title="Held-out words: more topics, always better")
axes[1].plot(choose.index, choose.iloc[:, 1], "o-", color="C1")
axes[1].set(xscale="log", xlabel="number of topics K", ylabel="cross-validated R$^2$",
            title="Click-through: flat from K = 4")
for ax in axes:
    ax.minorticks_off()
    ax.set_xticks(choose.index, choose.index);

import nutpie

se_logit = np.sqrt(1 / (tests["clicks"] + 0.5) + 1 / (tests["impressions"] - tests["clicks"] + 0.5)).to_numpy()
coords_ctr = {"test": np.arange(len(logit_ctr)), "month": month_codes, "topic": labels_txt}
with pm.Model(coords=coords_ctr) as ctr_model:
    theta_data = pm.Data("theta_data", theta_txt[0].mean(0), dims=("test", "topic"))
    y_data = pm.Data("y_data", logit_ctr, dims="test")
    mu_month = pm.GaussianRandomWalk("mu_month", sigma=0.1, init_dist=pm.Normal.dist(-4.0, 1.0),
                                     dims="month")
    gamma = pm.ZeroSumNormal("gamma", sigma=0.5, dims="topic")
    sigma = pm.HalfNormal("sigma", 0.5)
    pm.Normal("y", mu_month[month_idx] + theta_data @ gamma, pt.sqrt(sigma**2 + se_logit**2),
              observed=y_data, dims="test")

compiled_ctr = nutpie.compile_pymc_model(ctr_model)


def fit_gamma(theta_in, y_in, draws=1000, seed=0):
    fit = nutpie.sample(compiled_ctr.with_data(theta_data=theta_in, y_data=y_in), draws=draws, seed=seed,
                        save_warmup=False, progress_bar=False)
    return fit


def plugin_and_imputed(theta_draws, y_in, n_imp=8, seed=0):
    """Posterior draws of gamma: plug-in (posterior-mean mixtures) and pooled over n_imp mixture draws."""
    plug = fit_gamma(theta_draws.mean(0), y_in, seed=seed)
    picks = np.linspace(0, len(theta_draws) - 1, n_imp).astype(int)
    imp = [fit_gamma(theta_draws[p], y_in, draws=250, seed=seed + 1 + i).posterior["gamma"].values.reshape(-1, K_TXT)
           for i, p in enumerate(picks)]
    return plug, plug.posterior["gamma"].values.reshape(-1, K_TXT), np.concatenate(imp)


t0 = time.perf_counter()
fit_plugin, gamma_plugin, gamma_mi = plugin_and_imputed(theta_txt[0], logit_ctr, seed=RANDOM_SEED)
print(f"9 regressions: {time.perf_counter() - t0:.0f} s")
print("plug-in fit: divergences", int(fit_plugin.sample_stats["diverging"].sum()),
      "| max r_hat", round(float(az.rhat(fit_plugin.posterior["gamma"]).max()), 3))
print(pd.DataFrame({"plug-in mean": gamma_plugin.mean(0), "plug-in sd": gamma_plugin.std(0),
                    "imputed mean": gamma_mi.mean(0), "imputed sd": gamma_mi.std(0)},
                   index=labels_txt).round(3))

M_XYZ = np.array([[0.4124, 0.3576, 0.1805], [0.2126, 0.7152, 0.0722], [0.0193, 0.1192, 0.9505]])
WHITE = np.array([0.95047, 1.0, 1.08883])


def rgb_to_lab(rgb):
    c = np.asarray(rgb, float) / 255
    c = np.where(c > 0.04045, ((c + 0.055) / 1.055) ** 2.4, c / 12.92)
    f = c @ M_XYZ.T / WHITE
    f = np.where(f > (6 / 29) ** 3, np.cbrt(f), f / (3 * (6 / 29) ** 2) + 4 / 29)
    return np.stack([116 * f[..., 1] - 16, 500 * (f[..., 0] - f[..., 1]), 200 * (f[..., 1] - f[..., 2])], -1)


def lab_to_rgb(lab):
    fy = (lab[..., 0] + 16) / 116
    f = np.stack([fy + lab[..., 1] / 500, fy, fy - lab[..., 2] / 200], -1)
    xyz = np.where(f > 6 / 29, f**3, 3 * (6 / 29) ** 2 * (f - 4 / 29)) * WHITE
    c = np.clip(xyz @ np.linalg.inv(M_XYZ).T, 0, 1)
    return np.clip(np.where(c > 0.0031308, 1.055 * c ** (1 / 2.4) - 0.055, 12.92 * c), 0, 1)


def crop_box(img, threshold=12.0, frac=0.6):
    """Rows/columns where most pixels differ (Delta E > threshold) from the median corner colour."""
    lab = rgb_to_lab(img)
    corners = np.concatenate([lab[:3, :3], lab[:3, -3:], lab[-3:, :3], lab[-3:, -3:]]).reshape(-1, 3)
    fg = np.linalg.norm(lab - np.median(corners, 0), axis=-1) > threshold
    rows, cols = np.flatnonzero(fg.mean(1) > frac), np.flatnonzero(fg.mean(0) > frac)
    if len(rows) < 16 or len(cols) < 12:          # nothing clearly separable: keep the whole image
        return 0, img.shape[0], 0, img.shape[1]
    return rows[0] + 1, rows[-1], cols[0] + 1, cols[-1]   # one more pixel in, for resampling halos


def nearest(lab_pixels, codebook, chunk=100_000):
    out = np.empty(len(lab_pixels), np.int64)
    for s in range(0, len(lab_pixels), chunk):   # chunks: a pixels x codebook x 3 array gets big
        d2 = ((lab_pixels[s:s + chunk, None, :] - codebook[None]) ** 2).sum(-1)
        out[s:s + chunk] = d2.argmin(1)
    return out


def kmeans(x, k, g, iters=50):
    centres = [x[g.integers(len(x))]]
    for _ in range(k - 1):                        # k-means++ seeding
        d2 = ((x[:, None] - np.array(centres)[None]) ** 2).sum(-1).min(1)
        centres.append(x[g.choice(len(x), p=d2 / d2.sum())])
    centres = np.array(centres)
    for _ in range(iters):
        lab = nearest(x, centres)
        centres = np.array([x[lab == j].mean(0) if np.any(lab == j) else centres[j] for j in range(k)])
    return centres


boxes = np.array([crop_box(t) for t in thumbs])
kept_share = (boxes[:, 1] - boxes[:, 0]) * (boxes[:, 3] - boxes[:, 2]) / (64 * 48)
print("share of the thumbnail kept, median by era:",
      {e: round(float(np.median(kept_share[era_idx == i])), 2) for i, e in enumerate(ERAS)})
poster_lab = [rgb_to_lab(t[r0:r1, c0:c1]).reshape(-1, 3) for t, (r0, r1, c0, c1) in zip(thumbs, boxes)]

V_COL = 32
g_col = np.random.default_rng(3)
pixel_sample = np.concatenate([p[g_col.choice(len(p), 60, replace=False)] for p in poster_lab])
codebook = kmeans(pixel_sample, V_COL, g_col)
chroma, hue = np.hypot(codebook[:, 1], codebook[:, 2]), np.arctan2(codebook[:, 2], codebook[:, 1])
codebook = codebook[np.lexsort((codebook[:, 0], np.where(chroma < 12, -10, np.round(hue * 2))))]  # greys, then by hue
swatch = lab_to_rgb(codebook)
chroma, hue = np.hypot(codebook[:, 1], codebook[:, 2]), np.arctan2(codebook[:, 2], codebook[:, 1])
poster_words = [nearest(p, codebook) for p in poster_lab]

# For display, colours are grouped into families: the chromatic ones in hue order (red, orange,
# yellow, green, blue), then the earth tones and neutrals that dominate old paper (beige/tan,
# brown, grey), sorted light to dark within a family. CIELab hue alone would call every beige and
# brown "yellow" (they sit at hue angles of 70-95 degrees), so lightness and chroma split them off.
FAMILIES = ["red", "orange", "yellow", "green", "blue", "beige", "brown", "grey"]
hue_deg = np.degrees(hue) % 360
code_family = np.select(
    [chroma < 13,                                            # grey: too little colour to have a hue
     (hue_deg >= 105) & (hue_deg < 200),                     # green
     hue_deg >= 200,                                         # blue (no violets in this vocabulary)
     (hue_deg < 55) & (chroma >= 22) & (codebook[:, 0] < 58),  # red, brick, maroon
     (chroma >= 40) & (hue_deg < 75),                        # orange
     chroma >= 40,                                           # yellow: gold and ochre
     codebook[:, 0] < 50],                                   # brown: dark, muted warm
    [7, 3, 4, 0, 1, 2, 6], 5)                                # otherwise beige / tan
print({f: ", ".join(map(str, np.flatnonzero(code_family == i))) for i, f in enumerate(FAMILIES)})
display_order = np.lexsort((-codebook[:, 0], code_family))


def palette_squares(ax, weights, labels=True, min_label=0.1):
    """A colour distribution as a unit square: one column per colour family (width = the family's
    share; chromatic families in hue order, then beige, brown, grey), its colours stacked light (top)
    to dark (height = share within the family). Every rectangle's area is the colour's share."""
    w = np.asarray(weights, float) / np.sum(weights)
    x = 0.0
    for f in range(len(FAMILIES)):
        members = display_order[code_family[display_order] == f]
        total = w[members].sum()
        if total < 1e-4:
            continue
        y = 1.0
        for j in members:
            h = w[j] / total
            ax.add_patch(plt.Rectangle((x, y - h), total, h, facecolor=swatch[j], edgecolor="white", lw=0.6))
            y -= h
        if labels and total >= min_label:
            ax.text(x + total / 2, -0.02, FAMILIES[f], ha="center", va="top", fontsize=7, color="0.3")
        x += total
    ax.set(xlim=(0, 1), ylim=(0, 1), aspect="equal")
    ax.axis("off")


share_all = np.bincount(np.concatenate(poster_words), minlength=V_COL) / sum(map(len, poster_words))
fig = plt.figure(figsize=(12, 3.4))
left, right = fig.subfigures(1, 2, width_ratios=[2.6, 1])
axes = left.subplots(2, 1, height_ratios=[1, 1.4])
axes[0].imshow(swatch[display_order][None], aspect="auto")
axes[0].set(yticks=[])
axes[0].set_title("The colour vocabulary: 32 k-means colours by family, light to dark", fontsize=10)
fam_sorted = code_family[display_order]
starts = [np.flatnonzero(fam_sorted == f) for f in range(len(FAMILIES))]
axes[0].set_xticks([s_.mean() for s_ in starts if len(s_)], [FAMILIES[f] for f, s_ in enumerate(starts) if len(s_)],
                   fontsize=8)
for s_ in starts[1:]:
    if len(s_):
        axes[0].axvline(s_[0] - 0.5, color="white", lw=2)
axes[1].bar(range(V_COL), share_all[display_order], color=swatch[display_order], edgecolor="0.4", lw=0.4)
axes[1].set(xlim=(-0.5, V_COL - 0.5), xticks=[], ylabel="share of pixels")
palette_squares(right.subplots(), share_all, min_label=0.05)
right.suptitle("All 750 posters, as areas", fontsize=10);