"""
MSHT20 Hessian via central second differences of likelihood VALUES, eager mode.
Rationale: every XLA-compiled route (fused jax.hessian, HVP columns, jitted
gradients) hit pathological compile times on this 52-param likelihood.
Value-only finite differences in eager mode compile NOTHING: ~2*n^2 + 2n
evaluations. Accuracy ~1e-4 relative — far beyond eigenvector-band needs
(tolerance T=1). Symmetric by construction.
"""
import sys
import jax
import jax.numpy as jnp


def hessian_fd_values(f, *a, **k):
    def hess(p, *args):
        fp = lambda q: f(q, *args)
        n = p.shape[0]
        h = 5e-4 * jnp.maximum(1.0, jnp.abs(p))
        H = [[0.0] * n for _ in range(n)]
        f0 = float(fp(p))
        # diagonal
        for i in range(n):
            e = jnp.zeros_like(p).at[i].set(h[i])
            H[i][i] = (float(fp(p + e)) - 2.0 * f0 + float(fp(p - e))) / float(h[i]) ** 2
        # off-diagonal
        for i in range(n):
            for j in range(i + 1, n):
                ei = jnp.zeros_like(p).at[i].set(h[i])
                ej = jnp.zeros_like(p).at[j].set(h[j])
                val = (float(fp(p + ei + ej)) - float(fp(p + ei - ej))
                       - float(fp(p - ei + ej)) + float(fp(p - ei - ej))) / (4.0 * float(h[i]) * float(h[j]))
                H[i][j] = H[j][i] = val
        print(f"[fdval] Hessian assembled: {n}x{n}, f0={f0:.6f}", flush=True)
        return jnp.array(H)

    return hess


jax.hessian = hessian_fd_values

from msht20_model.app import main  # noqa: E402

if __name__ == "__main__":
    sys.argv = ["msht20_exe"] + sys.argv[1:]
    main()