/*
 * sanity.c -- smoke tests for the SMT solver.
 */
#include "smt.h"
#include <stdio.h>
#include <stdlib.h>
#include <string.h>

static int pass = 0, total = 0;
#define CHECK(c, name) do { \
    total++; \
    if (c) { pass++; printf("  [ok ]  %s\n", name); } \
    else   { printf("  [FAIL] %s\n", name); } \
} while (0)

/* Test 1: a=b, b=c, a=c -- SAT. */
static void t_chain_sat(void) {
    SmtCtx *ctx = smt_new();
    int a = smt_mk_const(ctx, "a");
    int b = smt_mk_const(ctx, "b");
    int c = smt_mk_const(ctx, "c");
    smt_assert_eq(ctx, a, b);
    smt_assert_eq(ctx, b, c);
    smt_assert_eq(ctx, a, c);
    CHECK(smt_check(ctx) == SMT_SAT, "transitive equality, SAT");
    smt_free(ctx);
}

/* Test 2: a=b, b=c, a!=c -- UNSAT. */
static void t_chain_unsat(void) {
    SmtCtx *ctx = smt_new();
    int a = smt_mk_const(ctx, "a");
    int b = smt_mk_const(ctx, "b");
    int c = smt_mk_const(ctx, "c");
    smt_assert_eq(ctx, a, b);
    smt_assert_eq(ctx, b, c);
    smt_assert_neq(ctx, a, c);
    CHECK(smt_check(ctx) == SMT_UNSAT, "transitivity violated -> UNSAT");
    smt_free(ctx);
}

/* Test 3: a=b, f(a) != f(b) -- UNSAT by congruence. */
static void t_congruence(void) {
    SmtCtx *ctx = smt_new();
    int a = smt_mk_const(ctx, "a");
    int b = smt_mk_const(ctx, "b");
    int fa = smt_mk_app(ctx, "f", &a, 1);
    int fb = smt_mk_app(ctx, "f", &b, 1);
    smt_assert_eq(ctx, a, b);
    smt_assert_neq(ctx, fa, fb);
    CHECK(smt_check(ctx) == SMT_UNSAT, "f-congruence: a=b => f(a)=f(b)");
    smt_free(ctx);
}

/* Test 4: pure SAT branching. (a=b) OR (a!=b), nothing else. -- SAT. */
static void t_or(void) {
    SmtCtx *ctx = smt_new();
    int a = smt_mk_const(ctx, "a");
    int b = smt_mk_const(ctx, "b");
    int eq = smt_mk_eq_atom(ctx, a, b);
    int clause[2] = { eq, -eq };  /* always SAT (tautology) */
    smt_assert_clause(ctx, clause, 2);
    CHECK(smt_check(ctx) == SMT_SAT, "trivial disjunction, SAT");
    smt_free(ctx);
}

/* Test 5: branching across theory atoms.
 *   (a=b) OR (c=d), and a!=b, and c!=d  -- UNSAT.
 * Each branch should be killed by the theory. */
static void t_branch_unsat(void) {
    SmtCtx *ctx = smt_new();
    int a = smt_mk_const(ctx, "a");
    int b = smt_mk_const(ctx, "b");
    int c = smt_mk_const(ctx, "c");
    int d = smt_mk_const(ctx, "d");
    int eab = smt_mk_eq_atom(ctx, a, b);
    int ecd = smt_mk_eq_atom(ctx, c, d);
    int cl[2] = { eab, ecd };
    smt_assert_clause(ctx, cl, 2);
    smt_assert_neq(ctx, a, b);
    smt_assert_neq(ctx, c, d);
    CHECK(smt_check(ctx) == SMT_UNSAT, "both branches refuted by theory");
    smt_free(ctx);
}

/* Test 6: deeper congruence.
 *   a=b, g(f(a)) != g(f(b))  -- UNSAT (two-level congruence). */
static void t_deep(void) {
    SmtCtx *ctx = smt_new();
    int a = smt_mk_const(ctx, "a");
    int b = smt_mk_const(ctx, "b");
    int fa = smt_mk_app(ctx, "f", &a, 1);
    int fb = smt_mk_app(ctx, "f", &b, 1);
    int gfa = smt_mk_app(ctx, "g", &fa, 1);
    int gfb = smt_mk_app(ctx, "g", &fb, 1);
    smt_assert_eq(ctx, a, b);
    smt_assert_neq(ctx, gfa, gfb);
    CHECK(smt_check(ctx) == SMT_UNSAT, "deep congruence: a=b => g(f(a))=g(f(b))");
    smt_free(ctx);
}

