/*
 * smt.c -- SMT integration: builds the SAT solver, hooks both
 *          theories (EUF + LRA) to it, manages atom IDs, accumulates
 *          user clauses, and drives smt_check.
 *
 * Phase 3 changed the integration from "one theory" to "a combined
 * theory dispatch." Each atom is owned by exactly one theory at
 * registration time:
 *
 *   - EUF atoms are created by smt_mk_eq_atom and bind two SmtTerms.
 *   - LRA atoms are created by smt_mk_lra_le_atom / smt_mk_lra_lt_atom
 *     and bind a linear combination of SmtVar to a rational right-hand
 *     side.
 *
 * The combined theory's assign / unassign / check / propagate hooks
 * dispatch each call to the owning theory. The SAT solver sees a
 * single TheoryHooks bundle.
 *
 * This shape generalises to Nelson--Oppen in Phase 4: a cross-theory
 * step inserted between check and propagate exchanges derived
 * equalities between EUF and LRA on shared sub-terms.
 */
#include "smt.h"
#include "sat.h"
#include "euf.h"
#include "lra.h"
#include "rat.h"

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

enum { OWNER_NONE = 0, OWNER_EUF = 1, OWNER_LRA = 2 };

struct SmtCtx {
    EUF *euf;
    LRA *lra;
    Sat *sat;     /* created on first smt_check */

    /* atom -> owner theory. Indexed 1..n_atoms. */
    int *atom_owner;
    int  cap_atom_owner;

    /* EUF atom registry: each atom that is an equality of two terms. */
    struct {
        int t1, t2;
    } *eq_atoms;
    int n_atoms;          /* total atoms allocated; ids are 1..n_atoms */
    int cap_eq_atoms;

    /* Phase 4: shared variables. shareds[i].euf is an EUF constant,
     * shareds[i].lra is an LRA variable. Pair-bridge atoms are
     * created at smt_mk_shared time against every previous shared. */
    SmtShared *shareds;
    int        n_shareds, cap_shareds;
    /* For each pair (i < j) of shareds, the EUF eq atom. Stored
     * as a flat lower-triangular array: pair_eq[j*(j-1)/2 + i]. */
    int *pair_eq;
    int  cap_pair_eq;

    /* Pending clauses (same as Phase 2). */
    int *clause_buf;
    int  clause_buf_n, clause_buf_cap;
    int *clause_starts;
    int  n_clauses, cap_clauses;

    int verbose;
    int solved;
};

/* ====== forward decls of combined-theory shims ====== */
static void combined_assign  (void *data, int lit);
static void combined_unassign(void *data, int lit);
static int  combined_check   (void *data, int **conflict_out, int *n_conflict_out);
static int  combined_propagate(void *data, int *prop_out,
                               int **reason_out, int *n_reason_out);

/* ====== atom owner table ====== */
static void set_atom_owner(SmtCtx *ctx, int atom, int owner) {
    if (atom <= 0) return;
    if (atom >= ctx->cap_atom_owner) {
        int nc = ctx->cap_atom_owner ? ctx->cap_atom_owner * 2 : 16;
        while (nc <= atom) nc *= 2;
        ctx->atom_owner = realloc(ctx->atom_owner, (size_t)nc * sizeof(int));
        for (int i = ctx->cap_atom_owner; i < nc; i++) ctx->atom_owner[i] = OWNER_NONE;
        ctx->cap_atom_owner = nc;
    }
    ctx->atom_owner[atom] = owner;
}

static int atom_owner(const SmtCtx *ctx, int atom) {
    if (atom <= 0 || atom >= ctx->cap_atom_owner) return OWNER_NONE;
    return ctx->atom_owner[atom];
}

/* Allocate a fresh atom id and grow all per-atom side-tables together.
 * Without this, an LRA atom would bump n_atoms but leave eq_atoms
 * un-grown, so the next EUF allocation would index past the end. */
