"""Small positive-definite, nonnegative contact QP; standard library only.
Run from notes/series_html: python3 examples/convex_contact.py
For study: exponential face enumeration, intended for at most 10 contacts.
"""
from itertools import combinations
import json


def solve(A, b):
    """Pivoted Gaussian elimination; input arrays are preserved."""
    n = len(b)
    rows = [list(map(float, row)) + [float(value)] for row, value in zip(A, b)]
    for k in range(n):
        pivot = max(range(k, n), key=lambda i: abs(rows[i][k]))
        rows[k], rows[pivot] = rows[pivot], rows[k]
        if abs(rows[k][k]) < 1e-14:
            raise ValueError('singular system')
        d = rows[k][k]
        rows[k] = [x/d for x in rows[k]]
        for i in range(n):
            if i != k:
                d = rows[i][k]
                rows[i] = [x-d*y for x, y in zip(rows[i], rows[k])]
    return [row[-1] for row in rows]


def candidate(H, b, support):
    n = len(b)
    f = [0.0] * n
    block = [[H[i][j] for j in support] for i in support]
    values = solve(block, [-b[i] for i in support])
    for i, value in zip(support, values):
        f[i] = value
    w = [sum(H[i][j]*f[j] for j in range(n)) + b[i]
         for i in range(n)]
    return f, w


def objective(H, b, f):
    return sum(f[i]*(sum(H[i][j]*f[j] for j in range(len(f)))/2+b[i])
               for i in range(len(f)))


def solve_contact(H, b):
    """H must be symmetric positive definite. Returns force and KKT residual."""
    n = len(b)
    if n > 10:
        raise ValueError('enumeration is limited to 10 variables')
    records = []
    for count in range(n+1):
        for support in combinations(range(n), count):
            f, w = candidate(H, b, support)
            valid = min(f, default=0) >= -1e-10 and min(w, default=0) >= -1e-10
            records.append(dict(support=support, force=f, residual=w, valid=valid,
                                objective=objective(H, b, f)))
    accepted = [r for r in records if r['valid']]
    if not accepted:
        raise ValueError('no feasible KKT candidate; check H')
    return dict(min(accepted, key=lambda r: r['objective']), candidates=records)


if __name__ == '__main__':
    H, b = [[1, -1], [-1, 2]], [-10, 15]
    result = solve_contact(H, b)
    free = solve(H, [-x for x in b])
    clipped = [max(0, x) for x in free]
    w = [sum(H[i][j]*clipped[j] for j in range(2))+b[i] for i in range(2)]
    assert max(abs(a-z) for a, z in zip(result['force'], [10, 0])) < 1e-12
    print(json.dumps({'solution': result, 'component_clipping': {
        'force': clipped, 'residual': w,
        'complementarity': max(abs(x*y) for x, y in zip(clipped, w))}}, indent=2))