"""
Regenerates the figures in EM_ADVANCED_claude_1015.tex into images_claude/.

Run from this directory:  python3 build_em_adv_figs.py

Figures produced:
  model_selection.png -- on HW4's 3gaussian.txt, K = 1..8: training log-likelihood (rises
                     forever), held-out log-likelihood (peaks), and BIC (picks a K).  Why
                     "maximize the likelihood" cannot choose the number of components.
  local_optima.png -- 300 random restarts of EM on 3gaussian.txt with K=3: the final
                     log-likelihoods cluster into a few distinct values, not one.  The
                     practical content of "EM is a hill-climber".
"""
import os
import numpy as np
import matplotlib
matplotlib.use("Agg")
import matplotlib.pyplot as plt
from scipy.special import logsumexp

HERE = os.path.dirname(os.path.abspath(__file__))
OUT = os.path.join(HERE, "images_claude")
HW3 = os.path.join(HERE, "..", "..", "HW3")
os.makedirs(OUT, exist_ok=True)

INK, C0, C1, C2 = "#1A1D21", "#2A6FB5", "#B5452A", "#1E7A5F"
plt.rcParams.update({
    "font.size": 10, "axes.edgecolor": INK, "axes.labelcolor": INK,
    "text.color": INK, "xtick.color": INK, "ytick.color": INK,
    "axes.spines.top": False, "axes.spines.right": False, "figure.dpi": 160,
})

# reuse the EM implementation the basic note's figures were built with, so the two
# notes cannot disagree about what the algorithm does
import sys
sys.path.insert(0, HERE)
from build_em_figs import em_gmm, log_gauss          # noqa: E402


def total_loglik(X, pi, mu, Sigma):
    K = len(pi)
    logw = np.column_stack([np.log(pi[k]) + log_gauss(X, mu[k], Sigma[k])
                            for k in range(K)])
    return logsumexp(logw, axis=1).sum()


def best_of(X, K, restarts=12, iters=250):
    best = None
    for s in range(restarts):
        out = em_gmm(X, K, iters=iters, seed=s)
        if best is None or out[3][-1] > best[3][-1]:
            best = out
    return best[0], best[1], best[2]


def fig_model_selection():
    X = np.loadtxt(os.path.join(HW3, "3gaussian.txt"))
    rng = np.random.RandomState(0)
    perm = rng.permutation(len(X))
    cut = int(0.7 * len(X))
    Xtr, Xte = X[perm[:cut]], X[perm[cut:]]
    D = X.shape[1]

    Ks = np.arange(1, 9)
    tr, te, bic = [], [], []
    for K in Ks:
        pi, mu, S = best_of(Xtr, K)
        tr.append(total_loglik(Xtr, pi, mu, S))
        te.append(total_loglik(Xte, pi, mu, S))
        # free parameters: K-1 weights + K means (D each) + K covariances (D(D+1)/2 each)
        P = (K - 1) + K * D + K * D * (D + 1) // 2
        bic.append(-2 * tr[-1] + P * np.log(len(Xtr)))
    tr, te, bic = np.array(tr), np.array(te), np.array(bic)

    fig, axes = plt.subplots(1, 3, figsize=(12.2, 3.6))
    for ax, v, col, title in [
        (axes[0], tr, C1, "training log-likelihood"),
        (axes[1], te, C0, "held-out log-likelihood"),
        (axes[2], bic, C2, "BIC  (lower is better)"),
    ]:
        ax.plot(Ks, v, color=col, lw=2.0, marker="o", ms=5)
        ax.set_xlabel("number of components $K$")
        ax.set_title(title, fontsize=10.5)
        ax.set_xticks(Ks)
    axes[0].annotate("still rising at $K=8$ (by $%.0f$),\njust imperceptibly — so this\ncurve can never choose $K$"
                     % (tr[-1] - tr[2]),
                     xy=(7.2, tr[6]), xytext=(2.4, tr[0] + 0.32 * (tr[-1] - tr[0])),
                     fontsize=9, color=C1,
                     arrowprops=dict(arrowstyle="->", color=C1, lw=1.0))
    print("   training LL:", np.array2string(tr, precision=0))
    kb = int(Ks[np.argmax(te)])
    axes[1].scatter([kb], [te.max()], s=70, facecolor="white", edgecolor=C0,
                    linewidths=1.8, zorder=5)
    axes[1].annotate(f"peaks at $K={kb}$", xy=(kb, te.max()),
                     xytext=(kb + 0.6, te.min() + 0.35 * (te.max() - te.min())),
                     fontsize=9, color=C0,
                     arrowprops=dict(arrowstyle="->", color=C0, lw=1.0))
    kc = int(Ks[np.argmin(bic)])
    axes[2].scatter([kc], [bic.min()], s=70, facecolor="white", edgecolor=C2,
                    linewidths=1.8, zorder=5)
    axes[2].annotate(f"minimum at $K={kc}$", xy=(kc, bic.min()),
                     xytext=(kc + 0.5, bic.min() + 0.42 * (bic.max() - bic.min())),
                     fontsize=9, color=C2,
                     arrowprops=dict(arrowstyle="->", color=C2, lw=1.0))
    fig.tight_layout()
    fig.savefig(os.path.join(OUT, "model_selection.png"), bbox_inches="tight")
    plt.close(fig)
    print(f"wrote model_selection.png   held-out best K={kb}, BIC best K={kc}")


