"""JAX hard quantization with explicitly tested surrogate cotangents.

python -m pip install --pre "omnibias-binary==0.1.0a2" "omnibias-jax==0.5.0rc1"
"""
import math
import jax
jax.config.update("jax_enable_x64", True)
import jax.numpy as jnp
from omnibias.binary.jax.ops import (
    binarize, binarize01, heaviside, ternarize, kbit_quantize,
    riccati_tanh_derivative, riccati_sigmoid_derivative,
    surrogate_tower, surrogate_jet, curvature_corrected_slope, binarize_curvature,
)

z = jnp.array([-0.75, -0.25, 0.0, 0.25, 0.75], dtype=jnp.float64)
beta = 2.0
assert jnp.array_equal(binarize(z, beta), jnp.array([-1., -1., 1., 1., 1.]))
grad = jax.grad(lambda x: binarize(x, beta).sum())(z)
assert jnp.allclose(grad, beta*(1-jnp.tanh(beta*z)**2), rtol=1e-12, atol=1e-12)
assert jnp.array_equal(heaviside(z, beta), binarize01(z, beta))
s = jax.nn.sigmoid(beta*z)
step_grad = jax.grad(lambda x: heaviside(x, beta).sum())(z)
assert jnp.allclose(step_grad, beta*s*(1-s), rtol=1e-12, atol=1e-12)
assert jnp.array_equal(ternarize(z, beta, delta=0.5), jnp.array([-1., 0., 0., 0., 1.]))
q = kbit_quantize(z, bits=2, lo=-1., hi=1., beta=beta)
assert q.shape == z.shape and jnp.isfinite(q).all() and q.min() >= -1 and q.max() <= 1
q_grad = jax.grad(lambda x: kbit_quantize(x, bits=2, beta=beta).sum())(z)
assert jnp.isfinite(q_grad).all() and jnp.abs(q_grad).sum() > 0

t = jnp.tanh(z)
assert jnp.allclose(riccati_tanh_derivative(t, 1), 1-t*t)
s = jax.nn.sigmoid(z)
assert jnp.allclose(riccati_sigmoid_derivative(s, 1), s*(1-s))
tower = surrogate_tower(z, beta, order=3)
jet = surrogate_jet(z, beta, order=3)
assert tower.shape == jet.shape == (4, 5)
for order in range(4):
    assert jnp.allclose(jet[order]*math.factorial(order), tower[order], rtol=1e-12, atol=1e-12)
assert jnp.allclose(curvature_corrected_slope(z, beta, window=0.), tower[1])
curvature_grad = jax.grad(lambda x: binarize_curvature(x, beta).sum())(z)
assert jnp.allclose(curvature_grad, curvature_corrected_slope(z, beta), rtol=1e-12, atol=1e-12)
print("quantizers: (5,); surrogate tower/jet: (4, 5); surrogate gradients passed")