"""CSE 541 Autumn 2026, Homework 2: starter code.

Problems 2 and 6 are empirical. Fill in every function marked TODO; everything else is provided so that
your plots are comparable to ours. Run   python hw2_starter.py   to produce all figures (about a minute).
Tip: the matrices here are tiny, so multithreaded BLAS only slows things down; run with OMP_NUM_THREADS=1.

Conventions: n arms, T rounds, rewards X_t = theta[I_t] + N(0,1); regret curves are pseudo-regret
R_t = sum_{s<=t} (max_i theta_i - theta_{I_s}), averaged over seeds.
"""
import numpy as np
import matplotlib.pyplot as plt

# ======================================================================
#                 Problem 2: UCB, Thompson sampling, ETC, EXP3
# ======================================================================
def instance_a():
    theta = np.zeros(10); theta[0] = 0.1
    return theta


def instance_b():
    i = np.arange(1, 41)
    theta = 1.0 - 1.0 / np.sqrt(np.maximum(i - 1, 1))
    theta[0] = 1.0
    return theta


def regret_curve(theta, choices):
    """Pseudo-regret R_t = sum_{s<=t} Delta_{I_s} for a sequence of arm indices."""
    return np.cumsum(theta.max() - theta[np.asarray(choices)])


def ucb(theta, T, rng):
    """UCB with width sqrt(2 log(2 n T^2) / T_i(t)); pull each arm once first. Return the array of arms played."""
    n = len(theta)
    # TODO
    raise NotImplementedError


def thompson(theta, T, rng):
    """Thompson sampling with prior N(0,1) on each mean and N(theta_i, 1) likelihood (so the posterior of arm i
    after s pulls with sum S is N(S/(s+1), 1/(s+1))). Return the array of arms played."""
    n = len(theta)
    # TODO
    raise NotImplementedError


def etc(theta, T, m, rng):
    """Explore-then-commit: play arms round-robin until each has m pulls, then commit to the empirical best."""
    n = len(theta)
    # TODO
    raise NotImplementedError


def exp3(theta, T, rng):
    """EXP3 from lecture with uniform exploration lambda, eta = sqrt(log(n)/(3 n T)), gamma = eta n.
    EXP3 needs bounded losses: feed it ell_t = clip(-X_t / 4, -1, 1), which is almost never clipped here.
    Return the array of arms played."""
    n = len(theta)
    # TODO
    raise NotImplementedError


def problem2(T=100_000, seeds=range(5), etc_m=(10, 100, 1000)):
    fig, axes = plt.subplots(1, 2, figsize=(12, 4.3))
    for ax, (name, theta) in zip(axes, [("(a) n=10, one arm at 0.1", instance_a()), ("(b) n=40, Delta_i = 1/sqrt(i-1)", instance_b())]):
        algs = {"UCB": lambda th, rng: ucb(th, T, rng), "Thompson": lambda th, rng: thompson(th, T, rng),
                "EXP3": lambda th, rng: exp3(th, T, rng)}
        for m in etc_m:
            algs[f"ETC m={m}"] = (lambda mm: (lambda th, rng: etc(th, T, mm, rng)))(m)
        for label, fn in algs.items():
            R = np.mean([regret_curve(theta, fn(theta, np.random.default_rng(s))) for s in seeds], axis=0)
            ax.plot(R, label=label, linestyle="--" if "ETC" in label else "-")
            print(f"  [{name}] {label:12s} R_T = {R[-1]:9.1f}")
        ax.set_title(name); ax.set_xlabel("round t"); ax.set_ylabel("pseudo-regret"); ax.legend(fontsize=7)
    fig.tight_layout(); fig.savefig("problem2.pdf"); print("wrote problem2.pdf")


# ======================================================================
#                 Problem 6.1: G-optimal design
# ======================================================================
def design_matrix(X, lam, reg=1e-9):
    """A(lambda) = sum_i lambda_i x_i x_i^T (+ a hair of regularization)."""
    return (X * lam[:, None]).T @ X + reg * np.eye(X.shape[1])


def g_criterion(X, lam):
    """f(lambda) = max_i ||x_i||^2_{A(lambda)^{-1}}."""
    Ainv = np.linalg.inv(design_matrix(X, lam))
    return np.max(np.sum(X @ Ainv * X, axis=1))


