Lecture 10: Sequence-space DEQNs

# Run-mode switch (smoke = CPU-bounded for CI, teaching = laptop figures, production = full reproduction).
RUN_MODE = "smoke"  # one of: "smoke", "teaching", "production"
SEED = 0

import jax
import jax.numpy as jnp
from jax import lax
import matplotlib.pyplot as plt
import numpy as np
import optax
from scipy.interpolate import BSpline

print(f"JAX version: {jax.__version__}")
print(f"Devices: {jax.devices()}")

# Derived from the kit-mandated RUN_MODE in cell 1.
TUTORIAL_MODE = "full" if RUN_MODE == "production" else "quick"
assert TUTORIAL_MODE in {"quick", "full"}
print(f"Tutorial mode: {TUTORIAL_MODE}")

#===Model Parameters===
ALPHA = 0.36   # capital share
DELTA = 0.025  # depreciation
BETA = 0.93    # discount factor (log utility is hardcoded throughout)

# Idiosyncratic shocks: low and high efficiency
EPS_VALS = jnp.array([0.5, 1.5])
PI_EPS = jnp.array([[0.9, 0.1],  # P(eps'=L | eps=L) = 0.9
                     [0.1, 0.9]]) # P(eps'=H | eps=H) = 0.9

# Aggregate shocks: low and high productivity
Z_VALS = jnp.array([0.93, 1.07])
PI_Z = jnp.array([[0.7, 0.3],  # persistence = 0.7
                  [0.3, 0.7]])

N_EPS = len(EPS_VALS)
N_Z = len(Z_VALS)

# Steady-state capital (from the deterministic Euler equation)
R_ss = 1.0 / BETA
K_ss = ((R_ss - 1.0 + DELTA) / ALPHA) ** (1.0 / (ALPHA - 1.0))
print(f"Steady-state capital: K_ss = {K_ss:.2f}")

# Price functions
def compute_prices(K, L, Z):
    """
    Cobb-Douglas factor prices.
    """
    KL = K / L
    R = 1.0 - DELTA + ALPHA * Z * KL ** (ALPHA - 1.0)
    w = (1.0 - ALPHA) * Z * KL ** ALPHA
    return R, w

# === Capital Grid ===
GRID_CONFIG = {
    "quick": {"N": 60, "K_MAX": 60.0},
    "full": {"N": 100, "K_MAX": 100.0},
}[TUTORIAL_MODE]

N = GRID_CONFIG["N"]          # number of intervals (N+1 grid points)
K_MAX = GRID_CONFIG["K_MAX"]  # upper bound on capital

K_GRID_SHIFT = 0.5
BASIS_SHIFT = 0.5  

x_grid = np.linspace(np.log(K_GRID_SHIFT), np.log(K_MAX + K_GRID_SHIFT), N + 1)
k_grid = jnp.array(np.exp(x_grid) - K_GRID_SHIFT)
dk = jnp.diff(k_grid)
x0 = float(x_grid[0])
delta_x = float((x_grid[-1] - x_grid[0]) / N)

# Visualize the grid spacing
fig, ax = plt.subplots(figsize=(10, 2))
ax.scatter(np.asarray(k_grid), np.zeros(N + 1), s=6, alpha=0.7)
ax.set_xlabel("Capital ")
ax.set_title(f"Log-spaced grid: {N+1} points in [0, {K_MAX}]")
ax.set_yticks([])
ax.axvline(K_ss, color='red', linestyle='--', alpha=0.5, label=f"deterministic $K_{{ss}}={K_ss:.1f}$")
ax.legend()
plt.tight_layout()
plt.show()
print(f"Grid spacing: {float(dk[0]):.3f} (near k=0) to {float(dk[-1]):.1f} (near k_max)")

# === Build I-Spline Basis ===
N_INTERIOR = 26 if TUTORIAL_MODE == "quick" else 56
ORDER = 4           # cubic I-splines

def build_ispline_basis(k_grid, n_interior=56, order=4, basis_shift=None):
    """Precompute I-spline basis at the grid points."""
    if basis_shift is None:
        basis_shift = BASIS_SHIFT
    k_arr = np.asarray(k_grid)
    x = np.log(basis_shift + k_arr)
    x_min, x_max = float(x[0]), float(x[-1])

    # Equidistant interior knots in x-space
    knots_int = np.linspace(x_min, x_max, n_interior + 2)[1:-1]
    knots = np.concatenate([np.repeat(x_min, order), knots_int, np.repeat(x_max, order)])
    n_basis = len(knots) - order

    # Compute I-splines (= integral of B-splines)
    basis = np.zeros((n_basis, len(k_arr)))
    for idx in range(n_basis):
        c = np.zeros(n_basis)
        c[idx] = 1.0
        bspl = BSpline(knots, c, order - 1, extrapolate=False)
        t_span = knots[idx + order] - knots[idx]
        if t_span < 1e-12:
            continue
        anti = bspl.antiderivative()
        basis[idx] = (order / t_span) * np.nan_to_num(anti(x) - anti(x_min), 0.0)

    return jnp.array(np.clip(basis, 0.0, 1.0))

