/*
 * lra.c -- Dutertre / de Moura "General Simplex for SMT".
 *
 * Internally every constraint is represented as a bound on a single
 * variable. Multi-variable constraints get a fresh slack variable s
 * and the row  s - sum_k c_k x_k = 0  enters the tableau. The atom
 * becomes a bound on s.
 *
 * The tableau is stored as a dense (n_basic x n_vars) matrix because
 * the tutorial dimensions are tiny and dense code is easier to follow.
 * A real solver would store rows sparsely.
 *
 * State:
 *
 *   For each variable v:
 *     - is_basic[v]: 1 if v is currently a basic variable
 *     - row_of [v]:  if basic, the index into rows[] of its row;
 *                    if non-basic, the index into cols[] (== v)
 *     - val[v]:      current delta-rational value
 *     - has_lo, lo, lo_atom: lower bound (if has_lo) and the atom
 *                            that asserted it
 *     - has_hi, hi, hi_atom: upper bound
 *
 *   Rows: each row corresponds to a basic variable b and is of the
 *   form  b - sum_j a_j * x_j = 0  where j ranges over non-basics.
 *   We store a[i][j] for the coefficient of variable j in row i; if
 *   variable b is the basic variable for row i, a[i][b] is set to -1
 *   (representing the "-b" term moved to the LHS); for any other
 *   basic variable a[i][k] = 0. For non-basic variables, a[i][j] is
 *   the actual coefficient.
 *
 *   Bound trail: each bound assertion saves the previous (has, val,
 *   atom) triple of the relevant variable + side so unassign can
 *   restore it. Independent of the merge_trail in EUF.
 */
#include "lra.h"

#include <assert.h>
#include <stdio.h>
#include <stdlib.h>
#include <string.h>

/* ====== bound trail entry ====== */
typedef struct {
    int   var;
    int   side;       /* 0 = lower, 1 = upper */
    int   had;        /* old has_lo / has_hi */
    Drat  old_val;
    int   old_atom;
} BoundUndo;

/* ====== atom record ====== */
typedef struct {
    int  var;         /* the slack/variable being bounded */
    int  op;          /* LRA_OP_LE or LRA_OP_LT */
    Rat  rhs;
} AtomRec;

/* ====== solver state ====== */
struct LRA {
    /* Variables. */
    char **var_name;
    int    n_vars, cap_vars;

    int  *is_basic;     /* 0/1 per var */
    int  *row_of;       /* basic var: index in rows[]; non-basic: == var */
    Drat *val;          /* current delta-rational value */

    int  *has_lo;       Drat *lo;  int *lo_atom;
    int  *has_hi;       Drat *hi;  int *hi_atom;

    /* Tableau rows. rows[i] is the basic variable of the i-th row.
     * a[i * n_vars + j] is the coefficient of variable j in row i.
     * For row i: basic_var = rows[i]; a[i][basic_var] is unused (we
     * implicitly treat the basic side as -1, and reconstruct values
     * from a[i][non-basic] * val[non-basic]). */
    int  *rows;
    int   n_rows, cap_rows;
    Rat  *a;            /* (cap_rows * cap_vars) */

    /* Atom dispatch table. atom -> AtomRec or {var=-1} if not LRA. */
    AtomRec *atoms;
    int     *atom_assigned;  /* 0/1: is this atom currently on the trail? */
    int      cap_atoms;

    /* SAT-trail mirror: same shape as EUF -- record which atoms have
     * been assigned, and at what point each was, so unassign can
     * pop back the bound trail. */
    int  *asserts;          /* lits in order */
    int   n_asserts, cap_asserts;
    int  *pre_bound;        /* bound-trail size before assert i */
    int  *pre_propq;        /* propq size before assert i */
    int   cap_pre;

    /* Bound trail. */
    BoundUndo *btrail;
    int        n_btrail, cap_btrail;

    /* Pending check work. */
    int check_pos;

    /* Propagation queue (mirror of EUF). */
    struct PropEntry {
        int  lit;
        int *reason;
        int  n_reason;
        int  reason_cap;
    } *propq;
    int n_propq, cap_propq, propq_head;

    /* Conflict scratch. */
    int *conf_buf;
    int  cap_conf, conf_n;

    /* Stats. */
    long n_pivots;
    long n_propagations;
};

/* ====== array growth helpers ====== */
static void grow_int(int **buf, int *cap, int want) {
    if (want <= *cap) return;
    int nc = *cap ? *cap * 2 : 8;
    while (nc < want) nc *= 2;
    *buf = realloc(*buf, (size_t)nc * sizeof(int));
    *cap = nc;
}

/* Tableau cell accessor. Always grown together with rows/vars. */
static Rat *A(LRA *lra, int row, int col) {
    return &lra->a[row * lra->cap_vars + col];
}