/* Test 7 (Phase 2): hidden-root congruence.
 * Two function apps whose arguments end up in the same class via a
 * third "merge target" term that is itself not used as any argument.
 * Without use-list aggregation, scan_congruence_pairs misses the
 * congruence because uses[shared root] is empty.
 *   e=a, e=c, f(a) != f(c)  -- UNSAT. */
static void t_hidden_root(void) {
    SmtCtx *ctx = smt_new();
    int e = smt_mk_const(ctx, "e");  /* lowest term id -- wins union-by-size */
    int a = smt_mk_const(ctx, "a");
    int c = smt_mk_const(ctx, "c");
    int fa = smt_mk_app(ctx, "f", &a, 1);
    int fc = smt_mk_app(ctx, "f", &c, 1);
    smt_assert_eq (ctx, e, a);
    smt_assert_eq (ctx, e, c);
    smt_assert_neq(ctx, fa, fc);
    CHECK(smt_check(ctx) == SMT_UNSAT, "hidden-root congruence (use-list aggregation)");
    smt_free(ctx);
}

/* Test 8 (Phase 2): theory propagation observable.
 *
 * Setup: assert a=b, b=c, and *register* the atom (a=c) without
 * putting it in any clause. After the two merges the theory sees
 * find(a)=find(c) and propagates +(a=c) onto the SAT trail with a
 * proof-forest reason -- two equality lits negated, plus the atom.
 *
 * Without theory propagation the atom would stay unassigned and the
 * SAT solver would have to branch on it. With propagation: SAT, 0
 * decisions, at least one theory propagation. */
static void t_propagation_visible(void) {
    SmtCtx *ctx = smt_new();
    int a = smt_mk_const(ctx, "a");
    int b = smt_mk_const(ctx, "b");
    int c = smt_mk_const(ctx, "c");
    (void)smt_mk_eq_atom(ctx, a, c);   /* register (a=c) but do not assert */
    smt_assert_eq(ctx, a, b);
    smt_assert_eq(ctx, b, c);
    int r = smt_check(ctx);
    CHECK(r == SMT_SAT
          && smt_decisions(ctx) == 0
          && smt_theory_props(ctx) >= 1,
          "theory propagates (a=c) with no SAT decision");
    smt_free(ctx);
}

/* === Phase 3: LRA sanity tests === */

/* T9: single-variable bound consistency. x <= 5 AND x <= 3 -> SAT. */
static void t_lra_single_sat(void) {
    SmtCtx *ctx = smt_new();
    int x = smt_mk_real_var(ctx, "x");
    SmtLinTerm xt = {x, 1, 1};
    smt_assert_lra_le(ctx, &xt, 1, 5, 1);
    smt_assert_lra_le(ctx, &xt, 1, 3, 1);
    CHECK(smt_check(ctx) == SMT_SAT, "LRA: x<=5 AND x<=3 satisfiable");
    smt_free(ctx);
}

/* T10: single-variable contradiction across polarities. */
static void t_lra_single_unsat(void) {
    SmtCtx *ctx = smt_new();
    int x = smt_mk_real_var(ctx, "x");
    SmtLinTerm xt = {x, 1, 1};
    int a3 = smt_mk_lra_le_atom(ctx, &xt, 1, 3, 1);
    int a5 = smt_mk_lra_le_atom(ctx, &xt, 1, 5, 1);
    int u1 =  a3; smt_assert_clause(ctx, &u1, 1);   /* x <= 3 */
    int u2 = -a5; smt_assert_clause(ctx, &u2, 1);   /* x > 5  */
    CHECK(smt_check(ctx) == SMT_UNSAT, "LRA: x<=3 AND x>5 -> UNSAT");
    smt_free(ctx);
}

/* T11: multi-variable Farkas. x+y<=5, x>3, y>3 -- the classic. */
static void t_lra_farkas(void) {
    SmtCtx *ctx = smt_new();
    int x = smt_mk_real_var(ctx, "x");
    int y = smt_mk_real_var(ctx, "y");
    SmtLinTerm xy[2] = {{x, 1, 1}, {y, 1, 1}};
    smt_assert_lra_le(ctx, xy, 2, 5, 1);
    SmtLinTerm xt = {x, 1, 1}, yt = {y, 1, 1};
    int ax = smt_mk_lra_le_atom(ctx, &xt, 1, 3, 1);
    int ay = smt_mk_lra_le_atom(ctx, &yt, 1, 3, 1);
    int u1 = -ax; smt_assert_clause(ctx, &u1, 1);
    int u2 = -ay; smt_assert_clause(ctx, &u2, 1);
    int r = smt_check(ctx);
    CHECK(r == SMT_UNSAT && smt_theory_conflicts(ctx) >= 1,
          "LRA: x+y<=5, x>3, y>3 -> Farkas UNSAT");
    smt_free(ctx);
}