static int new_atom(SmtCtx *ctx, int owner) {
    int atom = ++ctx->n_atoms;
    set_atom_owner(ctx, atom, owner);
    if (atom > ctx->cap_eq_atoms) {
        int nc = ctx->cap_eq_atoms ? ctx->cap_eq_atoms * 2 : 16;
        while (nc < atom) nc *= 2;
        ctx->eq_atoms = realloc(ctx->eq_atoms,
                                (size_t)nc * sizeof(*ctx->eq_atoms));
        for (int i = ctx->cap_eq_atoms; i < nc; i++) {
            ctx->eq_atoms[i].t1 = -1;
            ctx->eq_atoms[i].t2 = -1;
        }
        ctx->cap_eq_atoms = nc;
    }
    return atom;
}

/* ====== terms (EUF) ====== */
SmtTerm smt_mk_const(SmtCtx *ctx, const char *name) {
    return euf_mk_const(ctx->euf, name);
}
SmtTerm smt_mk_app(SmtCtx *ctx, const char *fname,
                   const SmtTerm *args, int n_args)
{
    return euf_mk_app(ctx->euf, fname, args, n_args);
}

/* ====== EUF atom allocation ====== */
SmtAtom smt_mk_eq_atom(SmtCtx *ctx, SmtTerm a, SmtTerm b) {
    if (a > b) { int t = a; a = b; b = t; }
    /* dedupe */
    for (int i = 0; i < ctx->n_atoms; i++) {
        if (atom_owner(ctx, i + 1) != OWNER_EUF) continue;
        int t1 = ctx->eq_atoms[i].t1, t2 = ctx->eq_atoms[i].t2;
        if (t1 == a && t2 == b) return i + 1;
    }
    int atom = new_atom(ctx, OWNER_EUF);
    ctx->eq_atoms[atom - 1].t1 = a;
    ctx->eq_atoms[atom - 1].t2 = b;
    euf_register_eq(ctx->euf, atom, a, b);
    return atom;
}

/* ====== LRA: variables and atoms ====== */
SmtVar smt_mk_real_var(SmtCtx *ctx, const char *name) {
    return lra_mk_var(ctx->lra, name);
}

static SmtAtom new_lra_atom(SmtCtx *ctx, int op,
                            const SmtLinTerm *terms, int n_terms,
                            int64_t rhs_num, int64_t rhs_den)
{
    int atom = new_atom(ctx, OWNER_LRA);

    LraTerm *lt = malloc((size_t)n_terms * sizeof(LraTerm));
    for (int i = 0; i < n_terms; i++) {
        lt[i].var   = terms[i].var;
        lt[i].coeff = rat_make(terms[i].num, terms[i].den);
    }
    lra_register_atom(ctx->lra, atom, op, lt, n_terms,
                      rat_make(rhs_num, rhs_den));
    free(lt);
    return atom;
}

SmtAtom smt_mk_lra_le_atom(SmtCtx *ctx,
                           const SmtLinTerm *terms, int n_terms,
                           int64_t rhs_num, int64_t rhs_den)
{
    return new_lra_atom(ctx, LRA_OP_LE, terms, n_terms, rhs_num, rhs_den);
}

SmtAtom smt_mk_lra_lt_atom(SmtCtx *ctx,
                           const SmtLinTerm *terms, int n_terms,
                           int64_t rhs_num, int64_t rhs_den)
{
    return new_lra_atom(ctx, LRA_OP_LT, terms, n_terms, rhs_num, rhs_den);
}

/* ====== clause accumulation ====== */
static void push_clause(SmtCtx *ctx, const int *lits, int n) {
    if (ctx->n_clauses == ctx->cap_clauses) {
        ctx->cap_clauses = ctx->cap_clauses ? ctx->cap_clauses * 2 : 8;
        ctx->clause_starts = realloc(ctx->clause_starts,
                                     (size_t)ctx->cap_clauses * sizeof(int));
    }
    int want = ctx->clause_buf_n + n + 1;
    if (want > ctx->clause_buf_cap) {
        int nc = ctx->clause_buf_cap ? ctx->clause_buf_cap * 2 : 32;
        while (nc < want) nc *= 2;
        ctx->clause_buf = realloc(ctx->clause_buf, (size_t)nc * sizeof(int));
        ctx->clause_buf_cap = nc;
    }
    ctx->clause_starts[ctx->n_clauses++] = ctx->clause_buf_n;
    ctx->clause_buf[ctx->clause_buf_n++] = n;
    for (int i = 0; i < n; i++) ctx->clause_buf[ctx->clause_buf_n++] = lits[i];
}