static void grow_vars(LRA *lra, int needed) {
    if (needed <= lra->cap_vars) return;
    int oldc = lra->cap_vars;
    int nc = oldc ? oldc * 2 : 8;
    while (nc < needed) nc *= 2;

    lra->var_name = realloc(lra->var_name, (size_t)nc * sizeof(char*));
    lra->is_basic = realloc(lra->is_basic, (size_t)nc * sizeof(int));
    lra->row_of   = realloc(lra->row_of,   (size_t)nc * sizeof(int));
    lra->val      = realloc(lra->val,      (size_t)nc * sizeof(Drat));
    lra->has_lo   = realloc(lra->has_lo,   (size_t)nc * sizeof(int));
    lra->has_hi   = realloc(lra->has_hi,   (size_t)nc * sizeof(int));
    lra->lo       = realloc(lra->lo,       (size_t)nc * sizeof(Drat));
    lra->hi       = realloc(lra->hi,       (size_t)nc * sizeof(Drat));
    lra->lo_atom  = realloc(lra->lo_atom,  (size_t)nc * sizeof(int));
    lra->hi_atom  = realloc(lra->hi_atom,  (size_t)nc * sizeof(int));

    /* Resize tableau. Need to reshape because column dimension changed.
     * Allocate new buffer, copy row by row. */
    Rat *na = malloc((size_t)lra->cap_rows * (size_t)nc * sizeof(Rat));
    for (int i = 0; i < lra->cap_rows * nc; i++) na[i] = rat_int(0);
    for (int i = 0; i < lra->n_rows; i++) {
        for (int j = 0; j < oldc; j++) {
            na[i * nc + j] = lra->a[i * oldc + j];
        }
    }
    free(lra->a);
    lra->a = na;

    for (int i = oldc; i < nc; i++) {
        lra->var_name[i] = NULL;
        lra->is_basic[i] = 0;
        lra->row_of  [i] = i;
        lra->val     [i] = DRAT_ZERO;
        lra->has_lo  [i] = 0;
        lra->has_hi  [i] = 0;
        lra->lo      [i] = DRAT_ZERO;
        lra->hi      [i] = DRAT_ZERO;
        lra->lo_atom [i] = 0;
        lra->hi_atom [i] = 0;
    }
    lra->cap_vars = nc;
}

static void grow_rows(LRA *lra, int needed) {
    if (needed <= lra->cap_rows) return;
    int oldc = lra->cap_rows;
    int nc = oldc ? oldc * 2 : 4;
    while (nc < needed) nc *= 2;

    lra->rows = realloc(lra->rows, (size_t)nc * sizeof(int));
    Rat *na = malloc((size_t)nc * (size_t)lra->cap_vars * sizeof(Rat));
    for (int i = 0; i < nc * lra->cap_vars; i++) na[i] = rat_int(0);
    for (int i = 0; i < lra->n_rows; i++) {
        for (int j = 0; j < lra->cap_vars; j++) {
            na[i * lra->cap_vars + j] = lra->a[i * lra->cap_vars + j];
        }
    }
    free(lra->a);
    lra->a = na;
    lra->cap_rows = nc;
}

static void grow_atoms(LRA *lra, int atom) {
    if (atom < lra->cap_atoms) return;
    int nc = lra->cap_atoms ? lra->cap_atoms * 2 : 16;
    while (nc <= atom) nc *= 2;
    lra->atoms = realloc(lra->atoms, (size_t)nc * sizeof(AtomRec));
    lra->atom_assigned = realloc(lra->atom_assigned, (size_t)nc * sizeof(int));
    for (int i = lra->cap_atoms; i < nc; i++) {
        lra->atoms[i].var = -1;
        lra->atoms[i].op  = 0;
        lra->atoms[i].rhs = rat_int(0);
        lra->atom_assigned[i] = 0;
    }
    lra->cap_atoms = nc;
}

static void grow_pre(LRA *lra, int needed) {
    if (needed <= lra->cap_pre) return;
    int nc = lra->cap_pre ? lra->cap_pre * 2 : 16;
    while (nc < needed) nc *= 2;
    lra->pre_bound = realloc(lra->pre_bound, (size_t)nc * sizeof(int));
    lra->pre_propq = realloc(lra->pre_propq, (size_t)nc * sizeof(int));
    lra->cap_pre = nc;
}

/* ====== lifecycle ====== */
LRA *lra_new(void) {
    LRA *lra = calloc(1, sizeof(LRA));
    return lra;
}

void lra_free(LRA *lra) {
    if (!lra) return;
    for (int i = 0; i < lra->n_vars; i++) free(lra->var_name[i]);
    free(lra->var_name);
    free(lra->is_basic); free(lra->row_of); free(lra->val);
    free(lra->has_lo); free(lra->has_hi);
    free(lra->lo); free(lra->hi);
    free(lra->lo_atom); free(lra->hi_atom);
    free(lra->rows); free(lra->a);
    free(lra->atoms);
    free(lra->atom_assigned);
    free(lra->asserts);
    free(lra->pre_bound); free(lra->pre_propq);
    free(lra->btrail);
    for (int i = 0; i < lra->cap_propq; i++) free(lra->propq[i].reason);
    free(lra->propq);
    free(lra->conf_buf);
    free(lra);
}

/* ====== variable allocation ====== */
int lra_mk_var(LRA *lra, const char *name) {
    grow_vars(lra, lra->n_vars + 1);
    int v = lra->n_vars++;
    size_t L = strlen(name);
    lra->var_name[v] = malloc(L + 1);
    memcpy(lra->var_name[v], name, L + 1);
    /* All fresh fields zeroed by grow_vars. */
    return v;
}

