# %% [markdown]
# # 0109 - KL regularisation estimators k1, k2, k3: value bias/variance, what their GRADIENTS optimise, why beta = 0
# Companion to `notes/0109 - KL Estimators.pdf`.  Samples x ~ q = pi_theta,  p = pi_ref,  r = p/q,  KL(q||p) = E_q[-log r].
# %%
import sys, math
sys.path.insert(0, "..")
from _m1_lib import *
import numpy as np
# %% [markdown]
# ## 1. Values.  E_q[k1] = E_q[k3] = KL(q||p);  E_q[k2] = KL + bias,  bias = -sum_{n>=3} E[l^n]/n!,  l = log r
#   k1 = -log r,   k2 = (log r)^2 / 2,   k3 = (r-1) - log r        (Schulman, "Approximating KL divergence")
# %%
def estimators(logr):
    return -logr, 0.5 * logr ** 2, torch.expm1(logr) - logr
gen = torch.Generator().manual_seed(2)
for trial in range(3):
    K = 6
    q = torch.softmax(torch.randn(K, generator=gen) * 1.2, 0); p = torch.softmax(torch.randn(K, generator=gen) * 1.2, 0)
    logr = (p / q).log()
    kl = float((q * (q / p).log()).sum())
    k1, k2, k3 = estimators(logr)
    check(f"categorical #{trial}: E_q[k1] = KL(q||p)", abs(float((q * k1).sum()) - kl), 1e-14)
    check(f"categorical #{trial}: E_q[k3] = KL(q||p)", abs(float((q * k3).sum()) - kl), 1e-14)
    bias2 = float((q * k2).sum()) - kl
    series = -sum(float((q * logr ** n).sum()) / math.factorial(n) for n in range(3, 40))
    check(f"categorical #{trial}: bias(k2) = -sum_(n>=3) E[l^n]/n!  (series, 40 terms)", abs(bias2 - series), 1e-12)
    # the master identity KL = sum_{n>=2} E[l^n]/n!  (from E_q[e^l] = 1)
    ms = sum(float((q * logr ** n).sum()) / math.factorial(n) for n in range(2, 40))
    check(f"categorical #{trial}: KL = sum_(n>=2) E_q[l^n]/n!", abs(ms - kl), 1e-12)
    print(f"      KL = {kl:.5f}   E[k2] - KL = {bias2:+.5f}   (leading term -E[l^3]/6 = {-float((q*logr**3).sum())/6:+.5f})")
# pointwise positivity
rg = torch.exp(torch.linspace(-12, 12, 100001))
k1g, k2g, k3g = estimators(rg.log())
check("k3 >= 0 and k2 >= 0 pointwise (min over a log-grid of r)", max(0.0, -float(min(k3g.min(), k2g.min()))), 0.0)
print(f"   k1 < 0 on r > 1 (fraction of log-grid where k1<0: {float((k1g<0).double().mean()):.3f}) -> k1 is a noisy, sign-changing estimator")
# %% [markdown]
# ## 2. Gaussian closed forms (equal variance, shift mu => KL = mu^2/2).  With s2 = mu^2:
#   Var(k1) = s2,   Var(k2) = s2^3/4 + s2^2/2,   Var(k3) = e^{s2} - 1 - s2,   bias(k2) = s2^2/8.   Verified by Gauss-Hermite quadrature.
# %%
xs, ws = np.polynomial.hermite_e.hermegauss(120); ws = ws / ws.sum()
def gauss_moments(shift):
    x = torch.tensor(xs); w = torch.tensor(ws)
    logr = shift * x - 0.5 * shift ** 2             # q = N(0,1), p = N(shift,1): log p/q = shift*x - shift^2/2  (E_q[r] = 1)
    k = estimators(logr)
    kl = shift ** 2 / 2
    out = []
    for kk in k:
        m = float((w * kk).sum()); v = float((w * kk ** 2).sum()) - m ** 2
        out.append((m, v))
    return kl, out
print("   KL     est   mean-KL (bias)     std/KL      closed-form std/KL")
for shift in (0.1, math.sqrt(2 * 0.05), 1.0):
    kl, out = gauss_moments(shift)
    s2 = shift ** 2
    cf = [s2, s2 ** 3 / 4 + s2 ** 2 / 2, math.exp(s2) - 1 - s2]
    cb = [0.0, s2 ** 2 / 8, 0.0]
    for i, nm in enumerate(("k1", "k2", "k3")):
        m, v = out[i]
        print(f"   {kl:.4f}  {nm}  {m-kl:+.3e}   {math.sqrt(v)/kl:9.4f}   {math.sqrt(cf[i])/kl:9.4f}")
        check(f"Gaussian shift={shift:.3f}: {nm} bias == closed form", abs((m - kl) - cb[i]), 1e-12)
        check(f"Gaussian shift={shift:.3f}: {nm} variance == closed form", abs(v - cf[i]) / cf[i], 1e-10)