ispline_basis = build_ispline_basis(k_grid, N_INTERIOR, ORDER, basis_shift=BASIS_SHIFT)
J = ispline_basis.shape[0]
print(f"Built {J} I-spline basis functions on {N+1} grid points")


# What the basis matrix actually is
B = np.asarray(ispline_basis)
print("Basis matrix shape:", B.shape)
print(f"Entry B[j, n] = I_j(log({BASIS_SHIFT:.2f} + k_grid[n]))")

fig, ax = plt.subplots(figsize=(10, 4))
im = ax.imshow(
    B,
    aspect='auto',
    origin='lower',
    extent=[float(k_grid[0]), float(k_grid[-1]), 0, B.shape[0]-1]
)
ax.set_xlabel("Capital grid point $k_n$")
ax.set_ylabel("Basis index $j$")
ax.set_title("I-spline basis evaluated on the policy grid")
fig.colorbar(im, ax=ax, label="basis value")
plt.tight_layout()
plt.show()


fig, axes = plt.subplots(1, 2, figsize=(12, 4))
for j in range(J):
    axes[0].plot(np.asarray(k_grid), np.asarray(ispline_basis[j]), linewidth=0.8, alpha=0.5)
axes[0].set_xlabel("Capital $k$")
axes[0].set_ylabel("$I_j(k)$")
axes[0].set_title(f"I-spline basis in $k$-space ({J} functions)")
axes[0].grid(True, alpha=0.3)

x_vals = np.log(BASIS_SHIFT + np.asarray(k_grid))
for j in range(J):
    axes[1].plot(x_vals, np.asarray(ispline_basis[j]), linewidth=0.8, alpha=0.5)
axes[1].set_xlabel(rf"$x = \log({BASIS_SHIFT:.2f} + k)$")
axes[1].set_ylabel("$I_j(x)$")
axes[1].set_title("I-spline basis in transformed space")
axes[1].grid(True, alpha=0.3)

plt.tight_layout()
plt.show()

# === Neural Network Actor ===
H = 50            # history length (last H aggregate shocks)
HIDDEN_DIM = 64   # MLP hidden layer width
INPUT_DIM = H * (N_Z + 1)  # one-hot + value per lag

def init_actor(key, input_dim, hidden_dim, n_eps, J):
    """Initialize a 2-hidden-layer MLP."""
    k1, k2, k3 = jax.random.split(key, 3)
    output_dim = n_eps * (1 + J)  # per eps: 1 alpha logit + J weight logits
    return {
        "W_0": jnp.sqrt(2.0/input_dim) * jax.random.normal(k1, (input_dim, hidden_dim)),
        "b_0": jnp.zeros(hidden_dim),
        "W_1": jnp.sqrt(2.0/hidden_dim) * jax.random.normal(k2, (hidden_dim, hidden_dim)),
        "b_1": jnp.zeros(hidden_dim),
        "W_2": 0.01 * jax.random.normal(k3, (hidden_dim, output_dim)),
        "b_2": jnp.zeros(output_dim),
    }

def actor_c_grid(psi, agg_input, R, w):
    """Forward pass: aggregate input -> consumption on the k-grid.

    Returns: c_grid of shape (n_eps, N+1)
    """
    # MLP forward pass (tanh activation — sweep-validated, 3x better than relu)
    h = jnp.tanh(agg_input @ psi["W_0"] + psi["b_0"])
    h = jnp.tanh(h @ psi["W_1"] + psi["b_1"])
    raw = (h @ psi["W_2"] + psi["b_2"]).reshape(N_EPS, 1 + J)

    # Boundary MPC: sigmoid -> (0, 1)
    alpha = jax.nn.sigmoid(raw[:, 0])                        # (n_eps,)

    # Phantom-zero softmax: guarantees weights >= 0 and sum < 1
    logits = raw[:, 1:]                                       # (n_eps, J)
    augmented = jnp.concatenate([jnp.zeros((N_EPS, 1)), logits], axis=1)
    w_tilde = jax.nn.softmax(augmented, axis=1)[:, 1:]       # (n_eps, J)

    # MPC on the grid: alpha * (1 - weighted I-splines)
    mpc = alpha[:, None] * (1.0 - w_tilde @ ispline_basis)   # (n_eps, N+1)

    # Cumulate consumption on the grid from the MPC schedule.
    # Boundary: c(k_0) = MPC(k_0) · m(k_0), where m = wε + Rk is cash-on-hand
    c0 = mpc[:, 0] * (w * EPS_VALS + R * k_grid[0])
    # Increments: Δc_n = MPC(k_n) · R · Δk_n — each MPC value used exactly once
    dc = mpc[:, 1:] * (R * dk[None, :])
    c_grid = jnp.concatenate([c0[:, None], c0[:, None] + jnp.cumsum(dc, axis=1)], axis=1)
    return c_grid

# Initialize and check
key = jax.random.PRNGKey(0)
key, init_key = jax.random.split(key)
psi = init_actor(init_key, INPUT_DIM, HIDDEN_DIM, N_EPS, J)
n_params = sum(p.size for p in jax.tree_util.tree_leaves(psi))
print(f"Actor parameters: {n_params}")
print(f"Input dim: {INPUT_DIM} (H={H} x {N_Z+1} features per lag)")
print(f"Output dim: {N_EPS} x (1 + {J}) = {N_EPS * (1 + J)}")