def fig_local_optima():
    # The HW's 2gaussian/3gaussian files are too well separated to show this: every
    # restart finds the same optimum.  The three-hump 1-D data from the basic note's
    # first figure is the honest example -- its humps overlap, so restarts matter.
    rng = np.random.RandomState(7)
    X = np.concatenate([rng.normal(58, 4.0, 300),
                        rng.normal(78, 5.0, 500),
                        rng.normal(102, 6.0, 200)])[:, None]

    finals, mus = [], []
    for s in range(300):
        pi, mu, S, lls, _ = em_gmm(X, 3, iters=300, seed=s)
        finals.append(lls[-1])
        mus.append(np.sort(mu[:, 0]))
    finals, mus = np.array(finals), np.array(mus)
    best = finals.max()
    good = finals > best - 1.0

    fig, axes = plt.subplots(1, 2, figsize=(11.0, 3.7),
                             gridspec_kw={"width_ratios": [1.25, 1]})
    ax = axes[0]
    ax.hist(finals, bins=70, color=C0, alpha=0.75, edgecolor=C0, linewidth=0.5)
    ax.axvline(best, color=C1, lw=1.8)
    ax.annotate(f"global optimum\n({good.mean()*100:.0f}% of restarts)",
                xy=(best - 2, ax.get_ylim()[1] * 0.60),
                xytext=(best - 0.52 * (best - finals.min()),
                        ax.get_ylim()[1] * 0.72),
                fontsize=9, color=C1, ha="center",
                arrowprops=dict(arrowstyle="->", color=C1, lw=1.0))
    ax.annotate(f"{(~good).sum()} runs stop here\n(two humps merged)",
                xy=(finals.min() + 3, 32), xytext=(finals.min() + 14, 95),
                fontsize=9, color=INK,
                arrowprops=dict(arrowstyle="->", color=INK, lw=0.9))
    ax.set_xlabel(r"final $\sum_i \log p(x_i)$ after convergence")
    ax.set_ylabel("restarts")
    ax.set_title("300 random restarts, $K=3$", fontsize=10.5)

    ax = axes[1]
    ax.scatter(mus[~good, 0], mus[~good, 1], s=18, c="#9AA3AB", linewidths=0,
               alpha=0.7, label="a local optimum")
    ax.scatter(mus[good, 0], mus[good, 1], s=18, c=C1, linewidths=0,
               alpha=0.8, label="the global optimum")
    ax.scatter([58], [78], s=160, c=C2, marker="*", edgecolor="white",
               linewidths=0.8, zorder=5, label="truth")
    ax.set_xlabel(r"smallest fitted $\mu_k$")
    ax.set_ylabel(r"middle fitted $\mu_k$")
    ax.set_title("where the bad runs end up", fontsize=10.5)
    ax.set_xlim(56.5, 72.5)
    ax.set_ylim(70.3, 79.2)
    ax.legend(frameon=False, fontsize=8.5, loc="lower left")
    fig.tight_layout()
    fig.savefig(os.path.join(OUT, "local_optima.png"), bbox_inches="tight")
    plt.close(fig)
    print(f"wrote local_optima.png   best={best:.1f}  worst={finals.min():.1f}  "
          f"reaching best: {good.mean()*100:.0f}%  "
          f"distinct modes~{len(np.unique(finals.round(1)))}")


if __name__ == "__main__":
    fig_model_selection()
    fig_local_optima()
    print("all figures written to", OUT)