"""Sparse elimination mod p and projection onto error unknowns."""
import sys, time
from system import ERRDIMS, ERRORDER, err_offsets
from rewrite import context_matrix
from evalmap import eval_map, sfp


def to_int_rows(S, p):
    cols = {}
    nuc = []
    for r in S.rows:
        for c in r:
            if c[0] == "nu" and c not in cols:
                cols[c] = None
    nu_cols = sorted(cols, key=lambda c: (S.order.index(c[1]) if False else 0, str(c)))
    idx = {}
    for i, c in enumerate(cols):
        idx[c] = i
    N = len(idx)
    off, nerr = err_offsets()
    for n in ERRORDER:
        for i in range(ERRDIMS[n]):
            idx[("err", n, i)] = N + off[n] + i
    rows = []
    for r in S.rows:
        rows.append({idx[c]: int(v) % p for c, v in r.items() if int(v) % p})
    return rows, N, nerr


def echelon(rows, p, verbose=False):
    piv = {}
    t0 = time.time()
    for n, r in enumerate(rows):
        r = dict(r)
        while r:
            c = min(r)
            if c in piv:
                pr = piv[c]
                f = r[c]
                for cc, vv in pr.items():
                    nv = (r.get(cc, 0) - f * vv) % p
                    if nv:
                        r[cc] = nv
                    else:
                        r.pop(cc, None)
            else:
                inv = pow(r[c], p - 2, p)
                r = {cc: vv * inv % p for cc, vv in r.items()}
                piv[c] = r
                break
        if verbose and n % 5000 == 0:
            print(f"   elim {n}/{len(rows)} pivots {len(piv)} {time.time()-t0:.1f}s", file=sys.stderr)
    return piv


def err_constraints(piv, N):
    return {c - N: {cc - N: v for cc, v in r.items()} for c, r in piv.items() if c >= N}


def nullspace_mod(cons, n, p):
    """cons: dict pivot->row (reduced echelon not guaranteed); return basis of
    solutions of the constraint rows in F_p^n"""
    # full RREF
    rows = [dict(r) for r in cons.values()]
    piv = {}
    for r in rows:
        while r:
            c = min(r)
            if c in piv:
                pr = piv[c]
                f = r[c]
                for cc, vv in pr.items():
                    nv = (r.get(cc, 0) - f * vv) % p
                    if nv:
                        r[cc] = nv
                    else:
                        r.pop(cc, None)
            else:
                inv = pow(r[c], p - 2, p)
                r = {cc: vv * inv % p for cc, vv in r.items()}
                piv[c] = r
                break
    # back substitute
    for c in sorted(piv, reverse=True):
        r = piv[c]
        for c2 in list(piv):
            if c2 < c and c in piv[c2]:
                f = piv[c2][c]
                for cc, vv in r.items():
                    nv = (piv[c2].get(cc, 0) - f * vv) % p
                    if nv:
                        piv[c2][cc] = nv
                    else:
                        piv[c2].pop(cc, None)
    free = [j for j in range(n) if j not in piv]
    basis = []
    for f in free:
        v = [0] * n
        v[f] = 1
        for c, r in piv.items():
            v[c] = (-r.get(f, 0)) % p
        basis.append(v)
    return basis, piv


def check_vec(cons, vec, p):
    bad = 0
    for c, r in cons.items():
        s = sum(v * vec[j] for j, v in r.items()) % p
        if s:
            bad += 1
    return bad