# === Visualize Initial Policy ===
def encode_Z_history(z_hist):
    """Encode Z-history as one-hot + values."""
    onehot = jax.nn.one_hot(z_hist, N_Z).reshape(-1)
    values = Z_VALS[z_hist].reshape(-1)
    return jnp.concatenate([onehot, values])

def shape_report(c_grid, R, w, title):
    """Verify the four shape guarantees of the I-spline MPC policy.

    Checks (all should hold by construction):
      min Δc >= 0        — consumption is increasing in k
      max(c - m) <= 0    — consumption is feasible (c ≤ cash-on-hand)
      MPC ∈ [0, 1]       — marginal propensity in valid range
      max ΔMPC <= 0      — MPC is decreasing (consumption is concave)
    """
    m_grid = float(w) * np.asarray(EPS_VALS)[:, None] + float(R) * k_np[None, :]
    min_dc = float(np.min(np.diff(c_grid, axis=1)))
    max_excess = float(np.max(c_grid - m_grid))
    mpc_intervals = np.diff(c_grid, axis=1) / (float(R) * np.asarray(dk)[None, :])
    min_mpc = float(np.min(mpc_intervals))
    max_mpc = float(np.max(mpc_intervals))
    max_upturn = float(np.max(np.diff(mpc_intervals, axis=1)))
    print(f"Shape check for {title}:")
    print(f"  min Δc across grid intervals      = {min_dc: .3e}  (should be >= 0)")
    print(f"  max(c - cash on hand)             = {max_excess: .3e}  (should be <= 0)")
    print(f"  MPC range on intervals            = [{min_mpc: .3f}, {max_mpc: .3f}]")
    print(f"  max increase in MPC across k-grid = {max_upturn: .3e}  (should be <= 0)")

# Use a constant Z=Z_L history for visualization
z_hist_dummy = jnp.zeros(H, dtype=jnp.int32)
agg_input = encode_Z_history(z_hist_dummy)
R_vis, w_vis = compute_prices(K_ss, 1.0, float(Z_VALS[0]))
c_grid_init = np.asarray(actor_c_grid(psi, agg_input, float(R_vis), float(w_vis)))

k_np = np.asarray(k_grid)
fig, axes = plt.subplots(1, 2, figsize=(12, 4))
colors = ["#1f77b4", "#ff7f0e"]
labels = [r"$\varepsilon_L$", r"$\varepsilon_H$"]

for e in range(N_EPS):
    m = float(w_vis) * float(EPS_VALS[e]) + float(R_vis) * k_np
    axes[0].plot(k_np, c_grid_init[e], color=colors[e], label=labels[e])
    axes[0].plot(k_np, m, color=colors[e], linestyle=':', alpha=0.3, label=f"$m$ ({labels[e]})")

    # MPC at k=0: consumption as fraction of cash-on-hand (= boundary MPC α)
    mpc = np.zeros_like(k_np)
    mpc[0] = c_grid_init[e, 0] / max(m[0], 1e-10)
    # MPC at k>0: finite-difference Δc/Δm = Δc/(R·Δk), matches I-spline MPC by construction
    mpc[1:] = np.diff(c_grid_init[e]) / (float(R_vis) * np.asarray(dk))
    axes[1].plot(k_np, mpc, color=colors[e], label=labels[e])

axes[0].set_xlabel("Capital $k$")
axes[0].set_ylabel("Consumption $c$")
axes[0].set_title("Initial policy (random weights)")
axes[0].legend(fontsize=8)
axes[0].grid(True, alpha=0.3)
axes[1].set_xlabel("Capital $k$")
axes[1].set_ylabel("MPC")
axes[1].set_title("MPC (decreasing by construction)")
axes[1].set_ylim(-0.05, 1.05)
axes[1].grid(True, alpha=0.3)
axes[1].legend(fontsize=8)
plt.tight_layout()
plt.show()

shape_report(c_grid_init, float(R_vis), float(w_vis), "initial random policy")
print("Even before training, the grid policy is increasing, feasible, and concave by construction.")

# === Simulation: Young (2010) Non-Stochastic Method ===
SIM_CONFIG = {
    "quick": {"T": 100, "N_AGG": 8, "BUFFER_CAP": 128},
    "full":  {"T": 100, "N_AGG": 8, "BUFFER_CAP": 128},
}[TUTORIAL_MODE]

T = SIM_CONFIG["T"]
N_AGG = SIM_CONFIG["N_AGG"]
BUFFER_CAP = SIM_CONFIG["BUFFER_CAP"]

def interpolate(vals_grid, eps_idx, k):
    """Piecewise-linear interpolation on the k-grid."""
    x = jnp.log(K_GRID_SHIFT + k)
    j = jnp.clip(jnp.floor((x - x0) / delta_x).astype(jnp.int32), 0, N - 1)
    w_lo = (k_grid[j + 1] - k) / (k_grid[j + 1] - k_grid[j])
    return w_lo * vals_grid[eps_idx, j] + (1.0 - w_lo) * vals_grid[eps_idx, j + 1]