void smt_assert_clause(SmtCtx *ctx, const int *signed_atoms, int n) {
    push_clause(ctx, signed_atoms, n);
}

void smt_assert_eq(SmtCtx *ctx, SmtTerm a, SmtTerm b) {
    int atom = smt_mk_eq_atom(ctx, a, b);
    int u = atom;  push_clause(ctx, &u, 1);
}
void smt_assert_neq(SmtCtx *ctx, SmtTerm a, SmtTerm b) {
    int atom = smt_mk_eq_atom(ctx, a, b);
    int u = -atom; push_clause(ctx, &u, 1);
}

void smt_assert_lra_le(SmtCtx *ctx, const SmtLinTerm *terms, int n_terms,
                       int64_t rn, int64_t rd)
{
    int a = smt_mk_lra_le_atom(ctx, terms, n_terms, rn, rd);
    int u = a; push_clause(ctx, &u, 1);
}
void smt_assert_lra_lt(SmtCtx *ctx, const SmtLinTerm *terms, int n_terms,
                       int64_t rn, int64_t rd)
{
    int a = smt_mk_lra_lt_atom(ctx, terms, n_terms, rn, rd);
    int u = a; push_clause(ctx, &u, 1);
}
void smt_assert_lra_eq(SmtCtx *ctx, const SmtLinTerm *terms, int n_terms,
                       int64_t rn, int64_t rd)
{
    int a1 = smt_mk_lra_le_atom(ctx, terms, n_terms, rn, rd);
    int a2 = smt_mk_lra_lt_atom(ctx, terms, n_terms, rn, rd);
    int u1 =  a1; push_clause(ctx, &u1, 1);
    int u2 = -a2; push_clause(ctx, &u2, 1);
}

/* ====== Phase 4: shared variables and pair-bridge atoms ======
 *
 * Each shared s is an EUF constant and an LRA real-valued variable
 * with the same name. When a new shared is created, we eagerly pair
 * it with every previously-created shared and install:
 *   - one EUF equality atom eq_ij,
 *   - one LRA "<= 0" atom le_ij (for x_i - x_j),
 *   - one LRA "<= 0" atom ge_ij (for x_j - x_i, which encodes >= 0
 *     for x_i - x_j),
 * plus three linking clauses that wire them together at the SAT
 * level.
 *
 * Storage layout: pair_eq is the lower-triangle flattening of the
 * pair-atom matrix. For shareds (i, j) with i < j, the EUF atom is
 * at index j*(j-1)/2 + i. We don't actually need to remember the
 * LRA atoms ourselves -- once installed, they live in the LRA
 * solver and the link clauses already carry their ids. */

static int *pair_eq_slot(SmtCtx *ctx, int i, int j) {
    if (i > j) { int t = i; i = j; j = t; }
    int idx = j * (j - 1) / 2 + i;
    if (idx >= ctx->cap_pair_eq) {
        int nc = ctx->cap_pair_eq ? ctx->cap_pair_eq * 2 : 16;
        while (nc <= idx) nc *= 2;
        ctx->pair_eq = realloc(ctx->pair_eq, (size_t)nc * sizeof(int));
        for (int k = ctx->cap_pair_eq; k < nc; k++) ctx->pair_eq[k] = 0;
        ctx->cap_pair_eq = nc;
    }
    return &ctx->pair_eq[idx];
}

