"""Hard quantization, analytic surrogate derivatives, and an annealing schedule.

python -m pip install --pre "omnibias-binary==0.1.0a2" "omnibias-torch==0.5.0rc1"
"""
import math
import torch
from omnibias.binary import BetaAnnealScheduler
from omnibias.binary.torch.ops import (
    binarize, binarize01, heaviside, ternarize, kbit_quantize,
    riccati_tanh_derivative, riccati_sigmoid_derivative,
    surrogate_tower, surrogate_jet, curvature_corrected_slope, binarize_curvature,
)

z = torch.tensor([-0.75, -0.25, 0.0, 0.25, 0.75], dtype=torch.float64, requires_grad=True)
beta = 2.0
hard = binarize(z, beta)
torch.testing.assert_close(hard, torch.tensor([-1., -1., 1., 1., 1.], dtype=z.dtype))
# This is a specified surrogate gradient, not the classical derivative of sign.
(grad,) = torch.autograd.grad(hard.sum(), z)
torch.testing.assert_close(grad, beta * (1 - torch.tanh(beta*z).square()))
step = heaviside(z, beta)
torch.testing.assert_close(step, binarize01(z, beta))
(step_grad,) = torch.autograd.grad(step.sum(), z)
s = torch.sigmoid(beta*z)
torch.testing.assert_close(step_grad, beta*s*(1-s))
torch.testing.assert_close(ternarize(z, beta, delta=0.5), torch.tensor([-1., 0., 0., 0., 1.], dtype=z.dtype))
quantized = kbit_quantize(z, bits=2, lo=-1.0, hi=1.0, beta=beta)
assert quantized.shape == z.shape and torch.isfinite(quantized).all()
assert quantized.min() >= -1 and quantized.max() <= 1
(quantized_grad,) = torch.autograd.grad(quantized.sum(), z)
assert torch.isfinite(quantized_grad).all() and quantized_grad.abs().sum() > 0

t = torch.tanh(z)
torch.testing.assert_close(riccati_tanh_derivative(t, 1), 1-t.square())
s = torch.sigmoid(z)
torch.testing.assert_close(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):
    torch.testing.assert_close(jet[order]*math.factorial(order), tower[order])
torch.testing.assert_close(curvature_corrected_slope(z, beta, window=0.0), tower[1])
(curvature_grad,) = torch.autograd.grad(binarize_curvature(z, beta).sum(), z)
torch.testing.assert_close(curvature_grad, curvature_corrected_slope(z, beta))

schedule = BetaAnnealScheduler(1.0, 4.0, num_steps=2, schedule="linear")
assert [schedule.step() for _ in range(3)] == [1.0, 2.5, 4.0]
assert schedule.current_step == 3 and schedule.value(100) == 4.0
schedule.reset()
assert schedule.current_step == 0 and schedule.step() == 1.0
print("quantizers: (5,); surrogate tower/jet: (4, 5); gradients and schedule passed")