def draw_Z_path(key, T, z_init):
    """Draw T aggregate shocks. Returns (emitted_path, terminal_carry).

    emitted_path = [z_t, z_{t+1}, ..., z_{t+T-1}],
    terminal_carry = z_{t+T}.
    """
    keys = jax.random.split(key, T)
    def step(z, k):
        z_next = (jax.random.uniform(k) >= PI_Z[z, 0]).astype(jnp.int32)
        return z_next, z
    z_terminal, Z_path = lax.scan(step, z_init, keys)
    return Z_path, z_terminal

def make_agg_inputs(full_Z_seq, H, T):
    """Build encoded Z-history inputs for each simulation step."""
    offsets = jnp.arange(T)[:, None] + jnp.arange(H)[None, :]
    windows = full_Z_seq[offsets]
    return jax.vmap(encode_Z_history)(windows)

# --- Young (2010) distribution functions ---

def distribution_aggregates(mu):
    """Compute aggregate K and L from distribution mu(eps, k).

    mu has shape (N_EPS, N+1) and sums to 1.
    K = Σ k · μ(ε,k),  L = Σ ε · μ(ε,k).
    """
    K = jnp.sum(mu * k_grid[None, :])
    L = jnp.sum(mu * EPS_VALS[:, None])
    return jnp.maximum(K, 0.01), jnp.maximum(L, 0.01)

def distribution_step(mu, c_grid_t, R, w):
    """Advance the distribution one period (Young 2010).

    For each (eps, k) grid point:
      1. Compute next-period capital: k' = w·eps + R·k - c(eps, k)
      2. Find the two neighboring grid points around k'
      3. Split the mass to those neighbors (lottery)
      4. Apply the idiosyncratic transition matrix Pi_eps

    Args:
        mu:        (N_EPS, N+1) current distribution, sums to 1
        c_grid_t:  (N_EPS, N+1) consumption policy on the grid
        R, w:      current factor prices
    Returns:
        mu_next:   (N_EPS, N+1) next-period distribution
    """
    # Next-period capital for every (eps, k) grid point
    m_grid = w * EPS_VALS[:, None] + R * k_grid[None, :]
    k_next = jnp.clip(m_grid - c_grid_t, 0.0, k_grid[-1])

    # Find grid position of k' in log-space
    x_next = jnp.log(K_GRID_SHIFT + k_next)
    j_lo = jnp.clip(
        jnp.floor((x_next - x0) / delta_x).astype(jnp.int32), 0, N - 1)

    # Lottery weights: fraction of mass going to lower vs upper neighbor
    k_lo = k_grid[j_lo]
    k_hi = k_grid[j_lo + 1]
    w_lo = (k_hi - k_next) / jnp.maximum(k_hi - k_lo, 1e-12)
    w_hi = 1.0 - w_lo

    # Scatter mass to neighboring grid points for each eps state
    def scatter_one_eps(i):
        mass = mu[i]                    # (N+1,)
        j = j_lo[i]                     # (N+1,) integer indices
        result = jnp.zeros(N + 1)
        result = result.at[j].add(mass * w_lo[i])
        result = result.at[j + 1].add(mass * w_hi[i])
        return result

    mu_after = jax.vmap(scatter_one_eps)(jnp.arange(N_EPS))   # (N_EPS, N+1)

    # Apply idiosyncratic transition: mu_next[i'] = Σ_i Pi[i, i'] · mu_after[i]
    mu_next = PI_EPS.T @ mu_after
    return mu_next

def initial_distribution(K_target):
    """Create an initial distribution centered around K_target.

    eps: uniform across states; k: Gaussian around K_target.
    Returns mu of shape (N_EPS, N+1), summing to 1.
    """
    eps_weight = jnp.ones(N_EPS) / N_EPS
    sigma_k = K_target * 0.3
    density = jnp.exp(-0.5 * ((k_grid - K_target) / sigma_k) ** 2)
    density = density / jnp.sum(density)
    mu = eps_weight[:, None] * density[None, :]
    return mu / jnp.sum(mu)

def simulate_young(psi, mu_init, Z_path, agg_inputs):
    """Simulate T periods using Young's non-stochastic method.

    Instead of tracking individual households, evolves the full
    distribution mu(eps, k) on the grid. Aggregates K, L are exact
    weighted sums — no Monte Carlo sampling noise.

    Args:
        mu_init:    (N_EPS, N+1) initial distribution
        Z_path:     (T,) aggregate shock indices
        agg_inputs: (T, INPUT_DIM) encoded Z-histories
    Returns:
        mu_final:   (N_EPS, N+1) terminal distribution
    """
    def step(mu, inputs):
        Z_idx, agg_input_t = inputs
        K_t, L_t = distribution_aggregates(mu)
        K_t = jax.lax.stop_gradient(K_t)
        L_t = jax.lax.stop_gradient(L_t)
        R_t, w_t = compute_prices(K_t, L_t, Z_VALS[Z_idx])
        c_grid_t = actor_c_grid(psi, agg_input_t, R_t, w_t)
        mu_next = distribution_step(mu, c_grid_t, R_t, w_t)
        return mu_next, None

    mu_final, _ = lax.scan(step, mu_init, (Z_path, agg_inputs))
    return mu_final

print(f"Simulation config: T={T}, n_agg={N_AGG}, buffer={BUFFER_CAP}")
print("Using Young (2010) non-stochastic method — no MC noise in aggregates.")