static void install_pair_bridge(SmtCtx *ctx, int i, int j) {
    SmtShared a = ctx->shareds[i];
    SmtShared b = ctx->shareds[j];

    /* eq_ij: EUF equality. */
    int eq = smt_mk_eq_atom(ctx, a.euf, b.euf);
    *pair_eq_slot(ctx, i, j) = eq;

    /* le_ij: a.lra - b.lra <= 0 */
    SmtLinTerm t_le[2] = { {a.lra, 1, 1}, {b.lra, -1, 1} };
    int le = smt_mk_lra_le_atom(ctx, t_le, 2, 0, 1);

    /* ge_ij: encode a.lra - b.lra >= 0 as b.lra - a.lra <= 0 */
    SmtLinTerm t_ge[2] = { {a.lra, -1, 1}, {b.lra, 1, 1} };
    int ge = smt_mk_lra_le_atom(ctx, t_ge, 2, 0, 1);

    /* Bridge clauses:
     *   eq -> le         :  -eq v le
     *   eq -> ge         :  -eq v ge
     *   le & ge -> eq    :  -le v -ge v eq
     */
    int c1[2] = { -eq, le };          push_clause(ctx, c1, 2);
    int c2[2] = { -eq, ge };          push_clause(ctx, c2, 2);
    int c3[3] = { -le, -ge, eq };     push_clause(ctx, c3, 3);
}

SmtShared smt_mk_shared(SmtCtx *ctx, const char *name) {
    SmtShared s;
    s.euf = euf_mk_const(ctx->euf, name);
    s.lra = lra_mk_var(ctx->lra, name);

    if (ctx->n_shareds == ctx->cap_shareds) {
        ctx->cap_shareds = ctx->cap_shareds ? ctx->cap_shareds * 2 : 8;
        ctx->shareds = realloc(ctx->shareds,
                               (size_t)ctx->cap_shareds * sizeof(SmtShared));
    }
    int my_index = ctx->n_shareds;
    ctx->shareds[ctx->n_shareds++] = s;

    /* Install bridge atoms with every previously-created shared. */
    for (int k = 0; k < my_index; k++) install_pair_bridge(ctx, k, my_index);

    return s;
}

static int find_shared_index(SmtCtx *ctx, SmtShared s) {
    for (int k = 0; k < ctx->n_shareds; k++) {
        if (ctx->shareds[k].euf == s.euf && ctx->shareds[k].lra == s.lra) return k;
    }
    return -1;
}

void smt_assert_shared_eq(SmtCtx *ctx, SmtShared s1, SmtShared s2) {
    int i = find_shared_index(ctx, s1);
    int j = find_shared_index(ctx, s2);
    if (i < 0 || j < 0 || i == j) return;
    int eq = *pair_eq_slot(ctx, i, j);
    int u = eq;  push_clause(ctx, &u, 1);
}

void smt_assert_shared_neq(SmtCtx *ctx, SmtShared s1, SmtShared s2) {
    int i = find_shared_index(ctx, s1);
    int j = find_shared_index(ctx, s2);
    if (i < 0 || j < 0 || i == j) return;
    int eq = *pair_eq_slot(ctx, i, j);
    int u = -eq; push_clause(ctx, &u, 1);
}

/* ====== combined-theory shims ====== */
static void combined_assign(void *data, int lit) {
    SmtCtx *ctx = (SmtCtx*)data;
    int atom = lit > 0 ? lit : -lit;
    int o = atom_owner(ctx, atom);
    if (o == OWNER_EUF) euf_assign(ctx->euf, lit);
    else if (o == OWNER_LRA) lra_assign(ctx->lra, lit);
}
static void combined_unassign(void *data, int lit) {
    SmtCtx *ctx = (SmtCtx*)data;
    int atom = lit > 0 ? lit : -lit;
    int o = atom_owner(ctx, atom);
    if (o == OWNER_EUF) euf_unassign(ctx->euf, lit);
    else if (o == OWNER_LRA) lra_unassign(ctx->lra, lit);
}
static int combined_check(void *data, int **conflict_out, int *n_conflict_out) {
    SmtCtx *ctx = (SmtCtx*)data;
    if (euf_check(ctx->euf, conflict_out, n_conflict_out)) return 1;
    if (lra_check(ctx->lra, conflict_out, n_conflict_out)) return 1;
    return 0;
}
static int combined_propagate(void *data, int *prop_out,
                              int **reason_out, int *n_reason_out)
{
    SmtCtx *ctx = (SmtCtx*)data;
    if (euf_propagate(ctx->euf, prop_out, reason_out, n_reason_out)) return 1;
    if (lra_propagate(ctx->lra, prop_out, reason_out, n_reason_out)) return 1;
    return 0;
}