def g_optimal(X, N=1000, rng=None):
    """Return a design lambda in the simplex over the rows of X that approximately minimizes g_criterion.
    Implement ONE of: greedy on h in {f, g}, or Frank-Wolfe on h in {f, g} (see the handout). N is the number
    of iterations/allocations. Must work for X with a few hundred rows and d <= 30."""
    rng = np.random.default_rng(0) if rng is None else rng
    n, d = X.shape
    # TODO
    raise NotImplementedError


def problem6_1(N=1000, d=10):
    """Compare the G-criterion of the uniform allocation with that of your design as the point cloud grows,
    for Gaussian points and for heavy-tailed (Student-t, 2 degrees of freedom) points."""
    ns = [10 + 2 ** i for i in range(1, 11)]
    plt.figure(figsize=(6.5, 4))
    for shape, draw in [("Gaussian", lambda rng, n: rng.standard_normal((n, d))),
                        ("Student-t(2)", lambda rng, n: rng.standard_t(2, size=(n, d)))]:
        unif, opt = [], []
        for n in ns:
            rng = np.random.default_rng(n)
            X = draw(rng, n)
            unif.append(g_criterion(X, np.ones(n) / n))
            opt.append(g_criterion(X, g_optimal(X, N=N, rng=rng)))
        plt.plot(ns, unif, marker="o", linestyle="--", label=f"{shape}: uniform allocation")
        plt.plot(ns, opt, marker="o", label=f"{shape}: your design")
        print(f"  {shape:12s} uniform f = {np.round(unif, 1)}\n  {shape:12s} design  f = {np.round(opt, 2)}")
    plt.axhline(d, color="k", linestyle=":", label="d (Kiefer-Wolfowitz)")
    plt.xscale("log"); plt.yscale("log"); plt.xlabel("n (number of points)"); plt.ylabel("max_i ||x_i||^2_{A(lambda)^{-1}}")
    plt.legend(fontsize=7); plt.title(f"G-criterion, d = {d}, N = {N} iterations"); plt.tight_layout(); plt.savefig("problem6_1.pdf")
    print("wrote problem6_1.pdf")


# ======================================================================
#          Problem 6.2: experimental design for function estimation
# ======================================================================
def f_true(x):
    return -x ** 2 + x * np.cos(8 * x) + np.sin(15 * x)


def sobolev_features(n=300, d=30, seed=0):
    """Locations X in [0,1] (50 on a grid, 250 clustered near 0.25) and the d-dimensional features Phi built from
    the first-order Sobolev kernel k(x, x') = 1 + min(x, x'). Returns X (n,), Phi (n, d), f_star (n,)."""
    rng = np.random.default_rng(seed)
    X = np.concatenate((np.linspace(0, 1, 50), 0.25 + 0.01 * rng.standard_normal(250)))
    X = np.sort(X)
    K = 1 + np.minimum(X[:, None], X[None, :])
    e, v = np.linalg.eigh(K)                      # eigenvalues in increasing order
    Phi = np.real(v @ np.diag(np.sqrt(np.abs(e))))[:, (n - d):]
    return X, Phi, f_true(X)


def sample_and_estimate(Phi, f_star, lam, tau, rng, reg=1e-6):
    """Draw tau locations i.i.d. from lam, observe y = f(x) + N(0,1), return (least-squares fit on all of X, X^T X)."""
    n, d = Phi.shape
    idx = rng.choice(n, size=tau, p=lam)
    y = f_star[idx] + rng.standard_normal(tau)
    XtX = Phi[idx].T @ Phi[idx]
    theta = np.linalg.lstsq(XtX + reg * np.eye(d), Phi[idx].T @ y, rcond=None)[0]
    return Phi @ theta, XtX