# === Fischer-Burmeister Loss ===

def fb_loss_one_state(psi, z_history, mu):
    """Compute mean FB² on the full (eps, k) grid for one replay state.

    With Young's method, aggregates K, L come from the distribution mu
    (exact weighted sums). Next-period aggregates come from advancing
    mu one step with distribution_step — deterministic, no MC noise.
    """
    z_t = z_history[-1]

    # Current prices from distribution (exact aggregates)
    K_t, L_t = distribution_aggregates(mu)
    K_t, L_t = jax.lax.stop_gradient(K_t), jax.lax.stop_gradient(L_t)
    R_t, w_t = compute_prices(K_t, L_t, Z_VALS[z_t])
    R_t, w_t = jax.lax.stop_gradient(R_t), jax.lax.stop_gradient(w_t)

    # Current policy on the full grid
    agg_input = encode_Z_history(z_history)
    c_grid_t = actor_c_grid(psi, agg_input, R_t, w_t)

    # Next-period capital for each (eps, k) on the grid
    m_grid = w_t * EPS_VALS[:, None] + R_t * k_grid[None, :]
    k_next_grid = jnp.clip(m_grid - c_grid_t, 0.0, k_grid[-1])

    # Next-period aggregates from Young step (exact, deterministic)
    mu_next = distribution_step(mu, c_grid_t, R_t, w_t)
    K_next, L_next = distribution_aggregates(mu_next)
    K_next, L_next = jax.lax.stop_gradient(K_next), jax.lax.stop_gradient(L_next)

    # Continuation value: q = E[β R' u'(c')]
    # Outer loop over next-Z unrolled at trace time; inner eps vectorized.
    q = jnp.zeros_like(c_grid_t)
    for z_next in range(N_Z):
        z_hist_next = jnp.concatenate([z_history[1:], jnp.array([z_next])])
        agg_next = encode_Z_history(z_hist_next)
        R_next, w_next = compute_prices(K_next, L_next, Z_VALS[z_next])
        R_next, w_next = jax.lax.stop_gradient(R_next), jax.lax.stop_gradient(w_next)
        c_grid_next = actor_c_grid(psi, agg_next, R_next, w_next)

        # Vectorized next-eps integration
        def _interp_next(i_next):
            k_flat = k_next_grid.ravel()
            idx = jnp.full(k_flat.shape, i_next, dtype=jnp.int32)
            return interpolate(c_grid_next, idx, k_flat).reshape(N_EPS, N + 1)

        c_next_all = jax.vmap(_interp_next)(jnp.arange(N_EPS))
        mu_u = 1.0 / jnp.maximum(c_next_all, 1e-10)       # u'(c) = 1/c
        exp_uc = jnp.einsum('ij,jik->ik', PI_EPS, mu_u)
        q += PI_Z[z_t, z_next] * BETA * R_next * exp_uc

    # Euler-implied consumption and FB residual
    c_euler = 1.0 / jnp.maximum(q, 1e-10)
    c_scale = jnp.maximum(c_grid_t, 1e-10)
    g = (c_euler - c_grid_t) / c_scale          # relative Euler gap
    s = k_next_grid / c_scale                    # relative slack
    fb = jnp.sqrt(g**2 + s**2 + 1e-12) - g - s  # Fischer-Burmeister residual

    return jnp.mean(fb**2)

def fb_loss(psi, buf_z, buf_mu):
    """Mean FB² over all buffer states."""
    fb2_all = jax.vmap(lambda z, mu: fb_loss_one_state(psi, z, mu))(buf_z, buf_mu)
    return jnp.mean(fb2_all)

print("FB loss defined. With Young's method, aggregates are exact — no MC noise.")

# === Training Setup ===
# Hyperparameters from the sweep leaderboard (fastest to FB² < 1e-6):
# MLP 2×64 tanh, fb_batch=16, fb_steps=10, lr=5e-5, H=50, T=100
TRAIN_CONFIG = {
    "quick": {"LR": 5e-5, "FB_STEPS": 10, "FB_BATCH": 16, "N_EPOCHS": 5000,  "LOG_EVERY": 100},
    "full":  {"LR": 5e-5, "FB_STEPS": 10, "FB_BATCH": 16, "N_EPOCHS": 10_000, "LOG_EVERY": 200},
}[TUTORIAL_MODE]

LR = TRAIN_CONFIG["LR"]
FB_STEPS = TRAIN_CONFIG["FB_STEPS"]
FB_BATCH = TRAIN_CONFIG["FB_BATCH"]      # mini-batch size for FB gradient steps
N_EPOCHS = TRAIN_CONFIG["N_EPOCHS"]
LOG_EVERY = TRAIN_CONFIG["LOG_EVERY"]

optimizer = optax.chain(optax.clip_by_global_norm(1.0), optax.adam(LR))
opt_state = optimizer.init(psi)

@jax.jit
def fb_update(psi, opt_state, buf_z, buf_mu):
    loss, grads = jax.value_and_grad(fb_loss)(psi, buf_z, buf_mu)
    updates, new_opt = optimizer.update(grads, opt_state)
    new_psi = optax.apply_updates(psi, updates)
    return new_psi, new_opt, loss

