"""Two stacked unit-mass blocks. Forces in N, accelerations in m/s²."""
from math import isclose


def pgs_stack(pull=0.0, regularization=0.0, sweeps=16, initial=(0.0, 0.0)):
    """Yield the initial state and each row update at one fixed instant."""
    f = [max(0.0, x) for x in initial]
    H = ((1 + regularization, -1), (-1, 2 + regularization))
    b = (-10, pull)

    def record(row, trial=None):
        a = (-10 + f[0] - f[1], -10 + pull + f[1])
        w = tuple(sum(H[i][j] * f[j] for j in range(2)) + b[i] for i in range(2))
        error = max(abs(f[i] - max(0.0, f[i] - w[i] / H[i][i])) for i in range(2))
        return dict(row=row, force=tuple(f), acceleration=a, residual=w,
                    error=error, trial=trial)

    yield record(None)
    for _ in range(sweeps):
        for i in range(2):
            other = sum(H[i][j] * f[j] for j in range(2) if j != i)
            trial = -(b[i] + other) / H[i][i]
            f[i] = max(0.0, trial)
            yield record(i, trial)


if __name__ == '__main__':
    for row in pgs_stack(sweeps=2):
        print(row)
    resting = list(pgs_stack(sweeps=40))[-1]
    assert all(isclose(x, y, abs_tol=1e-9) for x, y in zip(resting['force'], (20, 10)))
    separating = list(pgs_stack(pull=15, sweeps=1))[-1]
    assert separating['force'] == (10, 0)
    assert separating['acceleration'] == (0, 5)
    soft = list(pgs_stack(regularization=.1, sweeps=40))[-1]
    assert all(isclose(x, y, abs_tol=1e-9) for x, y in zip(soft['force'], (2100/131, 1000/131)))