/* ====== row math ======
 * value of basic variable b on row i:
 *   val[b] = sum over non-basics j  a[i][j] * val[j]
 * (No constant term: constants are eliminated by the slack-variable
 * trick at row construction, since all RHS goes to the bound on the
 * slack.) */
static Drat recompute_basic(const LRA *lra, int row_i) {
    Drat sum = DRAT_ZERO;
    int basic = lra->rows[row_i];
    for (int j = 0; j < lra->n_vars; j++) {
        if (lra->is_basic[j]) continue;
        if (j == basic)        continue;  /* shouldn't happen */
        Rat c = lra->a[row_i * lra->cap_vars + j];
        if (rat_is_zero(c)) continue;
        sum = drat_add(sum, drat_mul_rat(lra->val[j], c));
    }
    return sum;
}

/* ====== bound tracking ====== */
static void push_bound_undo(LRA *lra, int var, int side,
                            int had, Drat old_val, int old_atom)
{
    if (lra->n_btrail == lra->cap_btrail) {
        lra->cap_btrail = lra->cap_btrail ? lra->cap_btrail * 2 : 16;
        lra->btrail = realloc(lra->btrail,
                              (size_t)lra->cap_btrail * sizeof(BoundUndo));
    }
    BoundUndo u = { var, side, had, old_val, old_atom };
    lra->btrail[lra->n_btrail++] = u;
}

/* Tighten lower bound on variable v to (newval, atom).
 * Returns 1 if changed, 0 if newval was no tighter. */
static int tighten_lo(LRA *lra, int v, Drat newval, int atom) {
    if (lra->has_lo[v] && drat_cmp(newval, lra->lo[v]) <= 0) return 0;
    push_bound_undo(lra, v, 0, lra->has_lo[v], lra->lo[v], lra->lo_atom[v]);
    lra->has_lo [v] = 1;
    lra->lo     [v] = newval;
    lra->lo_atom[v] = atom;
    return 1;
}

static int tighten_hi(LRA *lra, int v, Drat newval, int atom) {
    if (lra->has_hi[v] && drat_cmp(newval, lra->hi[v]) >= 0) return 0;
    push_bound_undo(lra, v, 1, lra->has_hi[v], lra->hi[v], lra->hi_atom[v]);
    lra->has_hi [v] = 1;
    lra->hi     [v] = newval;
    lra->hi_atom[v] = atom;
    return 1;
}

/* When a non-basic variable's value is changed (e.g. forced into its
 * bound), propagate the change to all basic variables that include it
 * in their row. */
static void update_value_nonbasic(LRA *lra, int v, Drat new_val) {
    Drat delta = drat_sub(new_val, lra->val[v]);
    lra->val[v] = new_val;
    for (int i = 0; i < lra->n_rows; i++) {
        int basic = lra->rows[i];
        Rat c = lra->a[i * lra->cap_vars + v];
        if (rat_is_zero(c)) continue;
        lra->val[basic] = drat_add(lra->val[basic], drat_mul_rat(delta, c));
    }
}

/* Force a non-basic variable into its bound if it's out. */
static void clamp_nonbasic(LRA *lra, int v) {
    if (lra->is_basic[v]) return;
    if (lra->has_lo[v] && drat_cmp(lra->val[v], lra->lo[v]) < 0) {
        update_value_nonbasic(lra, v, lra->lo[v]);
    } else if (lra->has_hi[v] && drat_cmp(lra->val[v], lra->hi[v]) > 0) {
        update_value_nonbasic(lra, v, lra->hi[v]);
    }
}

/* ====== atom registration ====== */
void lra_register_atom(LRA *lra, int atom, int op,
                       const LraTerm *terms, int n_terms,
                       Rat rhs)
{
    grow_atoms(lra, atom);

    /* Single-variable atom with positive coefficient: no slack needed.
     * Divide rhs by the coefficient and store the bound directly on
     * the variable. Negative coefficient flips inequality direction
     * (LE <-> GE), which we can't represent as one polarity of an LE
     * atom, so fall through to slack form. */
    if (n_terms == 1) {
        Rat c = terms[0].coeff;
        if (!rat_is_zero(c) && rat_sgn(c) > 0) {
            int v = terms[0].var;
            Rat new_rhs = rat_div(rhs, c);
            lra->atoms[atom].var = v;
            lra->atoms[atom].op  = op;
            lra->atoms[atom].rhs = new_rhs;
            return;
        }
    }

    /* General case: introduce a slack s with row  s = sum c_k * x_k.
     * Internally stored as the equation moved to one side; the
     * column for s in this row is implicit (the basic variable).
     * For all non-basics j, a[row][j] is the coefficient. */
    char nm[32];
    snprintf(nm, sizeof nm, "s%d", lra->n_vars);
    int s = lra_mk_var(lra, nm);

    grow_rows(lra, lra->n_rows + 1);
    int ri = lra->n_rows++;
    lra->rows[ri] = s;
    lra->is_basic[s] = 1;
    lra->row_of  [s] = ri;

    /* Zero the row (was zeroed by grow_rows on first growth, but
     * subsequent reuse needs an explicit clear). */
    for (int j = 0; j < lra->cap_vars; j++)
        lra->a[ri * lra->cap_vars + j] = rat_int(0);

    /* Fill in the coefficients of the input terms. They may refer
     * to basic variables; if so, substitute the basic var's row in.
     * For tutorial simplicity we just record raw coefficients on
     * possibly-basic variables and rely on recompute_basic to use
     * only non-basics. Actually that's wrong: rows must reference
     * non-basics only.
     *
     * Substitution: if term refers to basic variable b, replace it
     * by sum_j a[row_of[b]][j] * x_j. */
    for (int k = 0; k < n_terms; k++) {
        int   v = terms[k].var;
        Rat   c = terms[k].coeff;
        if (lra->is_basic[v]) {
            int rb = lra->row_of[v];
            for (int j = 0; j < lra->n_vars; j++) {
                if (lra->is_basic[j]) continue;
                Rat aij = lra->a[rb * lra->cap_vars + j];
                if (rat_is_zero(aij)) continue;
                Rat add = rat_mul(c, aij);
                lra->a[ri * lra->cap_vars + j] =
                    rat_add(lra->a[ri * lra->cap_vars + j], add);
            }
        } else {
            lra->a[ri * lra->cap_vars + v] =
                rat_add(lra->a[ri * lra->cap_vars + v], c);
        }
    }

    /* Initial value of s. */
    lra->val[s] = recompute_basic(lra, ri);

    lra->atoms[atom].var = s;
    lra->atoms[atom].op  = op;
    lra->atoms[atom].rhs = rhs;
}