@jax.jit
def run_simulations(key, psi, z_hists, mu_starts, z_inits):
    """Generate shocks, simulate distributions forward, return terminal states."""
    key, k1 = jax.random.split(key)
    z_keys = jax.random.split(k1, N_AGG)
    Z_paths, Z_terms = jax.vmap(lambda k, z0: draw_Z_path(k, T, z0))(z_keys, z_inits)

    full_Z = jnp.concatenate([z_hists, Z_paths[:, 1:]], axis=1)
    agg_inputs = jax.vmap(lambda seq: make_agg_inputs(seq, H, T))(full_Z)

    def sim1(mu, zp, ai):
        return simulate_young(jax.lax.stop_gradient(psi), mu, zp, ai)
    mu_finals = jax.vmap(sim1)(mu_starts, Z_paths, agg_inputs)

    return key, mu_finals, Z_paths, Z_terms

rng = np.random.default_rng(0)
buffer_z = []
buffer_mu = []

mu_init = initial_distribution(K_ss)
for _ in range(min(BUFFER_CAP, 16)):
    buffer_z.append(np.zeros(H, dtype=np.int32))
    buffer_mu.append(np.asarray(mu_init))

print(f"Ready to train: {N_EPOCHS} epochs, lr={LR}, fb_steps={FB_STEPS}, fb_batch={FB_BATCH}")
print(f"Buffer seeded with {len(buffer_z)} distributions around K_ss={K_ss:.2f}")

import time

history = []
t0 = time.time()

for epoch in range(N_EPOCHS):
    idx = rng.integers(0, len(buffer_z), size=N_AGG)
    z_hists = jnp.stack([jnp.array(buffer_z[i]) for i in idx])
    mu_starts = jnp.stack([jnp.array(buffer_mu[i]) for i in idx])
    z_inits = jnp.array([buffer_z[i][-1] for i in idx], dtype=jnp.int32)

    key, mu_finals, Z_paths, Z_terms = run_simulations(key, psi, z_hists, mu_starts, z_inits)

    mu_f = np.array(mu_finals)
    Zp = np.array(Z_paths)
    Zt = np.array(Z_terms)
    for i in range(N_AGG):
        # Zp[i, 0] duplicates the parent's terminal Z (draw_Z_path emits z_init
        # before transitioning), so skip it with [1:]. Append z_terminal to
        # complete the chain so the stored history ends at the correct period.
        full_Z = np.concatenate([buffer_z[idx[i]], Zp[i, 1:], [Zt[i]]])
        buffer_z.append(full_Z[-H:])
        buffer_mu.append(mu_f[i])
    while len(buffer_z) > BUFFER_CAP:
        buffer_z.pop(0)
        buffer_mu.pop(0)

    all_z = jnp.stack([jnp.array(z) for z in buffer_z])
    all_mu = jnp.stack([jnp.array(m) for m in buffer_mu])

    # Mini-batch FB: subsample buffer entries each gradient step.
    # With Young's method aggregates are exact, so the only noise comes from
    # the mini-batch — this is beneficial SGD stochasticity.
    fb2 = 0.0
    for _ in range(FB_STEPS):
        mb = rng.integers(0, all_z.shape[0], size=min(FB_BATCH, all_z.shape[0]))
        psi, opt_state, fb2 = fb_update(psi, opt_state, all_z[mb], all_mu[mb])

    if (epoch + 1) % LOG_EVERY == 0:
        fb2_val = float(fb2)
        # Mean K from buffer distributions (exact weighted sum)
        buf_K = float(np.mean([
            np.sum(np.asarray(m) * np.asarray(k_grid)[None, :]) for m in buffer_mu
        ]))
        wall = time.time() - t0
        history.append((epoch + 1, fb2_val, buf_K, wall))
        print(f"Epoch {epoch + 1:5d} | FB² = {fb2_val:.2e} | K̄ = {buf_K:.2f} | {wall:.0f}s")

wall_total = time.time() - t0
print(f"\nDone! {N_EPOCHS} epochs in {wall_total:.0f}s ({wall_total/60:.1f} min)")
print(f"Final FB² = {float(fb2):.2e}")

# === Plot 1: Convergence ===
epochs_h = [h[0] for h in history]
fb2_h = [h[1] for h in history]

fig, ax = plt.subplots(figsize=(8, 4))
ax.semilogy(epochs_h, fb2_h, linewidth=2, color='#1f77b4')
ax.axhline(1e-6, color='red', linestyle='--', alpha=0.5, label='$10^{-6}$ target')
ax.set_xlabel("Epoch", fontsize=12)
ax.set_ylabel("FB$^2$ (log scale)", fontsize=12)
ax.set_title("Euler Equation Accuracy During Training", fontsize=14)
ax.legend(fontsize=11)
ax.grid(True, alpha=0.3)
plt.tight_layout()
plt.show()

# === Plot 2: Learned Policy Functions ===
fig, axes = plt.subplots(1, 2, figsize=(13, 5))
colors = ["#1f77b4", "#ff7f0e"]
labels_eps = [r"$\varepsilon_L$ (low income)", r"$\varepsilon_H$ (high income)"]
labels_Z = ["$Z_L$ (recession)", "$Z_H$ (boom)"]

