import math
from catopt_grid.core import LocalProblem
from catopt_grid.solver import admm_lite


def test_admm_lite_two_1d_problems_converge_to_same_solution():
    # Problem 1: minimize 0.5*Q1*x^2 + c1*x with Q1=2, c1=-6  -> unconstrained minimizer x*=3
    Q1 = [[2.0]]
    c1 = [-6.0]
    p1 = LocalProblem(id="p1", domain="test", n=1, Q=Q1, c=c1)

    # Problem 2: minimize 0.5*Q2*x^2 + c2*x with Q2=3, c2=-9  -> unconstrained minimizer x*=3
    Q2 = [[3.0]]
    c2 = [-9.0]
    p2 = LocalProblem(id="p2", domain="test", n=1, Q=Q2, c=c2)

    z, history = admm_lite([p1, p2], rho=1.0, max_iters=200, tol=1e-9)
    # The converged shared variable should be close to 3.0
    assert abs(z[0] - 3.0) < 1e-6
    assert len(history) > 0