/* ====== assert / unassert ====== */
static void apply_bound(LRA *lra, int lit) {
    int atom = lit > 0 ? lit : -lit;
    if (atom >= lra->cap_atoms || lra->atoms[atom].var < 0) return;
    AtomRec *ar = &lra->atoms[atom];
    int  v = ar->var;
    Rat  r = ar->rhs;

    /* Determine which bound and whether strict.
     * Convention:
     *   pos LE: v <= r            -> hi := r          (weak)
     *   neg LE: v > r             -> lo := r + eps    (strict)
     *   pos LT: v <  r            -> hi := r - eps    (strict)
     *   neg LT: v >= r            -> lo := r          (weak)
     */
    Drat bound;
    int  side; /* 0 lo, 1 hi */
    if (lit > 0) {
        side = 1;
        if (ar->op == LRA_OP_LE) bound = drat_rat(r);
        else                     bound = drat_make(r, rat_int(-1));
    } else {
        side = 0;
        if (ar->op == LRA_OP_LE) bound = drat_make(r, rat_int(1));
        else                     bound = drat_rat(r);
    }

    if (side == 1) tighten_hi(lra, v, bound, lit);
    else           tighten_lo(lra, v, bound, lit);

    /* Make sure non-basics still sit inside their bounds. */
    if (!lra->is_basic[v]) clamp_nonbasic(lra, v);
}

void lra_assign(void *data, int lit) {
    LRA *lra = (LRA*)data;
    int atom = lit > 0 ? lit : -lit;
    if (atom >= lra->cap_atoms || lra->atoms[atom].var < 0) {
        /* Not an LRA atom; nothing to do here. The outer dispatch
         * shouldn't have called us, but stay safe. */
        return;
    }
    grow_int(&lra->asserts, &lra->cap_asserts, lra->n_asserts + 1);
    grow_pre(lra, lra->n_asserts + 1);
    lra->pre_bound[lra->n_asserts] = lra->n_btrail;
    lra->pre_propq[lra->n_asserts] = lra->n_propq;
    lra->asserts  [lra->n_asserts++] = lit;
    lra->atom_assigned[atom] = 1;
    apply_bound(lra, lit);
}

static void prop_free(struct PropEntry *p) {
    free(p->reason);
    p->reason = NULL;
    p->reason_cap = 0;
    p->n_reason = 0;
    p->lit = 0;
}

static void undo_one_bound(LRA *lra) {
    assert(lra->n_btrail > 0);
    BoundUndo u = lra->btrail[--lra->n_btrail];
    if (u.side == 0) {
        lra->has_lo [u.var] = u.had;
        lra->lo     [u.var] = u.old_val;
        lra->lo_atom[u.var] = u.old_atom;
    } else {
        lra->has_hi [u.var] = u.had;
        lra->hi     [u.var] = u.old_val;
        lra->hi_atom[u.var] = u.old_atom;
    }
}

void lra_unassign(void *data, int lit) {
    LRA *lra = (LRA*)data;
    int atom = lit > 0 ? lit : -lit;
    if (atom >= lra->cap_atoms || lra->atoms[atom].var < 0) return;
    assert(lra->n_asserts > 0);
    lra->n_asserts--;
    assert(lra->asserts[lra->n_asserts] == lit);
    lra->atom_assigned[atom] = 0;

    if (lra->check_pos > lra->n_asserts) {
        lra->check_pos = lra->n_asserts;
        int tb = lra->pre_bound[lra->n_asserts];
        while (lra->n_btrail > tb) undo_one_bound(lra);
        int tq = lra->pre_propq[lra->n_asserts];
        while (lra->n_propq  > tq) prop_free(&lra->propq[--lra->n_propq]);
        if (lra->propq_head > lra->n_propq) lra->propq_head = lra->n_propq;
    }
}