/* ====== build SAT and solve ====== */
int smt_check(SmtCtx *ctx) {
    if (ctx->solved) return SMT_UNKNOWN;
    ctx->solved = 1;

    if (ctx->n_atoms == 0) {
        for (int c = 0; c < ctx->n_clauses; c++) {
            int off = ctx->clause_starts[c];
            int n = ctx->clause_buf[off];
            if (n == 0) return SMT_UNSAT;
        }
        return SMT_SAT;
    }

    ctx->sat = sat_new(ctx->n_atoms);

    static TheoryHooks hooks;
    hooks.data      = ctx;
    hooks.assign    = combined_assign;
    hooks.unassign  = combined_unassign;
    hooks.check     = combined_check;
    hooks.propagate = combined_propagate;
    sat_set_theory(ctx->sat, &hooks);

    for (int c = 0; c < ctx->n_clauses; c++) {
        int off = ctx->clause_starts[c];
        int n   = ctx->clause_buf[off];
        int *lits = &ctx->clause_buf[off + 1];
        if (sat_add_clause(ctx->sat, lits, n)) return SMT_UNSAT;
    }

    int r = sat_solve(ctx->sat);
    if (r == SAT_SATISFIABLE)   return SMT_SAT;
    if (r == SAT_UNSATISFIABLE) return SMT_UNSAT;
    return SMT_UNKNOWN;
}

/* ====== model queries ====== */
int smt_model_eq(SmtCtx *ctx, SmtTerm a, SmtTerm b) {
    return euf_find(ctx->euf, a) == euf_find(ctx->euf, b);
}
int smt_term_print(const SmtCtx *ctx, SmtTerm t, char *buf, int sz) {
    return euf_term_print(ctx->euf, t, buf, sz);
}

/* ====== lifecycle and stats ====== */
SmtCtx *smt_new(void) {
    SmtCtx *ctx = calloc(1, sizeof(SmtCtx));
    ctx->euf = euf_new();
    ctx->lra = lra_new();
    return ctx;
}

void smt_free(SmtCtx *ctx) {
    if (!ctx) return;
    if (ctx->sat) sat_free(ctx->sat);
    euf_free(ctx->euf);
    lra_free(ctx->lra);
    free(ctx->atom_owner);
    free(ctx->eq_atoms);
    free(ctx->shareds);
    free(ctx->pair_eq);
    free(ctx->clause_buf);
    free(ctx->clause_starts);
    free(ctx);
}

int smt_n_atoms(const SmtCtx *ctx) { return ctx->n_atoms; }
int smt_n_terms(const SmtCtx *ctx) { return euf_n_terms(ctx->euf); }

long smt_decisions       (const SmtCtx *c) { return c->sat ? sat_decisions(c->sat) : 0; }
long smt_conflicts       (const SmtCtx *c) { return c->sat ? sat_conflicts(c->sat) : 0; }
long smt_theory_conflicts(const SmtCtx *c) { return c->sat ? sat_theory_conflicts(c->sat) : 0; }
long smt_theory_props    (const SmtCtx *c) { return c->sat ? sat_theory_props(c->sat) : 0; }
long smt_propagations    (const SmtCtx *c) { return c->sat ? sat_propagations(c->sat) : 0; }
long smt_lra_pivots      (const SmtCtx *c) { return c->lra ? lra_n_pivots(c->lra) : 0; }

void smt_set_verbose(SmtCtx *ctx, int v) {
    ctx->verbose = v;
    if (ctx->sat) sat_set_verbose(ctx->sat, v);
}
