# sweep.py: for assemblies built by the page's lock.mjs (via dump.mjs), find the friction per contact
# at which the assembly first admits a non-zero tension state, and optionally save certificates on
# both sides of it. Uses HiGHS through scipy. The certificates are what the verifier checks; this
# solver is never trusted by the page.
#   python3 sweep.py N seed [--cert out.json]
import json, subprocess, sys, numpy as np
from scipy.optimize import linprog
from scipy.sparse import coo_matrix, vstack, csr_matrix
HERE = __file__.rsplit('/', 1)[0] or '.'
def load(N, seed, m=20, r=4):
    out = subprocess.run(['node', f'{HERE}/dump.mjs', str(N), str(seed), str(m), str(r)], capture_output=True, text=True, check=True).stdout
    return json.loads(out)
def sp(rows, n):
    I, J, V = [], [], []
    for k, row in enumerate(rows):
        for j, c in row: I.append(k); J.append(j); V.append(c)
    return csr_matrix(coo_matrix((V, (I, J)), shape=(len(rows), n)))
def solve(P, lam):
    n = P['n']; E = sp(P['eq'], n); D = sp(P['capD'], n); M = sp(P['capM'], n)
    Aub = vstack([D - lam * M, -D - lam * M, csr_matrix(np.ones((1, n)))]).tocsr()
    b = np.zeros(Aub.shape[0]); b[-1] = 1
    tries = [('highs', {}), ('highs-ds', {'presolve': False}), ('highs-ipm', {'presolve': False}),
             ('highs-ds', {'primal_feasibility_tolerance': 1e-10, 'dual_feasibility_tolerance': 1e-10})]
    for method, opts in tries:
        res = linprog(-np.ones(n), A_ub=Aub, b_ub=b, A_eq=E, b_eq=np.zeros(E.shape[0]), bounds=(0, None), method=method, options=opts)
        if res.status == 0: break
    if res.status != 0:
        # HiGHS occasionally reports an unknown status or a solve error on these degenerate homogeneous
        # problems near lambda = 2 on short fibres, under every option above. Last resort: GLPK in exact
        # rational arithmetic (glpsol --exact), which is slow but settles it.
        return glpk_exact(Aub, b, E, n, lam), None, E, D, M
    return -res.fun, res, E, D, M
def glpk_exact(Aub, b, E, n, lam):
    import tempfile, os, re
    lines = ['Maximize', ' obj: ' + ' + '.join(f'x{j}' for j in range(n)), 'Subject To']
    def row(name, r, sense, rhs):
        r = r.tocoo(); terms = ' '.join(f"{'-' if v < 0 else '+'} {abs(v):.17g} x{j}" for j, v in zip(r.col, r.data) if v != 0)
        if terms: lines.append(f' {name}: {terms} {sense} {rhs}')
    for i in range(Aub.shape[0]): row(f'u{i}', Aub[i], '<=', b[i])
    for i in range(E.shape[0]): row(f'e{i}', E[i], '=', 0)
    lines.append('End')
    d = tempfile.mkdtemp(); p = os.path.join(d, 'p.lp'); o = os.path.join(d, 'p.out')
    open(p, 'w').write('\n'.join(lines) + '\n')
    try: subprocess.run(['glpsol', '--exact', '--lp', p, '-o', o], capture_output=True, text=True, timeout=240)
    except subprocess.TimeoutExpired: raise RuntimeError(f'lambda={lam}: glpk exact did not finish in 240 s')
    txt = open(o).read(); m = re.search(r'Status:\s+(\S+)', txt); z = re.search(r'Objective:\s+obj = ([-0-9.e+]+)', txt)
    if not m or m.group(1) != 'OPTIMAL' or not z: raise RuntimeError(f'lambda={lam}: glpk exact also failed')
    return float(z.group(1))
class Partial(Exception):
    def __init__(self, lo, hi, why): super().__init__(why); self.lo, self.hi, self.why = lo, hi, why
def critical(P, tol=2e-4):
    # at lambda = 2 every cap |dT| <= (T_before + T_after) holds for any non-negative tensions, so the
    # assembly is trivially locked there; a threshold at 2 means friction never helped.
    lo, hi = 0.0, 2.0
    while (hi - lo) > tol * hi:
        mid = (lo + hi) / 2
        try: z = solve(P, mid)[0]
        except RuntimeError as e: raise Partial(lo, hi, str(e))
        if z > 1e-9: hi = mid
        else: lo = mid
    return lo, hi
def certs(P, lam_lo, lam_hi):
    z, res, E, D, M = solve(P, lam_hi)
    T = np.maximum(res.x, 0) / max(res.x)
    z0, r0, E, D, M = solve(P, lam_lo)
    assert z0 < 1e-9
    nc = D.shape[0]
    y = -r0.ineqlin.marginals; p = np.maximum(y[:nc], 0); q = np.maximum(y[nc:2*nc], 0); u = -r0.eqlin.marginals
    g = E.T @ u + (D - lam_lo * M).T @ p + (-D - lam_lo * M).T @ q
    return {'lockedAt': lam_hi, 'T': [float(f'{t:.12g}') for t in T], 'slidesAt': lam_lo, 'u': [float(f'{x:.12g}') for x in u], 'p': [float(f'{x:.12g}') for x in p], 'q': [float(f'{x:.12g}') for x in q], 'gmin': float(g.min())}
if __name__ == '__main__':
    N, seed = int(sys.argv[1]), int(sys.argv[2])
    P = load(N, seed)
    try:
        lo, hi = critical(P)
        rec = {'N': N, 'seed': seed, 'lamLo': lo, 'lamHi': hi, 'NlamC': N * (lo + hi) / 2}
    except Partial as e:
        # every solver gave up at some lambda inside the bracket: keep the bracket it had established,
        # which is still true (lo was shown to slide, hi to lock), and say so
        print(json.dumps({'N': N, 'seed': seed, 'lamLo': e.lo, 'lamHi': e.hi, 'NlamC': N * (e.lo + e.hi) / 2, 'partial': True, 'why': e.why})); sys.exit(0)
    if '--cert' in sys.argv:
        c = certs(P, lo * 0.99, hi * 1.01)
        rec.update({'cert': c})
        json.dump(rec, open(sys.argv[sys.argv.index('--cert') + 1], 'w'))
        print(json.dumps({k: v for k, v in rec.items() if k != 'cert'}), 'gmin', c['gmin'])
    else:
        print(json.dumps(rec))