# compare with the table in Schulman's note (KL=0.005 and KL=0.5): std/KL = 20, 1.42, 1.42 and 2, 1.73, 1.7 ; bias(k2)/KL = 0.002, 0.25
kl, out = gauss_moments(0.1)
check("matches Schulman's table at KL=0.005: std(k1)/KL ~ 20", abs(math.sqrt(out[0][1]) / kl - 20.0), 0.2)
kl, out = gauss_moments(1.0)
check("matches Schulman's table at KL=0.5: bias(k2)/KL = 0.25", abs((out[1][0] - kl) / kl - 0.25), 1e-9)
# %% [markdown]
# ## 3. Heavy tails: Var(k3) is infinite when chi^2(p||q) is; one token the reference loves and the policy nearly never samples
# %%
q = torch.tensor([0.001, 0.999]); p = torch.tensor([0.30, 0.70])
logr = (p / q).log(); k1, k2, k3 = estimators(logr)
kl = float((q * (q / p).log()).sum())
for nm, k in (("k1", k1), ("k2", k2), ("k3", k3)):
    m = float((q * k).sum()); v = float((q * k ** 2).sum()) - m ** 2
    print(f"   {nm}: mean {m:+.5f}  std {math.sqrt(v):.5f}   (KL = {kl:.5f})   max value {float(k.max()):.1f}")
chi2_pq = float((p ** 2 / q).sum() - 1)
print(f"   chi2(p||q) = E_q[r^2]-1 = {chi2_pq:.1f}  -> Var(k3) is dominated by one rare token with r = {float((p/q)[0]):.0f}")
check("E_q[r^2] - 1 = chi2(p||q) exactly", abs(float((q * (p / q) ** 2).sum()) - 1 - chi2_pq), 1e-9)
# %% [markdown]
# ## 4. GRADIENTS (on-policy, single-step categorical policy with logits theta; samples x ~ q_theta DETACHED, gradient flows through q_theta in r)
#   k1-as-loss:  E[grad k1] = E[grad log q] = 0                        -> useless as a regulariser
#   k2-as-loss:  E[grad k2] = E[(log q - log p) grad log q] = grad KL(q||p)           (the reverse KL we meant)
#   k3-as-loss:  E[grad k3] = E[(1 - r) grad log q]        = -sum_x p grad log q = grad KL(p||q)   (the FORWARD KL)
#   k1-in-reward (detached coefficient):  E[(log q - log p) grad log q] = grad KL(q||p)
# %%
gen = torch.Generator().manual_seed(7)
K = 6
th_ref = torch.randn(K, generator=gen) * 1.0
def grads(delta):
    th = (th_ref + delta * torch.randn(K, generator=gen)).clone().requires_grad_(True)
    q = torch.softmax(th, 0); p = torch.softmax(th_ref, 0)
    lq, lp = torch.log_softmax(th, 0), torch.log_softmax(th_ref, 0)
    qd = q.detach()
    r = (lp - lq).exp()
    # expectation over x ~ q (weights detached) of the gradient of each estimator
    g = {}
    g['k1'] = torch.autograd.grad((qd * (-(lp - lq))).sum(), th, retain_graph=True)[0]
    g['k2'] = torch.autograd.grad((qd * 0.5 * (lp - lq) ** 2).sum(), th, retain_graph=True)[0]
    g['k3'] = torch.autograd.grad((qd * (r - 1 - (lp - lq))).sum(), th, retain_graph=True)[0]
    # k1 in reward: detached coefficient (log q - log p) times score
    g['k1r'] = torch.autograd.grad((qd * (lq - lp).detach() * lq).sum(), th, retain_graph=True)[0]
    g['rev'] = torch.autograd.grad((q * (lq - lp)).sum(), th, retain_graph=True)[0]          # grad KL(q||p), exact
    g['fwd'] = torch.autograd.grad((p * (lp - lq)).sum(), th, retain_graph=True)[0]          # grad KL(p||q), exact
    return g