# Mean K from buffer distributions
buf_K = float(np.mean([
    np.sum(np.asarray(m) * np.asarray(k_grid)[None, :]) for m in buffer_mu
]))

for z_idx in range(N_Z):
    z_hist = jnp.full(H, z_idx, dtype=jnp.int32)
    agg_input = encode_Z_history(z_hist)
    R, w = compute_prices(max(buf_K, 0.01), 1.0, float(Z_VALS[z_idx]))
    c_grid_learned = np.asarray(actor_c_grid(psi, agg_input, float(R), float(w)))

    ls = '-' if z_idx == 0 else '--'
    for e in range(N_EPS):
        m = float(w) * float(EPS_VALS[e]) + float(R) * k_np
        axes[0].plot(k_np, c_grid_learned[e], color=colors[e], linestyle=ls,
                     linewidth=2, label=f"{labels_eps[e]}, {labels_Z[z_idx]}")

        # MPC at k=0: consumption as fraction of cash-on-hand (= boundary MPC α)
        mpc = np.zeros_like(k_np)
        mpc[0] = c_grid_learned[e, 0] / max(m[0], 1e-10)
        # MPC at k>0: finite-difference Δc/Δm = Δc/(R·Δk), matches I-spline MPC by construction
        mpc[1:] = np.diff(c_grid_learned[e]) / (float(R) * np.asarray(dk))
        axes[1].plot(k_np, mpc, color=colors[e], linestyle=ls, linewidth=2)

axes[0].set_xlabel("Capital $k$", fontsize=12)
axes[0].set_ylabel("Consumption $c$", fontsize=12)
axes[0].set_title("Learned consumption policy", fontsize=14)
axes[0].legend(fontsize=8, loc='lower right')
axes[0].grid(True, alpha=0.3)
axes[0].set_xlim(0, min(50, K_MAX))

axes[1].set_xlabel("Capital $k$", fontsize=12)
axes[1].set_ylabel("MPC", fontsize=12)
axes[1].set_title("Marginal propensity to consume", fontsize=14)
axes[1].set_ylim(-0.05, 1.05)
axes[1].grid(True, alpha=0.3)
axes[1].set_xlim(0, min(50, K_MAX))

plt.tight_layout()
plt.show()

# shape_report (defined in the initial-visualization cell above) verifies
# the four I-spline guarantees: c increasing, feasible, MPC ∈ [0,1], MPC decreasing
shape_report(c_grid_learned, float(R), float(w), "trained policy in the final plotted state")

# === Euler Error Diagnostics ===

def euler_diagnostics_one_state(psi, z_history, mu):
    """Per-grid-point Euler gap (g) and FB residual for one buffer state.

    Same computation as fb_loss_one_state but returns the full (eps, k) grids
    instead of a scalar mean, plus the consumption and MPC schedules.
    """
    z_t = z_history[-1]

    K_t, L_t = distribution_aggregates(mu)
    R_t, w_t = compute_prices(float(K_t), float(L_t), Z_VALS[z_t])
    R_t_f, w_t_f = float(R_t), float(w_t)

    agg_input = encode_Z_history(z_history)
    c_grid_t = actor_c_grid(psi, agg_input, R_t_f, w_t_f)

    m_grid = w_t_f * EPS_VALS[:, None] + R_t_f * k_grid[None, :]
    k_next_grid = jnp.clip(m_grid - c_grid_t, 0.0, k_grid[-1])

    # Next-period aggregates from Young step
    mu_next = distribution_step(mu, c_grid_t, R_t_f, w_t_f)
    K_next, L_next = distribution_aggregates(mu_next)
    K_next_f, L_next_f = float(K_next), float(L_next)

    q = jnp.zeros_like(c_grid_t)
    for z_next in range(N_Z):
        z_hist_next = jnp.concatenate([z_history[1:], jnp.array([z_next])])
        agg_next = encode_Z_history(z_hist_next)
        R_next, w_next = compute_prices(K_next_f, L_next_f, Z_VALS[z_next])
        R_nf, w_nf = float(R_next), float(w_next)
        c_grid_next = actor_c_grid(psi, agg_next, R_nf, w_nf)

        def _interp_next(i_next):
            k_flat = k_next_grid.ravel()
            idx = jnp.full(k_flat.shape, i_next, dtype=jnp.int32)
            return interpolate(c_grid_next, idx, k_flat).reshape(N_EPS, N + 1)

        c_next_all = jax.vmap(_interp_next)(jnp.arange(N_EPS))
        mu_u = 1.0 / jnp.maximum(c_next_all, 1e-10)
        exp_uc = jnp.einsum('ij,jik->ik', PI_EPS, mu_u)
        q += PI_Z[z_t, z_next] * BETA * R_nf * exp_uc

    c_euler = 1.0 / jnp.maximum(q, 1e-10)
    c_scale = jnp.maximum(c_grid_t, 1e-10)
    g_grid = (c_euler - c_grid_t) / c_scale
    s_grid = k_next_grid / c_scale
    fb_grid = jnp.sqrt(g_grid**2 + s_grid**2 + 1e-12) - g_grid - s_grid

    # MPC via finite-difference (= I-spline MPC at grid points by construction)
    c_np = np.asarray(c_grid_t)
    m_np = np.asarray(m_grid)
    mpc = np.zeros_like(c_np)
    mpc[:, 0] = c_np[:, 0] / np.maximum(m_np[:, 0], 1e-10)
    mpc[:, 1:] = np.diff(c_np, axis=1) / (R_t_f * np.asarray(dk)[None, :])

    return {
        "K": float(K_t), "c_grid": c_np, "mpc_grid": mpc,
        "g_grid": np.asarray(g_grid), "fb_grid": np.asarray(fb_grid),
    }