/* ====== pivoting ======
 * Swap basic variable b (in row ri) with non-basic variable n.
 * Before: row ri encodes  b = a_n * n + sum_{j != n, non-basic} a_j * j
 * After:  row ri encodes  n = (1/a_n) * b - sum_{j != n, non-basic} (a_j/a_n) * j
 * And every other row that mentions n needs to substitute the new
 * expression for n in.
 */
static void pivot(LRA *lra, int ri, int n) {
    int b = lra->rows[ri];
    Rat an = *A(lra, ri, n);
    assert(!rat_is_zero(an));

    /* Step 1: rewrite row ri. Divide row by a_n (this makes the
     * coefficient of n equal to 1), then negate everything (the
     * variable now on the LHS is n, originally was -b... handled by
     * sign convention).
     *
     * Concretely: the row currently represents
     *   b - sum_j a_j * j = 0
     * (where j ranges over non-basics including n).
     * Solving for n:
     *   n = b/a_n - sum_{j != n} (a_j / a_n) * j
     *
     * After the pivot, the row's basic var is n, and the coefficients
     * stored at positions other than n are the coefficients of the
     * RHS expression in n's solved form, EXCEPT we store -1 for b
     * (representing the moved -b term)... we just store directly.
     *
     * The convention I'll use here: a[ri][j] is the coefficient of
     * non-basic j in the expression for the basic var of row ri.
     * So for the rewritten row:
     *   a[ri][b] = 1/a_n        (b became non-basic)
     *   a[ri][j] = -a_j/a_n     (j != n, non-basic, j != b)
     *   a[ri][n] is unused (n is now basic)
     */
    for (int j = 0; j < lra->n_vars; j++) {
        if (j == n) continue;
        Rat aj = lra->a[ri * lra->cap_vars + j];
        if (j == b) {
            lra->a[ri * lra->cap_vars + j] = rat_div(rat_int(1), an);
        } else {
            /* Want a[ri][j] := -a_j / a_n  in the rewritten row. */
            lra->a[ri * lra->cap_vars + j] = rat_neg(rat_div(aj, an));
        }
    }
    /* Wait -- but the original row had b on the LHS, not -b. Let me
     * re-derive. Convention: a_ij stored in row i for non-basic j is
     * the coefficient of j in the expression for basic[i]:
     *
     *      basic[i] = sum_j a_ij * j
     *
     * (no minus sign in the relation -- the relation is direct.)
     *
     * To pivot, take n out of the sum and solve for n:
     *      basic = a_n * n + rest                (rest = sum_{j!=n} a_j j)
     *      n = (basic - rest) / a_n
     *      n = (1/a_n) * basic + sum_{j!=n} (-a_j/a_n) * j
     *
     * So after pivot, with n as new basic:
     *      a[ri][basic] = 1/a_n           (basic is now non-basic)
     *      a[ri][j]     = -a_j/a_n        (other non-basics, j!=n)
     *
     * That matches what I wrote. Good. */

    /* Now make n the basic variable of this row. */
    lra->is_basic[b] = 0;
    lra->row_of  [b] = b;   /* convention: non-basic's row_of is itself */
    lra->is_basic[n] = 1;
    lra->row_of  [n] = ri;
    lra->rows[ri]   = n;

    /* Step 2: every other row that mentions n must substitute. For
     * each other row k with c = a[k][n], replace:
     *
     *   basic[k] = c * n + (rest of row k)
     *
     * with
     *
     *   basic[k] = c * (a[ri][...]) + (rest of row k - c * n term)
     *
     * i.e. add c * (row ri's new coefficients) to row k, except for
     * column n which is zeroed.
     */
    for (int k = 0; k < lra->n_rows; k++) {
        if (k == ri) continue;
        Rat c = lra->a[k * lra->cap_vars + n];
        if (rat_is_zero(c)) continue;
        /* Zero the n column for row k. */
        lra->a[k * lra->cap_vars + n] = rat_int(0);
        /* Add c * row[ri] to row k, columns != n. */
        for (int j = 0; j < lra->n_vars; j++) {
            if (j == n) continue;
            Rat aij_ri = lra->a[ri * lra->cap_vars + j];
            if (rat_is_zero(aij_ri)) continue;
            lra->a[k * lra->cap_vars + j] =
                rat_add(lra->a[k * lra->cap_vars + j], rat_mul(c, aij_ri));
        }
    }

    lra->n_pivots++;
}

/* Update the value of a freshly-pivoted basic variable's row to the
 * given target (its bound). Adjust the new non-basic accordingly. */
static void pivot_and_update(LRA *lra, int basic_row, int new_basic_var,
                             Drat target)
{
    /* The OLD basic is the one currently violating its bound. After
     * pivoting, the old basic becomes non-basic and is freely settable;
     * we set it to its violated bound (target) and let the change
     * propagate to all basics. */
    int old_basic = lra->rows[basic_row];
    pivot(lra, basic_row, new_basic_var);
    /* old_basic is now non-basic; pin it to target. */
    update_value_nonbasic(lra, old_basic, target);
    (void)0;
}

