# %% [markdown]
# # 0103 - Importance sampling for off-policy updates; ratio variance; the product-of-token-ratios explosion
# Companion to `notes/0103 - Importance Sampling and Off-Policy Updates.pdf`.  Exact enumeration, fp64.
# %%
import sys, math
sys.path.insert(0, "..")
from _m1_lib import *
# %% [markdown]
# ## 1. IS identity, unit mean, and Var_q[w] = chi^2 = E_p[w] - 1   (note Sec. 2.1)
# q = pi_old (sampler), p = pi_theta (target).  Sequence-level weight w(o) = p(o)/q(o).
# %%
lm = TinyLM(V=4, L=4, seed=0, scale=1.0)
th_old = lm.theta0.clone()
th_new = perturb(th_old, 0.5, seed=3)
q, p = lm.probs(th_old), lm.probs(th_new)
w = p / q
f = default_reward(lm) + 0.3 * lm.lengths                    # an arbitrary test function f(o)
check("E_q[w f] = E_p[f]  (IS identity)", abs(float((q * w * f).sum() - (p * f).sum())), 1e-14)
check("E_q[w] = 1", abs(float((q * w).sum()) - 1.0), 1e-14)
chi2 = float((p * w).sum()) - 1.0
check("Var_q[w] = E_p[w] - 1 = chi^2(p||q)", abs(float((q * w ** 2).sum() - 1.0) - chi2), 1e-14)
# variance of the IS estimator of E_p[f]:  Var_q[w f] = E_p[w f^2] - (E_p f)^2
v_direct = float((q * (w * f) ** 2).sum() - (q * w * f).sum() ** 2)
v_formula = float((p * w * f ** 2).sum() - (p * f).sum() ** 2)
check("Var_q[w f] = E_p[w f^2] - (E_p f)^2", abs(v_direct - v_formula), 1e-13)
# %% [markdown]
# ## 2. Self-normalised IS: bias O(1/n) and the effective sample size  ESS = (sum w)^2 / sum w^2  ~  n / (1 + chi^2)
# %%
gen = torch.Generator().manual_seed(21)
def snis(n, reps):
    est, ess = [], []
    for _ in range(reps):
        idx = torch.multinomial(q, n, replacement=True, generator=gen)
        wi = w[idx]
        est.append(float((wi * f[idx]).sum() / wi.sum()))
        ess.append(float(wi.sum() ** 2 / (wi ** 2).sum()))
    return torch.tensor(est), torch.tensor(ess)