# Evaluate on all buffer states
print(f"Running Euler diagnostics on {len(buffer_z)} buffer states...")
diag_results = []
for i in range(len(buffer_z)):
    d = euler_diagnostics_one_state(
        psi, jnp.array(buffer_z[i]), jnp.array(buffer_mu[i])
    )
    diag_results.append(d)

all_K = np.array([d["K"] for d in diag_results])
all_c = np.stack([d["c_grid"] for d in diag_results])        # (n_states, N_EPS, N+1)
all_mpc = np.stack([d["mpc_grid"] for d in diag_results])    # (n_states, N_EPS, N+1)
abs_g = np.abs(np.stack([d["g_grid"] for d in diag_results]))
abs_fb = np.abs(np.stack([d["fb_grid"] for d in diag_results]))

print(f"K range: [{all_K.min():.2f}, {all_K.max():.2f}]")
print(f"Mean |g| = {abs_g.mean():.2e},  Mean FB² = {(abs_fb**2).mean():.2e}")

# === Scatter: Consumption and MPC vs Aggregate Capital K ===
# At selected individual capital levels, how do c and MPC vary with the
# aggregate state?  Tight scatter clouds = approximate aggregation holds.

target_k_vals = [0.0, 1.0, 3.0, 7.0, 15.0, 30.0, float(k_np[-1])]
target_k_vals = [tk for tk in target_k_vals if tk <= float(k_np[-1])]
k_indices = [int(np.argmin(np.abs(k_np - tk))) for tk in target_k_vals]

colors_eps = ["#1f77b4", "#ff7f0e"]
eps_labels = [r"$\varepsilon_L$", r"$\varepsilon_H$"]

fig, axes = plt.subplots(2, len(k_indices), figsize=(3.5 * len(k_indices), 8),
                         sharex=True)
for col, ki in enumerate(k_indices):
    for e in range(N_EPS):
        axes[0, col].scatter(all_K, all_c[:, e, ki], s=3, alpha=0.3,
                             color=colors_eps[e],
                             label=eps_labels[e] if col == 0 else None)
        axes[1, col].scatter(all_K, all_mpc[:, e, ki], s=3, alpha=0.3,
                             color=colors_eps[e])

    axes[0, col].set_title(f"$k = {k_np[ki]:.1f}$", fontsize=9)
    axes[1, col].set_xlabel("$K$", fontsize=9)
    axes[0, col].grid(True, alpha=0.3)
    axes[1, col].grid(True, alpha=0.3)
    axes[0, col].tick_params(labelsize=7)
    axes[1, col].tick_params(labelsize=7)

axes[0, 0].set_ylabel("Consumption $c$")
axes[1, 0].set_ylabel("MPC")
axes[0, 0].legend(fontsize=7, markerscale=3)
fig.suptitle("Consumption and MPC vs aggregate capital $K$", fontsize=13, y=1.02)
plt.tight_layout()
plt.show()

print("Each column is a fixed individual capital level k.")
print("Variation across K comes from the encoder mapping different Z-histories to different policy coefficients.")

# === Euler / FB Error Bands ===
# Median (line) and 10th–90th percentile (band) of |g| and |FB(g,s)|
# across replay buffer states, for each income state.

n_states = abs_g.shape[0]
fig, axes = plt.subplots(1, N_EPS, figsize=(6.2 * N_EPS, 8.2), sharex=True)

for i in range(N_EPS):
    fb_med = np.median(abs_fb[:, i, :], axis=0)
    fb_p10 = np.percentile(abs_fb[:, i, :], 10.0, axis=0)
    fb_p90 = np.percentile(abs_fb[:, i, :], 90.0, axis=0)

    # |FB| panel
    axes[i].fill_between(k_np, np.maximum(fb_p10, 1e-12),
                            np.maximum(fb_p90, 1e-12),
                            color=colors_eps[i], alpha=0.18)
    axes[i].plot(k_np, np.maximum(fb_med, 1e-12),
                    color=colors_eps[i], linewidth=2.0)
    axes[i].set_yscale("log")
    axes[i].set_ylabel(r"$|FB(g,s)|$")
    axes[i].set_xlabel("Capital $k$")
    axes[i].grid(True, alpha=0.3)



fig.suptitle(
    f"Replay-state one-step Euler / FB diagnostic ({n_states} states)",
    fontsize=13, y=1.02)
plt.tight_layout()
plt.show()

print("The band shows the 10th–90th percentile range across buffer states.")
print("Wider bands indicate more variation in accuracy across aggregate states.")
print("High |g| near k=0 reflects the borrowing constraint; high |g| at high k")
print("is often on grid points with zero household mass.")