/* T12: branching over LRA disjunction.
 *   ((x <= 1) v (x >= 10))  AND  2 <= x <= 9.
 * Forces the SAT layer to enumerate both branches; each refuted by LRA. */
static void t_lra_branching(void) {
    SmtCtx *ctx = smt_new();
    int x = smt_mk_real_var(ctx, "x");
    SmtLinTerm xt = {x, 1, 1};
    int a_le1  = smt_mk_lra_le_atom(ctx, &xt, 1, 1,  1);
    int a_lt10 = smt_mk_lra_lt_atom(ctx, &xt, 1, 10, 1);
    int cl[2] = { a_le1, -a_lt10 };       /* (x<=1) v (x>=10) */
    smt_assert_clause(ctx, cl, 2);
    smt_assert_lra_le(ctx, &xt, 1, 9, 1); /* x <= 9 */
    int a_lt2 = smt_mk_lra_lt_atom(ctx, &xt, 1, 2, 1);
    int u = -a_lt2; smt_assert_clause(ctx, &u, 1);   /* x >= 2 */
    CHECK(smt_check(ctx) == SMT_UNSAT,
          "LRA: disjunction refuted on both branches");
    smt_free(ctx);
}

/* T13: mixed EUF + LRA, no theory interaction (no shared sub-terms).
 * EUF side is SAT; LRA side is UNSAT; combined is UNSAT. */
static void t_lra_mixed(void) {
    SmtCtx *ctx = smt_new();
    int x = smt_mk_real_var(ctx, "x");
    int y = smt_mk_real_var(ctx, "y");
    int a = smt_mk_const(ctx, "a");
    int b = smt_mk_const(ctx, "b");
    SmtLinTerm xy[2] = {{x, 1, 1}, {y, 1, 1}};
    smt_assert_lra_le(ctx, xy, 2, 5, 1);
    SmtLinTerm xt = {x, 1, 1}, yt = {y, 1, 1};
    int ax = smt_mk_lra_le_atom(ctx, &xt, 1, 3, 1);
    int ay = smt_mk_lra_le_atom(ctx, &yt, 1, 3, 1);
    int u1 = -ax; smt_assert_clause(ctx, &u1, 1);
    int u2 = -ay; smt_assert_clause(ctx, &u2, 1);
    smt_assert_eq(ctx, a, b);   /* irrelevant, EUF-side */
    CHECK(smt_check(ctx) == SMT_UNSAT, "LRA UNSAT + EUF SAT -> UNSAT");
    smt_free(ctx);
}

/* T14: strict inequalities and delta. x < 5 AND x >= 5 -- UNSAT. */
static void t_lra_strict(void) {
    SmtCtx *ctx = smt_new();
    int x = smt_mk_real_var(ctx, "x");
    SmtLinTerm xt = {x, 1, 1};
    smt_assert_lra_lt(ctx, &xt, 1, 5, 1);    /* x < 5 */
    int a5 = smt_mk_lra_lt_atom(ctx, &xt, 1, 5, 1);
    int u = -a5; smt_assert_clause(ctx, &u, 1); /* x >= 5 */
    CHECK(smt_check(ctx) == SMT_UNSAT, "LRA: x<5 AND x>=5 -> UNSAT (strict)");
    smt_free(ctx);
}

/* === Phase 4: Nelson-Oppen combination === */

/* T15: simplest cross-theory: shared_eq + congruence -> UNSAT.
 * Exercises only the EUF->LRA bridge clauses (the merge propagates
 * eq_xy to LRA, but LRA needs nothing more once x.euf and y.euf are
 * merged; congruence on f closes it). */
static void t_no_shared_eq_simple(void) {
    SmtCtx *ctx = smt_new();
    SmtShared x = smt_mk_shared(ctx, "x");
    SmtShared y = smt_mk_shared(ctx, "y");
    int fx = smt_mk_app(ctx, "f", &x.euf, 1);
    int fy = smt_mk_app(ctx, "f", &y.euf, 1);
    smt_assert_neq(ctx, fx, fy);
    smt_assert_shared_eq(ctx, x, y);
    CHECK(smt_check(ctx) == SMT_UNSAT,
          "NO: f(x)!=f(y) AND shared_eq(x,y) -> UNSAT");
    smt_free(ctx);
}

/* T16: shared_eq + LRA bound conflict.
 *   shared_eq(x, y), x <= 4, y >= 5.
 * EUF merges x.euf, y.euf; the bridge clauses force le_xy and ge_xy
 * to be asserted as bounds (x - y <= 0 and y - x <= 0). Then x <= 4
 * and y >= 5 force the bridge slack x - y to be <= 4 - 5 = -1, which
 * contradicts y - x <= 0 (i.e. x - y >= 0). UNSAT. */