/* ====== conflict construction ======
 * When a basic variable x_i violates its bound and the row has no
 * non-basic with the "freedom" to absorb the discrepancy, we have a
 * conflict. The reason:
 *
 *   * If x_i < lo[x_i]: we want to increase x_i. Row gives
 *       x_i = sum_j a_ij * x_j.
 *     For each j with a_ij > 0: increasing x_j increases x_i. If
 *     x_j is already at its upper bound, we cannot increase it.
 *     For each j with a_ij < 0: decreasing x_j increases x_i. If
 *     x_j is at its lower bound, we cannot decrease it.
 *     If every j is "stuck", the row's value can't be raised: every
 *     a_ij > 0 non-basic at hi, every a_ij < 0 non-basic at lo.
 *     Reason clause: negation of lo[x_i] atom AND, for each blocked
 *     non-basic, the bound atom that blocks it.
 *
 *   * If x_i > hi[x_i]: symmetric (signs flipped).
 *
 * The "reason" is the conjunction. Negated, it's the conflict clause.
 *
 * Each bound atom may have positive or negative SAT polarity:
 *   pos LE: v <= r            (atom is +, hi)
 *   neg LE: v > r             (atom is -, lo)
 *   pos LT: v < r             (atom is +, hi)
 *   neg LT: v >= r            (atom is -, lo)
 * The bound atom for hi/lo is stored as the lit (positive or negative).
 * Negation flips the sign for the conflict clause.
 */
static void grow_conf(LRA *lra, int needed) {
    if (needed <= lra->cap_conf) return;
    int nc = lra->cap_conf ? lra->cap_conf * 2 : 16;
    while (nc < needed) nc *= 2;
    lra->conf_buf = realloc(lra->conf_buf, (size_t)nc * sizeof(int));
    lra->cap_conf = nc;
}

/* push -lit (negation as DIMACS) into conf_buf */
static void conf_push_neg(LRA *lra, int lit) {
    grow_conf(lra, lra->conf_n + 1);
    lra->conf_buf[lra->conf_n++] = -lit;
}

/* Construct the conflict clause for "basic x_i can't go up to lo[x_i]". */
static void build_conflict_lo(LRA *lra, int ri) {
    int xi = lra->rows[ri];
    lra->conf_n = 0;
    /* Need to negate: the asserted lo[xi] caused this; the asserted
     * hi[j] (a_ij>0) or lo[j] (a_ij<0) for blocked non-basics. */
    conf_push_neg(lra, lra->lo_atom[xi]);
    for (int j = 0; j < lra->n_vars; j++) {
        if (lra->is_basic[j]) continue;
        if (j == xi)           continue;
        Rat c = lra->a[ri * lra->cap_vars + j];
        if (rat_is_zero(c)) continue;
        if (rat_sgn(c) > 0) {
            /* j stuck at upper bound. */
            conf_push_neg(lra, lra->hi_atom[j]);
        } else {
            conf_push_neg(lra, lra->lo_atom[j]);
        }
    }
}

static void build_conflict_hi(LRA *lra, int ri) {
    int xi = lra->rows[ri];
    lra->conf_n = 0;
    conf_push_neg(lra, lra->hi_atom[xi]);
    for (int j = 0; j < lra->n_vars; j++) {
        if (lra->is_basic[j]) continue;
        if (j == xi)           continue;
        Rat c = lra->a[ri * lra->cap_vars + j];
        if (rat_is_zero(c)) continue;
        if (rat_sgn(c) > 0) {
            conf_push_neg(lra, lra->lo_atom[j]);
        } else {
            conf_push_neg(lra, lra->hi_atom[j]);
        }
    }
}

/* ====== check entry ======
 * Process newly asserted lits up to lra->n_asserts, then run the
 * simplex pivot loop.
 *
 * On entry, basic variable values may already be off due to
 * non-basic clamping during apply_bound. Recompute all basics first
 * (cheap), then look for any basic var out of bound.
 */
static void recompute_all_basics(LRA *lra) {
    for (int i = 0; i < lra->n_rows; i++) {
        lra->val[lra->rows[i]] = recompute_basic(lra, i);
    }
}

/* ====== Phase 4: bound propagation ======
 *
 * After the simplex declares feasibility, scan every registered LRA
 * atom and check whether the current bounds-plus-tableau force it
 * one way or the other. If yes, queue a theory-propagation with a
 * Farkas-style reason.
 *
 * For atom (v OP r), we compute the implied upper bound (resp. lower
 * bound) on v from current state:
 *   - if v is non-basic: just hi[v] (resp. lo[v]).
 *   - if v is basic: traverse v's row; positive coefficients use
 *     hi[non-basic], negative use lo[non-basic]; the sum is the
 *     implied upper bound on v. Symmetric for lower.
 * Both branches return a list of atom lits that established the
 * bounds used; those are the reasons.
 *
 * Atom A=true is implied iff implied_upper(v) <= r (weak) or < r
 * (strict). A=false is implied iff implied_lower(v) > r or >= r.
 */

/* Returns 1 if v has a derivable upper bound under current state,
 * 0 otherwise. On success *out is the bound and reasons[0..n) are
 * the lits responsible. */