target = float((p * f).sum())
for n in (4, 16, 64, 256):
    est, ess = snis(n, 40000 // (1 if n <= 16 else 4))
    bias = float(est.mean() - target); se = float(est.std() / len(est) ** 0.5)
    print(f"n={n:4d}  SNIS bias={bias:+.5f} (se {se:.5f})   mean ESS={float(ess.mean()):.2f}   n/(1+chi2)={n/(1+chi2):.2f}")
est, ess = snis(256, 4000)
check("MC: ESS at n=256 within 15% of n/(1+chi^2)", abs(float(ess.mean()) / (256 / (1 + chi2)) - 1), 0.15, kind="mc")
# %% [markdown]
# ## 3. Product of token ratios: exact chi^2 growth with horizon T  (long responses)
# Sequence weight over the first T tokens: w_T = prod_{t<=T} r_t.  E_q[w_T] = 1 for every T but E_q[w_T^2] grows geometrically.
# %%
lmL = TinyLM(V=3, L=12, seed=4, scale=1.0)
tok = lmL.theta0[:lmL.n_tok].reshape(lmL.V + 1, lmL.V); tok[:, lmL.EOS] -= 2.5      # make EOS rare -> long sequences
thq = lmL.theta0.clone(); thp = perturb(thq, 0.6, seed=9)
lq, lp = lmL.token_logprobs(thq), lmL.token_logprobs(thp)
qL = lq.sum(1).exp()
check("long-horizon model: sum_o q(o) = 1", abs(float(qL.sum()) - 1), 1e-12)
print("N sequences:", lmL.N, " mean length under q:", float((qL * lmL.lengths).sum()))
logw_cum = (lp - lq).cumsum(1)                                # log w_T for T = 1..L  (masked tokens contribute 0 -> ratio 1 after EOS)
rows = []
for T in range(1, lmL.L + 1):
    wT = logw_cum[:, T - 1].exp()
    m1 = float((qL * wT).sum()); m2 = float((qL * wT ** 2).sum())
    med = float((qL * (wT < 1).double()).sum())               # P_q(w_T < 1): the typical sample has weight < 1
    rows.append((T, m1, m2 - 1, med, float((qL * logw_cum[:, T - 1]).sum())))
    check(f"E_q[w_T] = 1 at T={T}", abs(m1 - 1), 1e-12) if T in (1, 6, 12) else None
print(" T   E[w]    chi2_T      P(w<1)   E[log w]")
for r in rows: print(f"{r[0]:2d}  {r[1]:.4f}  {r[2]:9.4f}   {r[3]:.4f}   {r[4]:+.4f}")
chi_T = torch.tensor([r[2] for r in rows])
assert bool((chi_T[1:] >= chi_T[:-1] - 1e-12).all())
check("chi^2_T non-decreasing in T (E_q[w_T^2] only grows)", 0.0, 0.0)
ratio_growth = float(chi_T[-1] / chi_T[0])
print(f"chi2_12 / chi2_1 = {ratio_growth:.1f}")
# chain rule:  E_q[log w_T] = - sum_{t<=T} E_q[KL_t(q || p)]   computed independently on the prefix tree
def tree_kl(thq, thp, T):
    tot = 0.0
    def rec(prefix, prob):
        nonlocal tot
        if len(prefix) >= T: return
        cq, cp = lmL.cond_probs(thq, prefix), lmL.cond_probs(thp, prefix)
        tot += prob * float((cq * (cq.log() - cp.log())).sum())
        for a in range(lmL.V - 1):
            rec(prefix + (a,), prob * float(cq[a]))
    rec((), 1.0)
    return tot
for T in (3, 8):
    check(f"E_q[log w_T] = -sum_t E[KL_t] (tree recursion), T={T}", abs(rows[T - 1][4] + tree_kl(thq, thp, T)), 1e-12)
# %% [markdown]
# ## 4. iid-token closed form  E_q[w_T^2] = (1 + chi2_1)^T   and ESS collapse
# %%
Vv = 3
qa = torch.tensor([0.5, 0.3, 0.2]); pa = torch.tensor([0.45, 0.33, 0.22])
chi1 = float((pa ** 2 / qa).sum()) - 1
def brute(T):
    seqs = torch.cartesian_prod(*[torch.arange(Vv)] * T) if T > 1 else torch.arange(Vv)[:, None]
    lqs = qa[seqs].log().sum(1); lps = pa[seqs].log().sum(1)
    return float((lqs.exp() * (lps - lqs).exp() ** 2).sum())
for T in (1, 3, 6, 9):
    check(f"iid tokens: E_q[w_T^2] = (1+chi2_1)^T, T={T}", abs(brute(T) - (1 + chi1) ** T) / (1 + chi1) ** T, 1e-12)
print(f"per-token chi2_1 = {chi1:.5f}")
for T in (1, 10, 100, 1000, 4000):
    print(f"T={T:5d}   chi2_T = {(1+chi1)**T-1:12.3f}   ESS/n ~ 1/(1+chi2) = {(1+chi1)**(-T):.3e}")
# %% [markdown]
# ## 5. Naive vs stable weight computation (fp32 running product underflows; log-space is exact)
# %%
gen = torch.Generator().manual_seed(77)
def weights(T, sigma, dtype):
    logr = (-0.5 * sigma ** 2 + sigma * torch.randn(T, generator=gen)).to(torch.float64)
    naive = torch.tensor(1.0, dtype=dtype)
    for x in logr.exp().to(dtype): naive = naive * x          # naive running product of token ratios
    stable = logr.sum().exp()                                  # log-space
    return float(naive), float(stable), float(logr.sum())
n64, s64, _ = weights(50, 0.1, torch.float64)
check("T=50: naive fp64 product == exp(sum log r)", abs(n64 - s64) / s64, 1e-12)
n32, s64b, lw = weights(4000, 0.5, torch.float32)
print(f"T=4000 sigma=0.5:  log w = {lw:.1f};  naive fp32 product = {n32};  log-space fp64 = {s64b:.3e}")
assert n32 == 0.0 and math.isfinite(lw)
# %% [markdown]
# ## 6. Token-level surrogate (PPO-style ratio) vs the exact off-policy gradient
# Exact:   grad J(theta) = E_old[ w(o) (R-b) sum_t grad log pi_theta(o_t|s_t) ]  (sequence-level weight, all tokens).
# Token-level surrogate:  L(theta) = E_old[ sum_t r_t(theta) (R-b) ] with r_t = pi_theta(o_t|s_t)/pi_old(o_t|s_t).
# They agree at theta_old and drift apart as ||theta - theta_old|| grows: value error O(delta^2), gradient error O(delta).
# %%
lm = TinyLM(V=4, L=4, seed=0, scale=1.0)
R = default_reward(lm)
th_old = lm.theta0.clone()
q = lm.probs(th_old)
b = float(q @ R)
def exact_grad(th):
    t = th.clone().requires_grad_(True)
    return torch.autograd.grad((lm.probs(t) * R).sum(), t)[0]
def seq_is_grad(th):
    t = th.clone().requires_grad_(True)
    lpn = lm.seq_logprob(t)
    wts = (lpn - lm.seq_logprob(th_old)).exp().detach()
    return torch.autograd.grad((q * wts * (R - b) * lpn).sum(), t)[0]
def token_surrogate(th):
    t = th.clone().requires_grad_(True)
    r = (lm.token_logprobs(t) - lm.token_logprobs(th_old)).exp() * lm.mask       # r_t(theta), 0 where masked
    L = (q[:, None] * r * (R - b)[:, None]).sum()
    return L, torch.autograd.grad(L, t)[0]
deltas = [0.0, 0.02, 0.04, 0.08, 0.16, 0.32]
L0 = float((q * lm.lengths * (R - b)).sum())
gerr, verr = [], []
for d in deltas:
    th = perturb(th_old, d, seed=5) if d > 0 else th_old.clone()
    ge = exact_grad(th)
    check(f"sequence-level IS gradient == exact grad J, delta={d}", maxabs(seq_is_grad(th), ge), 1e-14) if d in (0.0, 0.32) else None
    Lv, gt = token_surrogate(th)
    gerr.append(float((gt - ge).norm())); verr.append(abs(float(Lv.detach()) - L0 - float((lm.probs(th) * R).sum() - (q * R).sum())))
    print(f"delta={d:5.2f}   |grad_tok - grad_J| = {gerr[-1]:.3e}   |dL - dJ| = {verr[-1]:.3e}")
check("token surrogate gradient == exact gradient at theta_old", gerr[0], 1e-14)
check("token surrogate value == J to first order at theta_old (error at delta=0)", verr[0], 1e-14)
import numpy as np
sl_g = np.polyfit(np.log(deltas[1:]), np.log(gerr[1:]), 1)[0]
sl_v = np.polyfit(np.log(deltas[1:]), np.log(verr[1:]), 1)[0]
print(f"log-log slopes: gradient error {sl_g:.2f} (expect ~1), value error {sl_v:.2f} (expect ~2)")
check("gradient-error slope ~ 1", abs(sl_g - 1.0), 0.15)
check("value-error slope ~ 2", abs(sl_v - 2.0), 0.15)
print("done; checks:", len(CHECKS))