g = grads(0.8)
check("E[grad k1] = 0  (k1 as a loss carries no regularising gradient)", float(g['k1'].abs().max()), 1e-14)
check("E[grad k2] == grad KL(q||p)  (reverse KL)", maxabs(g['k2'], g['rev']), 1e-14)
check("k1-in-reward (detached coefficient) == grad KL(q||p)", maxabs(g['k1r'], g['rev']), 1e-14)
check("E[grad k3] == grad KL(p||q)  (FORWARD KL, not the one k3 estimates)", maxabs(g['k3'], g['fwd']), 1e-14)
print(f"   ||grad KL(q||p) - grad KL(p||q)|| / ||grad KL(q||p)|| = {float((g['rev']-g['fwd']).norm()/g['rev'].norm()):.3f}   (different vectors)")
print("   relative gap between the two gradients as q -> p (first-order agreement, 'k3 is a first-order approximation'):")
prev = None
for delta in (0.4, 0.2, 0.1, 0.05, 0.025):
    gen.manual_seed(7); _ = torch.randn(K, generator=gen)          # same direction each time
    gg = grads(delta)
    gap = float((gg['rev'] - gg['fwd']).norm() / gg['rev'].norm())
    print(f"     delta={delta:5.3f}   gap = {gap:.4f}")
    prev = gap
assert prev < 0.05
check("gradient gap forward vs reverse KL vanishes as delta -> 0 (gap at delta=0.025 < 5%)", 0.0, 0.0)
# %% [markdown]
# ## 5. SEQUENCES: per-token k2/k3 'as loss' ignore that earlier tokens change the later contexts
#   true:   grad KL_seq(pi||ref) = E[ (sum_t l_t) * score(o) ] = E[ sum_t (sum_{t'>=t} l_t') s_t ]   ("k1 in reward" with reward-to-go)
#   per-token k2 loss, summed over t:   E[ sum_t l_t s_t ]  =  true  -  sum_t E[ KL_t(s_t) * score(o_<t) ]    (missing prefix-score term)
# %%
lm = TinyLM(V=4, L=4, seed=0, scale=1.0)
th = lm.theta0.clone(); th_r = perturb(th, 1.0, seed=2)
pi = lm.probs(th)
lpi, lref = lm.token_logprobs(th), lm.token_logprobs(th_r)
l_tok = (lpi - lref) * lm.mask                                     # l_t = log pi - log ref  = -log r_t
Ttens = token_score_tensor(lm, th)                                 # [N, L, P]
S = Ttens.sum(1)
KL_seq = float((pi * l_tok.sum(1)).sum())
t_ = th.clone().requires_grad_(True)
KLt = (lm.probs(t_) * (lm.seq_logprob(t_) - lm.seq_logprob(th_r))).sum()
g_true = torch.autograd.grad(KLt, t_)[0]
g_seq = (pi[:, None] * l_tok.sum(1)[:, None] * S).sum(0)
rtg = l_tok.flip(1).cumsum(1).flip(1) * lm.mask
g_rtg = (pi[:, None] * (Ttens * rtg[:, :, None]).sum(1)).sum(0)
g_tok = (pi[:, None] * (Ttens * l_tok[:, :, None]).sum(1)).sum(0)                  # per-token k2-as-loss gradient, summed over t
check("E[(sum_t l_t) * score] == grad KL_seq(pi||ref)", maxabs(g_seq, g_true), 1e-14)
check("reward-to-go form E[sum_t (sum_{t'>=t} l_t') s_t] == grad KL_seq", maxabs(g_rtg, g_true), 1e-14)
print(f"   per-token k2-as-loss gradient vs true sequence-KL gradient: relative error {float((g_tok-g_true).norm()/g_true.norm()):.4f}")
# the missing term: sum_t E[ KL_t(s_t) * score(o_<t) ] where KL_t(s) = KL(pi(.|s) || ref(.|s)) and score(o_<t) = sum_{s<t} s_s
prefix_score = torch.cumsum(Ttens, 1) - Ttens                                      # sum_{s<t} s_s
KLt_tok = torch.zeros(lm.N, lm.L)
for i, sq in enumerate(lm.seqs):
    for t in range(len(sq)):
        cp, cr = lm.cond_probs(th, sq[:t]), lm.cond_probs(th_r, sq[:t])
        KLt_tok[i, t] = (cp * (cp.log() - cr.log())).sum()
missing = (pi[:, None] * (KLt_tok[:, :, None] * prefix_score).sum(1)).sum(0)
check("per-token k2 loss gradient == true gradient - sum_t E[KL_t * prefix score]  (exact identity)", maxabs(g_tok, g_true - missing), 1e-13)
# per-token k3 loss gradient equals the gradient of the sum of FORWARD per-state KLs with frozen context distribution
r_tok = (lref - lpi).exp() * lm.mask + (1 - lm.mask)
g_k3 = (pi[:, None] * (Ttens * ((1 - r_tok) * lm.mask)[:, :, None]).sum(1)).sum(0)
t_ = th.clone().requires_grad_(True)
tot = 0.0
for s in lm.prefixes():
    pr = 1.0
    for k in range(len(s)): pr *= float(lm.cond_probs(th, s[:k])[s[k]])           # P_pi(context = s) (detached)
    cp_t, cr = lm.cond_probs(t_, s), lm.cond_probs(th_r, s)
    tot = tot + pr * (cr * (cr.log() - cp_t.log())).sum()                        # KL(ref(.|s) || pi_theta(.|s)) with frozen context weight