static void t_no_eq_with_lra_conflict(void) {
    SmtCtx *ctx = smt_new();
    SmtShared x = smt_mk_shared(ctx, "x");
    SmtShared y = smt_mk_shared(ctx, "y");
    smt_assert_shared_eq(ctx, x, y);
    SmtLinTerm xt = {x.lra, 1, 1};
    SmtLinTerm yt = {y.lra, 1, 1};
    smt_assert_lra_le(ctx, &xt, 1, 4, 1);     /* x <= 4 */
    int ay5 = smt_mk_lra_lt_atom(ctx, &yt, 1, 5, 1);
    int u = -ay5; smt_assert_clause(ctx, &u, 1);   /* y >= 5 */
    CHECK(smt_check(ctx) == SMT_UNSAT,
          "NO: shared_eq(x,y) AND x<=4 AND y>=5 -> UNSAT");
    smt_free(ctx);
}

/* T17: LRA->EUF arrow. The key test. x=5, y=5, f(x)!=f(y) -> UNSAT.
 * LRA pins x and y to 5, theory-propagates the two bridge atoms
 * (le_xy, ge_xy); the bridge clause forces eq_xy; EUF merges the
 * classes; congruence on f gives the conflict. */
static void t_no_lra_to_euf(void) {
    SmtCtx *ctx = smt_new();
    SmtShared x = smt_mk_shared(ctx, "x");
    SmtShared y = smt_mk_shared(ctx, "y");
    int fx = smt_mk_app(ctx, "f", &x.euf, 1);
    int fy = smt_mk_app(ctx, "f", &y.euf, 1);
    smt_assert_neq(ctx, fx, fy);
    SmtLinTerm xt = {x.lra, 1, 1};
    SmtLinTerm yt = {y.lra, 1, 1};
    smt_assert_lra_eq(ctx, &xt, 1, 5, 1);
    smt_assert_lra_eq(ctx, &yt, 1, 5, 1);
    int r = smt_check(ctx);
    CHECK(r == SMT_UNSAT && smt_theory_props(ctx) >= 1,
          "NO: x=5 AND y=5 AND f(x)!=f(y) -> UNSAT via LRA->EUF");
    smt_free(ctx);
}

/* T18: three shareds transitivity. x=y, y=z, x!=z -> UNSAT.
 * Pure EUF reasoning over shareds; tests the pair-bridge installation
 * for three shareds (three pair atoms). */
static void t_no_transitive(void) {
    SmtCtx *ctx = smt_new();
    SmtShared x = smt_mk_shared(ctx, "x");
    SmtShared y = smt_mk_shared(ctx, "y");
    SmtShared z = smt_mk_shared(ctx, "z");
    smt_assert_shared_eq (ctx, x, y);
    smt_assert_shared_eq (ctx, y, z);
    smt_assert_shared_neq(ctx, x, z);
    CHECK(smt_check(ctx) == SMT_UNSAT,
          "NO: x=y, y=z, x!=z over shareds -> UNSAT");
    smt_free(ctx);
}

/* T19: NO does NOT over-constrain. Two unrelated shareds with f(x)
 * different from f(y) and no equality assertion -> SAT.
 * The pair-bridge atoms exist but are not forced; SAT freely chooses
 * eq=false and the bridge clauses don't constrain le, ge. */
static void t_no_independent(void) {
    SmtCtx *ctx = smt_new();
    SmtShared x = smt_mk_shared(ctx, "x");
    SmtShared y = smt_mk_shared(ctx, "y");
    int fx = smt_mk_app(ctx, "f", &x.euf, 1);
    int fy = smt_mk_app(ctx, "f", &y.euf, 1);
    smt_assert_neq(ctx, fx, fy);
    CHECK(smt_check(ctx) == SMT_SAT,
          "NO: shareds with no eq assertion remain independent");
    smt_free(ctx);
}

int main(void) {
    t_chain_sat();
    t_chain_unsat();
    t_congruence();
    t_or();
    t_branch_unsat();
    t_deep();
    t_hidden_root();
    t_propagation_visible();
    t_lra_single_sat();
    t_lra_single_unsat();
    t_lra_farkas();
    t_lra_branching();
    t_lra_mixed();
    t_lra_strict();
    t_no_shared_eq_simple();
    t_no_eq_with_lra_conflict();
    t_no_lra_to_euf();
    t_no_transitive();
    t_no_independent();
    printf("\n%d / %d passed\n", pass, total);
    return pass == total ? 0 : 1;
}