static int implied_upper(LRA *lra, int v, Drat *out, int *reasons, int *n_reasons) {
    *n_reasons = 0;
    Drat best;
    int  best_lit = 0, has_best = 0;

    if (lra->has_hi[v]) {
        best = lra->hi[v];
        best_lit = lra->hi_atom[v];
        has_best = 1;
    }

    if (lra->is_basic[v]) {
        int ri = lra->row_of[v];
        Drat ub = DRAT_ZERO;
        int  scratch[256]; int n_scratch = 0;
        int  bounded = 1;
        for (int j = 0; j < lra->n_vars; j++) {
            if (j == v)              continue;
            if (lra->is_basic[j])    continue;
            Rat c = lra->a[ri * lra->cap_vars + j];
            if (rat_is_zero(c)) continue;
            if (rat_sgn(c) > 0) {
                if (!lra->has_hi[j]) { bounded = 0; break; }
                ub = drat_add(ub, drat_mul_rat(lra->hi[j], c));
                scratch[n_scratch++] = lra->hi_atom[j];
            } else {
                if (!lra->has_lo[j]) { bounded = 0; break; }
                ub = drat_add(ub, drat_mul_rat(lra->lo[j], c));
                scratch[n_scratch++] = lra->lo_atom[j];
            }
        }
        if (bounded && (!has_best || drat_cmp(ub, best) < 0)) {
            *out = ub;
            for (int k = 0; k < n_scratch; k++) reasons[k] = scratch[k];
            *n_reasons = n_scratch;
            return 1;
        }
    }
    if (has_best) {
        *out = best;
        reasons[0] = best_lit;
        *n_reasons = 1;
        return 1;
    }
    return 0;
}

static int implied_lower(LRA *lra, int v, Drat *out, int *reasons, int *n_reasons) {
    *n_reasons = 0;
    Drat best;
    int  best_lit = 0, has_best = 0;

    if (lra->has_lo[v]) {
        best = lra->lo[v];
        best_lit = lra->lo_atom[v];
        has_best = 1;
    }

    if (lra->is_basic[v]) {
        int ri = lra->row_of[v];
        Drat lb = DRAT_ZERO;
        int  scratch[256]; int n_scratch = 0;
        int  bounded = 1;
        for (int j = 0; j < lra->n_vars; j++) {
            if (j == v)              continue;
            if (lra->is_basic[j])    continue;
            Rat c = lra->a[ri * lra->cap_vars + j];
            if (rat_is_zero(c)) continue;
            if (rat_sgn(c) > 0) {
                if (!lra->has_lo[j]) { bounded = 0; break; }
                lb = drat_add(lb, drat_mul_rat(lra->lo[j], c));
                scratch[n_scratch++] = lra->lo_atom[j];
            } else {
                if (!lra->has_hi[j]) { bounded = 0; break; }
                lb = drat_add(lb, drat_mul_rat(lra->hi[j], c));
                scratch[n_scratch++] = lra->hi_atom[j];
            }
        }
        if (bounded && (!has_best || drat_cmp(lb, best) > 0)) {
            *out = lb;
            for (int k = 0; k < n_scratch; k++) reasons[k] = scratch[k];
            *n_reasons = n_scratch;
            return 1;
        }
    }
    if (has_best) {
        *out = best;
        reasons[0] = best_lit;
        *n_reasons = 1;
        return 1;
    }
    return 0;
}

static void queue_propagation(LRA *lra, int lit, const int *reasons, int n_reasons) {
    if (lra->n_propq == lra->cap_propq) {
        int old_cap = lra->cap_propq;
        lra->cap_propq = lra->cap_propq ? lra->cap_propq * 2 : 8;
        lra->propq = realloc(lra->propq,
                             (size_t)lra->cap_propq * sizeof(*lra->propq));
        for (int i = old_cap; i < lra->cap_propq; i++) {
            lra->propq[i].lit = 0;
            lra->propq[i].reason = NULL;
            lra->propq[i].n_reason = 0;
            lra->propq[i].reason_cap = 0;
        }
    }
    int n = 1 + n_reasons;
    struct PropEntry *p = &lra->propq[lra->n_propq++];
    if (p->reason_cap < n) {
        p->reason_cap = n;
        p->reason = realloc(p->reason, (size_t)n * sizeof(int));
    }
    p->reason[0] = lit;
    for (int i = 0; i < n_reasons; i++) p->reason[1 + i] = -reasons[i];
    p->n_reason = n;
    p->lit = lit;
    lra->n_propagations++;
}

static void scan_propagations(LRA *lra) {
    int reasons[256];
    int n_reasons;
    for (int atom = 1; atom < lra->cap_atoms; atom++) {
        if (lra->atoms[atom].var < 0) continue;
        if (lra->atom_assigned[atom]) continue;

        int v  = lra->atoms[atom].var;
        int op = lra->atoms[atom].op;
        Rat r  = lra->atoms[atom].rhs;

        /* A=true: (v OP r) is the bound (v <= r) or (v < r). */
        Drat upper;
        if (implied_upper(lra, v, &upper, reasons, &n_reasons)) {
            Drat thresh = (op == LRA_OP_LE) ? drat_rat(r)
                                            : drat_make(r, rat_int(-1));
            if (drat_cmp(upper, thresh) <= 0) {
                queue_propagation(lra, atom, reasons, n_reasons);
                continue;
            }
        }
        /* A=false: ~(v OP r) is (v > r) or (v >= r). */
        Drat lower;
        if (implied_lower(lra, v, &lower, reasons, &n_reasons)) {
            Drat thresh = (op == LRA_OP_LE) ? drat_make(r, rat_int(1))
                                            : drat_rat(r);
            if (drat_cmp(lower, thresh) >= 0) {
                queue_propagation(lra, -atom, reasons, n_reasons);
            }
        }
    }
}