def problem6_2(tau=1000, seed=0):
    X, Phi, f_star = sobolev_features(seed=seed)
    n, d = Phi.shape
    rng = np.random.default_rng(seed)
    lam_G = g_optimal(Phi, N=1000, rng=rng)
    lam_U = np.ones(n) / n
    f_G, A_G = sample_and_estimate(Phi, f_star, lam_G, tau, rng)
    f_U, A_U = sample_and_estimate(Phi, f_star, lam_U, tau, rng)
    conf_G = np.sqrt(np.sum(Phi @ np.linalg.inv(A_G) * Phi, axis=1))
    conf_U = np.sqrt(np.sum(Phi @ np.linalg.inv(A_U) * Phi, axis=1))
    fig, axes = plt.subplots(1, 3, figsize=(15, 4.2))
    axes[0].plot(X, np.cumsum(lam_U), label="uniform (CDF)"); axes[0].plot(X, np.cumsum(lam_G), label="G-optimal (CDF)")
    axes[0].set_title("Plot 1: allocations over X"); axes[0].set_xlabel("x"); axes[0].set_ylabel("CDF"); axes[0].legend()
    axes[1].plot(X, f_star, "k", label="f_star"); axes[1].plot(X, f_G, label="f_G_Phi"); axes[1].plot(X, f_U, label="f_unif_Phi")
    axes[1].set_title("Plot 2: estimates"); axes[1].set_xlabel("x"); axes[1].set_ylabel("f(x)"); axes[1].legend()
    axes[2].plot(X, np.abs(f_G - f_star), label="|f_G_Phi - f_star|"); axes[2].plot(X, np.abs(f_U - f_star), label="|f_unif_Phi - f_star|")
    axes[2].axhline(np.sqrt(d / tau), color="k", linestyle=":", label="sqrt(d/tau)")
    axes[2].plot(X, conf_G, "--", label="conf_G"); axes[2].plot(X, conf_U, "--", label="conf_unif")
    axes[2].set_title("Plot 3: errors and confidence widths"); axes[2].set_xlabel("x"); axes[2].set_yscale("log"); axes[2].legend(fontsize=7)
    fig.tight_layout(); fig.savefig("problem6_2.pdf"); print("wrote problem6_2.pdf")
    print(f"  max |f_G - f*| = {np.max(np.abs(f_G - f_star)):.3f}   max |f_U - f*| = {np.max(np.abs(f_U - f_star)):.3f}   "
          f"max conf_G = {conf_G.max():.3f}   max conf_U = {conf_U.max():.3f}   sqrt(d/tau) = {np.sqrt(d / tau):.3f}")


# ======================================================================
#                 Problem 6.3: linear bandits on the Sobolev features
# ======================================================================
def lin_elimination(Phi, f_star, T, rng, tau=100, gamma=1.0, U=1.0, beta=None):
    """Phased elimination with a G-optimal design on the surviving arms (handout, Algorithm 'Elimination').
    If beta is None use the beta_k of the handout; otherwise use the constant beta given. Return the arms played (length T)."""
    # TODO
    raise NotImplementedError


def lin_ucb(Phi, f_star, T, rng, gamma=1.0, U=1.0):
    """LinUCB with beta_t = sqrt(gamma) U + sqrt(2 log(1/delta) + log(|V_t|/|V_0|)), delta = 1/T."""
    # TODO
    raise NotImplementedError


def lin_thompson(Phi, f_star, T, rng, gamma=1.0):
    """Linear Thompson sampling: theta_tilde ~ N(theta_t, V_t^{-1})."""
    # TODO
    raise NotImplementedError


def problem6_3(T=40_000, seed=0):
    X, Phi, f_star = sobolev_features(seed=seed)
    plt.figure(figsize=(6, 4))
    for name, fn in [("elimination (beta_k from handout)", lin_elimination),
                     ("elimination (beta = 2)", lambda P, f, T, r: lin_elimination(P, f, T, r, beta=2.0)),
                     ("LinUCB", lin_ucb), ("LinTS", lin_thompson)]:
        choices = fn(Phi, f_star, T, np.random.default_rng(seed))
        R = np.cumsum(f_star.max() - f_star[np.asarray(choices)])
        plt.plot(R, label=name); print(f"  {name:34s} R_T = {R[-1]:8.1f}")
    plt.xlabel("round t"); plt.ylabel("pseudo-regret"); plt.title("Linear bandits on the Sobolev features (n=300, d=30)")
    plt.legend(fontsize=8); plt.tight_layout(); plt.savefig("problem6_3.pdf"); print("wrote problem6_3.pdf")


if __name__ == "__main__":
    print("=== Problem 2 ===");   problem2()
    print("=== Problem 6.1 ==="); problem6_1()
    print("=== Problem 6.1 ==="); problem6_1()
    print("=== Problem 6.2 ==="); problem6_2()
    print("=== Problem 6.3 ==="); problem6_3()