"""GaussNewton is functional: step returns a new flat parameter vector."""
import torch
from omnibias.torch.optim import GaussNewton

x = torch.tensor([-1.0, 0.0, 1.0], dtype=torch.float64)
y = 2 * x + 1

def residual(parameters):
    return parameters[0] * x + parameters[1] - y

parameters = torch.zeros(2, dtype=torch.float64)
optimizer = GaussNewton(solver="qr", damping=1e-3)
before = residual(parameters).square().mean()
for _ in range(4):
    parameters, info = optimizer.step(residual, parameters)
after = residual(parameters).square().mean()
assert after < before
assert torch.isfinite(parameters).all()
torch.testing.assert_close(parameters, torch.tensor([2.0, 1.0], dtype=torch.float64),
                           rtol=1e-6, atol=1e-6)
print("Recovered slope and intercept:", parameters.detach().tolist())