g_fwd_frozen = torch.autograd.grad(tot, t_)[0]
check("per-token k3-as-loss gradient == grad of sum_s P_pi(s) KL(ref(.|s)||pi_theta(.|s)) (context weights frozen)", maxabs(g_k3, g_fwd_frozen), 1e-13)
# values: sequence-level KL from the three per-token estimators
k1_seq = float((pi * l_tok.sum(1)).sum())
k3_tok = (torch.expm1(lref - lpi) - (lref - lpi)) * lm.mask
k2_tok = 0.5 * (lpi - lref) ** 2 * lm.mask
print(f"   sequence KL = {KL_seq:.5f};  E[sum_t k1] = {k1_seq:.5f};  E[sum_t k3] = {float((pi*k3_tok.sum(1)).sum()):.5f};  E[sum_t k2] = {float((pi*k2_tok.sum(1)).sum()):.5f}")
check("E[sum_t k3_t] = KL_seq (per-token estimator is unbiased for the sequence KL value)", abs(float((pi * k3_tok.sum(1)).sum()) - KL_seq), 1e-13)
# %% [markdown]
# ## 6. Naive vs stable evaluation of k3 (fp32): (r-1)-log r cancels catastrophically when |log r| is small
# %%
lr32 = torch.tensor([1e-1, 1e-2, 1e-3, 1e-4, 1e-5], dtype=torch.float32)
naive = (lr32.exp() - 1) - lr32
stable = torch.expm1(lr32) - lr32
truth = (torch.expm1(lr32.double()) - lr32.double())
print("   log r     naive fp32        expm1 fp32       fp64 truth       naive rel.err   stable rel.err")
for a, n_, s_, t_ in zip(lr32.tolist(), naive.tolist(), stable.tolist(), truth.tolist()):
    print(f"   {a:7.0e}  {n_:+.4e}   {s_:+.4e}   {t_:+.4e}    {abs(n_-t_)/t_:9.2e}     {abs(s_-t_)/t_:9.2e}")
check("expm1 form: relative error < 1e-3 for |log r| >= 1e-4 (fp32)", float(((stable.double() - truth).abs() / truth)[:4].max()), 1e-3)
assert float(((naive.double() - truth).abs() / truth)[3]) > 0.1
check("naive form is wrong by >10% at |log r| = 1e-4 (fp32)", 0.0, 0.0)
# The expm1 form removes the cancellation inside exp(x)-1 but NOT the one in expm1(x) - x ~ x^2/2: its worst-case relative error is ~ 2u/|x|.
# Worst case over a dense random grid per decade (fp32, both signs), vs the 2u/|x|_min law; and a series fallback for small |x|.
u32 = 2.0 ** -24
gen_n = torch.Generator().manual_seed(11)
def k3_expm1(x): return torch.expm1(x) - x
def k3_hybrid(x):
    ser = x * x / 2 * (1 + x / 3 + x * x / 12)                              # Taylor of expm1(x)-x, error ~ x^5/120 relative x^3/60
    return torch.where(x.abs() < 1e-2, ser, k3_expm1(x))
print("   |x| decade        worst rel.err expm1 form   2u/|x|min     worst rel.err hybrid (series for |x|<1e-2)")
worst_ratio, worst_hybrid = 0.0, 0.0
for lo in (1e-1, 1e-2, 1e-3, 1e-4, 1e-5, 1e-6):
    hi = 10 * lo
    xg = torch.exp(torch.rand(400_000, generator=gen_n, dtype=torch.float64) * math.log(10) + math.log(lo)).float()
    xg = torch.cat([xg, -xg]); tr = torch.expm1(xg.double()) - xg.double()
    re = float(((k3_expm1(xg).double() - tr).abs() / tr).max()); rh = float(((k3_hybrid(xg).double() - tr).abs() / tr).max())
    law = 2 * u32 / lo
    worst_ratio = max(worst_ratio, re / law); worst_hybrid = max(worst_hybrid, rh)
    print(f"   [{lo:.0e},{hi:.0e}]      {re:12.3e}              {law:10.2e}     {rh:12.3e}")
check("expm1 form: worst-case relative error <= 2u/|x| (max ratio to the law)", worst_ratio, 1.0)
check("hybrid (Taylor for |x|<1e-2, expm1 form above): worst relative error over |x| in [1e-6, 1)", worst_hybrid, 1e-4)
print("done; checks:", len(CHECKS))