int lra_check(void *data, int **conflict_out, int *n_conflict_out) {
    LRA *lra = (LRA*)data;

    /* The new-assertion loop already applied bounds; just advance
     * check_pos. */
    lra->check_pos = lra->n_asserts;

    /* Direct bound contradiction: if any variable has lo > hi, the two
     * atoms that set those bounds form the Farkas certificate (a
     * 2-literal conflict clause). This catches single-variable
     * contradictions that don't go through the simplex tableau. */
    for (int v = 0; v < lra->n_vars; v++) {
        if (!lra->has_lo[v] || !lra->has_hi[v]) continue;
        if (drat_cmp(lra->lo[v], lra->hi[v]) > 0) {
            lra->conf_n = 0;
            conf_push_neg(lra, lra->lo_atom[v]);
            conf_push_neg(lra, lra->hi_atom[v]);
            *conflict_out   = lra->conf_buf;
            *n_conflict_out = lra->conf_n;
            return 1;
        }
    }

    recompute_all_basics(lra);

    /* Bland-style: smallest-index violated basic first; smallest-index
     * valid non-basic to pivot with. Avoids cycling. */
    int max_iters = 10000;
    while (max_iters-- > 0) {
        int viol_row = -1;
        int viol_low = 0; /* 1 if violating lower; 0 if upper */
        for (int i = 0; i < lra->n_rows; i++) {
            int b = lra->rows[i];
            if (lra->has_lo[b] && drat_cmp(lra->val[b], lra->lo[b]) < 0) {
                viol_row = i; viol_low = 1; break;
            }
            if (lra->has_hi[b] && drat_cmp(lra->val[b], lra->hi[b]) > 0) {
                viol_row = i; viol_low = 0; break;
            }
        }
        if (viol_row < 0) {
            /* Feasible. Look for theory propagations to enqueue. */
            scan_propagations(lra);
            return 0;
        }

        int b = lra->rows[viol_row];

        /* Find a non-basic j that can absorb the deficit.
         *   want to increase b (viol_low): need
         *     a_bj > 0 and val[j] < hi[j],  OR
         *     a_bj < 0 and val[j] > lo[j].
         *   want to decrease b (viol_high): need
         *     a_bj > 0 and val[j] > lo[j],  OR
         *     a_bj < 0 and val[j] < hi[j].
         */
        int found = -1;
        for (int j = 0; j < lra->n_vars; j++) {
            if (lra->is_basic[j]) continue;
            Rat c = lra->a[viol_row * lra->cap_vars + j];
            if (rat_is_zero(c)) continue;
            int csgn = rat_sgn(c);
            if (viol_low) {
                if (csgn > 0 && (!lra->has_hi[j] || drat_cmp(lra->val[j], lra->hi[j]) < 0))
                    { found = j; break; }
                if (csgn < 0 && (!lra->has_lo[j] || drat_cmp(lra->val[j], lra->lo[j]) > 0))
                    { found = j; break; }
            } else {
                if (csgn > 0 && (!lra->has_lo[j] || drat_cmp(lra->val[j], lra->lo[j]) > 0))
                    { found = j; break; }
                if (csgn < 0 && (!lra->has_hi[j] || drat_cmp(lra->val[j], lra->hi[j]) < 0))
                    { found = j; break; }
            }
        }

        if (found < 0) {
            /* Conflict. */
            if (viol_low) build_conflict_lo(lra, viol_row);
            else          build_conflict_hi(lra, viol_row);
            *conflict_out   = lra->conf_buf;
            *n_conflict_out = lra->conf_n;
            return 1;
        }

        /* Pivot b with `found`. Target: bring val[b] to its bound. */
        Drat target = viol_low ? lra->lo[b] : lra->hi[b];
        pivot_and_update(lra, viol_row, found, target);
    }
    /* Should never happen with Bland's rule; treat as no conflict. */
    return 0;
}

/* ====== propagation (skeletal) ====== */
int lra_propagate(void *data, int *prop_out,
                  int **reason_out, int *n_reason_out)
{
    LRA *lra = (LRA*)data;
    while (lra->propq_head < lra->n_propq) {
        struct PropEntry *p = &lra->propq[lra->propq_head];
        /* skip already-assigned lits */
        lra->propq_head++;
        if (p->lit == 0) continue;
        *prop_out      = p->lit;
        *reason_out    = p->reason;
        *n_reason_out  = p->n_reason;
        return 1;
    }
    return 0;
}

/* ====== queries ====== */
int lra_owns_atom(const LRA *lra, int atom) {
    if (atom <= 0 || atom >= lra->cap_atoms) return 0;
    return lra->atoms[atom].var >= 0;
}

long lra_n_pivots       (const LRA *lra) { return lra->n_pivots; }
long lra_n_propagations (const LRA *lra) { return lra->n_propagations; }
