import numpy as _np
from catopt_grid import LocalProblem, ADMMLiteSolver


def test_admm_lite_two_agents_quadratic_convergence():
    # Two 1-D local problems with simple quadratics.
    # f1(x) = 0.5 * 2 * x^2 + (-6) * x  => Q1=2, b1=-6, bounds [0, 5]
    Q1 = _np.array([[2.0]])
    b1 = _np.array([-6.0])
    lp1 = LocalProblem(n=1, Q=Q1, b=b1, lb=_np.array([0.0]), ub=_np.array([5.0]), name="lp1")

    # f2(x) = 0.5 * 1 * x^2 + (-2) * x  => Q2=1, b2=-2, bounds [0, 5]
    Q2 = _np.array([[1.0]])
    b2 = _np.array([-2.0])
    lp2 = LocalProblem(n=1, Q=Q2, b=b2, lb=_np.array([0.0]), ub=_np.array([5.0]), name="lp2")

    solver = ADMMLiteSolver([lp1, lp2], rho=1.0, max_iter=1000, tol=1e-6)
    z, xs, us = solver.solve()

    # Check that the consensus value is within bounds and that x_i are close to each other
    assert z.shape == (1,)
    assert all(0.0 <= xi[0] <= 5.0 for xi in xs)
    # Since both x_i converge to the same z, their difference should be small
    diffs = [_np.linalg.norm(xs[i] - xs[0]) for i in range(1, len(xs))]
    assert max(diffs) < 1e-4
    # And u_i should be finite
    for ui in us:
        assert _np.all(_np.isfinite(ui))