/*
 * sat.c -- CDCL SAT solver with theory hooks.
 *
 * Inherits the core CDCL machinery from the standalone SAT solver in
 * Predictive Algorithms (two watched literals, first-UIP conflict
 * analysis, VSIDS on a binary heap, phase saving, Luby restarts). The
 * differences are:
 *
 *   - On every assignment / unassignment of a literal, the registered
 *     theory's `assign` / `unassign` callback is invoked, so the theory
 *     can mirror the SAT trail.
 *
 *   - After every round of unit propagation that does not itself produce
 *     a conflict, the theory's `check` callback is invoked. If the
 *     theory reports a conflict, the returned clause enters the solver's
 *     clause database as if conflict analysis had just produced it, and
 *     the analysis loop runs on it.
 *
 * Phase 1 does not implement theory propagation (the theory deriving
 * new literal implications during search). It will be added later.
 */

#include "sat.h"

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

/* ====== literal encoding ====== */
typedef int Lit;
typedef int Var;
#define LIT_UNDEF (-1)
#define MK_LIT(v, sign) (((v) << 1) | (sign))
#define NEG(l)          ((l) ^ 1)
#define VAR(l)          ((l) >> 1)
#define SIGN(l)         ((l) & 1)

/* External (DIMACS-style, 1-indexed signed int) <-> internal lit. */
static Lit ext2int(int e) {
    int v = (e > 0 ? e : -e) - 1;
    return MK_LIT(v, e < 0 ? 1 : 0);
}
static int int2ext(Lit l) {
    return SIGN(l) ? -(VAR(l) + 1) : (VAR(l) + 1);
}

/* ====== clauses ====== */
typedef struct Clause {
    int  size;
    int  learnt;
    Lit  lits[1];
} Clause;

static Clause *clause_new(int n, const Lit *lits, int learnt) {
    Clause *c = malloc(sizeof(Clause) + (size_t)(n - 1) * sizeof(Lit));
    c->size   = n;
    c->learnt = learnt;
    memcpy(c->lits, lits, (size_t)n * sizeof(Lit));
    return c;
}

typedef struct {
    Clause **buf;
    int      n;
    int      cap;
} CList;

static void cl_push(CList *l, Clause *c) {
    if (l->n == l->cap) {
        int nc = l->cap ? l->cap * 2 : 8;
        l->buf = realloc(l->buf, (size_t)nc * sizeof(Clause *));
        l->cap = nc;
    }
    l->buf[l->n++] = c;
}

/* ====== solver state ====== */
struct Sat {
    int nvars;
    int unsat;
    int verbose;

    int     *value;
    int     *level;
    Clause **reason;
    int     *phase;
    double  *activity;
    char    *seen;

    Lit *trail;
    int  trail_n, trail_cap;
    int  qhead;

    int *trail_lim;
    int  level_n, level_cap;

    CList *watches;
    CList  clauses;
    CList  learnts;

    Var    *heap;
    int    *heap_pos;
    int     heap_n;
    double  var_inc;
    double  var_decay;

    int   luby_i;
    long  restart_unit;

    long decisions, propagations, conflicts, restarts, learned_count;
    long theory_conflicts;
    long theory_props;

    long conflict_limit;

    const TheoryHooks *theory;
};

/* ====== VSIDS heap ====== */
static int heap_more(const Sat *s, Var a, Var b) {
    return s->activity[a] > s->activity[b];
}
static void heap_up(Sat *s, int i) {
    Var x = s->heap[i];
    while (i > 0) {
        int p = (i - 1) >> 1;
        if (!heap_more(s, x, s->heap[p])) break;
        s->heap[i] = s->heap[p];
        s->heap_pos[s->heap[i]] = i;
        i = p;
    }
    s->heap[i] = x;
    s->heap_pos[x] = i;
}
static void heap_down(Sat *s, int i) {
    Var x = s->heap[i];
    int c;
    while ((c = 2 * i + 1) < s->heap_n) {
        if (c + 1 < s->heap_n && heap_more(s, s->heap[c + 1], s->heap[c])) c++;
        if (!heap_more(s, s->heap[c], x)) break;
        s->heap[i] = s->heap[c];
        s->heap_pos[s->heap[i]] = i;
        i = c;
    }
    s->heap[i] = x;
    s->heap_pos[x] = i;
}
static void heap_push(Sat *s, Var v) {
    if (s->heap_pos[v] != -1) return;
    s->heap_pos[v] = s->heap_n;
    s->heap[s->heap_n++] = v;
    heap_up(s, s->heap_n - 1);
}
static Var heap_pop(Sat *s) {
    Var top = s->heap[0];
    s->heap_pos[top] = -1;
    s->heap_n--;
    if (s->heap_n > 0) {
        s->heap[0] = s->heap[s->heap_n];
        s->heap_pos[s->heap[0]] = 0;
        heap_down(s, 0);
    }
    return top;
}
static void heap_increase(Sat *s, Var v) {
    if (s->heap_pos[v] != -1) heap_up(s, s->heap_pos[v]);
}

static void bump_var(Sat *s, Var v) {
    s->activity[v] += s->var_inc;
    if (s->activity[v] > 1e100) {
        for (int i = 0; i < s->nvars; i++) s->activity[i] *= 1e-100;
        s->var_inc *= 1e-100;
    }
    heap_increase(s, v);
}
static void decay_var_act(Sat *s) {
    s->var_inc *= (1.0 / s->var_decay);
}

/* ====== assignment ====== */
static int lit_value(const Sat *s, Lit l) {
    int v = s->value[VAR(l)];
    if (v < 0) return -1;
    return v ^ SIGN(l);
}

static void enqueue(Sat *s, Lit l, Clause *reason) {
    Var v = VAR(l);
    int val = !SIGN(l);
    s->value[v]  = val;
    s->level[v]  = s->level_n;
    s->reason[v] = reason;
    s->phase[v]  = val;
    s->trail[s->trail_n++] = l;
    if (s->theory) s->theory->assign(s->theory->data, int2ext(l));
}

static void new_decision_level(Sat *s) {
    if (s->level_n >= s->level_cap) {
        int nc = s->level_cap ? s->level_cap * 2 : 16;
        s->trail_lim = realloc(s->trail_lim, (size_t)nc * sizeof(int));
        s->level_cap = nc;
    }
    s->trail_lim[s->level_n++] = s->trail_n;
}

static void backtrack_to(Sat *s, int level) {
    if (s->level_n <= level) return;
    int trail_start = s->trail_lim[level];
    for (int i = s->trail_n - 1; i >= trail_start; i--) {
        Lit  l = s->trail[i];
        Var  v = VAR(l);
        if (s->theory) s->theory->unassign(s->theory->data, int2ext(l));
        s->value[v]  = -1;
        s->reason[v] = NULL;
        heap_push(s, v);
    }
    s->trail_n = trail_start;
    s->qhead   = trail_start;
    s->level_n = level;
}

/* ====== watched-literal propagation ====== */
static Clause *propagate(Sat *s) {
    while (s->qhead < s->trail_n) {
        Lit p = s->trail[s->qhead++];
        Lit not_p = NEG(p);
        CList *ws = &s->watches[not_p];
        Clause **buf = ws->buf;
        int n = ws->n;
        int j = 0;

        for (int i = 0; i < n; i++) {
            Clause *c = buf[i];

            if (c->lits[0] == not_p) {
                Lit t = c->lits[0]; c->lits[0] = c->lits[1]; c->lits[1] = t;
            }
            Lit other = c->lits[0];

            if (lit_value(s, other) == 1) {
                buf[j++] = c;
                continue;
            }

            int found = 0;
            for (int k = 2; k < c->size; k++) {
                if (lit_value(s, c->lits[k]) != 0) {
                    Lit t = c->lits[1];
                    c->lits[1] = c->lits[k];
                    c->lits[k] = t;
                    cl_push(&s->watches[c->lits[1]], c);
                    found = 1;
                    break;
                }
            }
            if (found) continue;

            buf[j++] = c;

            if (lit_value(s, other) == 0) {
                while (++i < n) buf[j++] = buf[i];
                ws->n = j;
                s->qhead = s->trail_n;
                return c;
            }
            enqueue(s, other, c);
            s->propagations++;
        }
        ws->n = j;
    }
    return NULL;
}

/* ====== first-UIP conflict analysis ====== */
static int analyze(Sat *s, Clause *confl, Lit *out, int *btlevel) {
    int counter = 0;
    int trail_idx = s->trail_n - 1;
    int learnt_n = 1;
    int cur_level = s->level_n;
    Lit p = LIT_UNDEF;

    do {
        for (int j = (p == LIT_UNDEF ? 0 : 1); j < confl->size; j++) {
            Lit q = confl->lits[j];
            Var v = VAR(q);
            if (!s->seen[v] && s->level[v] > 0) {
                bump_var(s, v);
                s->seen[v] = 1;
                if (s->level[v] >= cur_level) counter++;
                else                          out[learnt_n++] = q;
            }
        }
        while (!s->seen[VAR(s->trail[trail_idx])]) trail_idx--;
        p = s->trail[trail_idx];
        confl = s->reason[VAR(p)];
        s->seen[VAR(p)] = 0;
        counter--;
        trail_idx--;
    } while (counter > 0);

    out[0] = NEG(p);

    *btlevel = 0;
    int max_i = 1;
    for (int i = 1; i < learnt_n; i++) {
        int lvl = s->level[VAR(out[i])];
        if (lvl > *btlevel) { *btlevel = lvl; max_i = i; }
    }
    if (learnt_n > 1 && max_i != 1) {
        Lit t = out[1]; out[1] = out[max_i]; out[max_i] = t;
    }

    for (int i = 1; i < learnt_n; i++) s->seen[VAR(out[i])] = 0;
    return learnt_n;
}

static Lit pick_branch_lit(Sat *s) {
    while (s->heap_n > 0) {
        Var v = heap_pop(s);
        if (s->value[v] == -1) {
            return MK_LIT(v, s->phase[v] ? 0 : 1);
        }
    }
    return LIT_UNDEF;
}

/* ====== Luby restart ====== */
static long luby(int i) {
    int size, seq;
    for (size = 1, seq = 0; size < i + 1; size = 2 * size + 1, seq++);
    while (size - 1 != i) {
        size = (size - 1) >> 1;
        seq--;
        i = i % size;
    }
    return 1L << seq;
}

/* ====== ingest a theory conflict ======
 * Returns a Clause* that conflict analysis can run on. The clause is added
 * to the learnts database (so it persists and participates in propagation
 * after backjump). The theory hands us a list of external lits whose
 * conjunction is theory-inconsistent; we negate to form a clause.
 */
static Clause *make_theory_clause(Sat *s, int *ext_lits, int n) {
    Lit *tmp = malloc((size_t)n * sizeof(Lit));
    for (int i = 0; i < n; i++) tmp[i] = ext2int(ext_lits[i]);

    /* The clause must be ordered so its two highest-level literals
     * are at positions 0 and 1, just like learned clauses. */
    int hi1 = 0, hi2 = 0;
    for (int i = 1; i < n; i++) {
        if (s->level[VAR(tmp[i])] > s->level[VAR(tmp[hi1])]) { hi2 = hi1; hi1 = i; }
        else if (s->level[VAR(tmp[i])] > s->level[VAR(tmp[hi2])]) hi2 = i;
    }
    if (hi1 != 0) { Lit t = tmp[0]; tmp[0] = tmp[hi1]; tmp[hi1] = t;
                    if (hi2 == 0) hi2 = hi1; }
    if (hi2 != 1) { Lit t = tmp[1]; tmp[1] = tmp[hi2]; tmp[hi2] = t; }

    Clause *c = clause_new(n, tmp, 1);
    free(tmp);
    cl_push(&s->learnts, c);
    if (n >= 2) {
        cl_push(&s->watches[c->lits[0]], c);
        cl_push(&s->watches[c->lits[1]], c);
    }
    s->learned_count++;
    return c;
}

/* ====== ingest a theory-propagated literal ======
 * The theory hands us a reason clause whose position 0 is the propagated
 * literal (currently undefined) and whose other positions are negations
 * of currently-asserted literals (i.e., all false on the trail).
 * The clause becomes a learned clause and the propagated lit's reason.
 * We watch lits[0] (the propagated lit) and the highest-level falsified
 * lit (which becomes lits[1]).
 */
static Clause *make_theory_reason(Sat *s, int *ext_lits, int n) {
    Lit *tmp = malloc((size_t)n * sizeof(Lit));
    for (int i = 0; i < n; i++) tmp[i] = ext2int(ext_lits[i]);

    /* tmp[0] is the propagated lit -- keep it at position 0.
     * Among tmp[1..n-1], move the highest-level falsified lit to
     * position 1, so the watched-literal invariant holds. */
    if (n >= 2) {
        int hi = 1;
        for (int i = 2; i < n; i++) {
            if (s->level[VAR(tmp[i])] > s->level[VAR(tmp[hi])]) hi = i;
        }
        if (hi != 1) { Lit t = tmp[1]; tmp[1] = tmp[hi]; tmp[hi] = t; }
    }

    Clause *c = clause_new(n, tmp, 1);
    free(tmp);
    cl_push(&s->learnts, c);
    if (n >= 2) {
        cl_push(&s->watches[c->lits[0]], c);
        cl_push(&s->watches[c->lits[1]], c);
    }
    s->learned_count++;
    return c;
}

/* ====== clause addition ====== */
static int lit_cmp(const void *a, const void *b) { return *(const Lit*)a - *(const Lit*)b; }

int sat_add_clause(Sat *s, const int *ext_lits, int n_lits) {
    if (s->unsat) return 1;
    if (n_lits == 0) { s->unsat = 1; return 1; }

    Lit *lits = malloc((size_t)n_lits * sizeof(Lit));
    int n = 0;
    for (int i = 0; i < n_lits; i++) {
        Lit l = ext2int(ext_lits[i]);
        if (VAR(l) < 0 || VAR(l) >= s->nvars) { free(lits); return 1; }
        lits[n++] = l;
    }

    qsort(lits, (size_t)n, sizeof(Lit), lit_cmp);
    int m = 0;
    for (int i = 0; i < n; i++) {
        if (i > 0 && lits[i] == lits[i - 1]) continue;
        if (m > 0 && lits[m - 1] == NEG(lits[i])) { free(lits); return 0; }
        lits[m++] = lits[i];
    }
    n = m;

    int k = 0;
    for (int i = 0; i < n; i++) {
        int v = lit_value(s, lits[i]);
        if (v == 1) { free(lits); return 0; }
        if (v == 0) continue;
        lits[k++] = lits[i];
    }
    n = k;

    if (n == 0) { free(lits); s->unsat = 1; return 1; }

    if (n == 1) {
        enqueue(s, lits[0], NULL);
        free(lits);
        Clause *c = propagate(s);
        if (c) { s->unsat = 1; return 1; }
        return 0;
    }

    Clause *c = clause_new(n, lits, 0);
    free(lits);
    cl_push(&s->clauses, c);
    cl_push(&s->watches[c->lits[0]], c);
    cl_push(&s->watches[c->lits[1]], c);
    return 0;
}

/* ====== search ====== */
static int search(Sat *s, long max_conflicts) {
    long conflicts_here = 0;
    Lit *learnt = malloc((size_t)(s->nvars + 1) * sizeof(Lit));

    while (1) {
        Clause *confl = propagate(s);

        /* If SAT-level propagation is clean, give the theory a turn. */
        if (!confl && s->theory) {
            int *t_clause = NULL;
            int  t_n = 0;
            if (s->theory->check(s->theory->data, &t_clause, &t_n)) {
                /* An empty theory conflict means the problem is
                 * structurally UNSAT regardless of the trail (the
                 * theory found a contradiction whose proof uses no
                 * SAT literals -- e.g., a Hall violation among
                 * unmodified domain-defined bounds). */
                if (t_n == 0) { free(learnt); return SAT_UNSATISFIABLE; }
                confl = make_theory_clause(s, t_clause, t_n);
                s->theory_conflicts++;
            } else if (s->theory->propagate) {
                /* No conflict; ask the theory for an implied literal.
                 * If we get one, install its reason clause, enqueue the
                 * literal, and re-run the propagation cycle. */
                int  prop_lit = 0;
                int *reason   = NULL;
                int  reason_n = 0;
                if (s->theory->propagate(s->theory->data,
                                         &prop_lit, &reason, &reason_n))
                {
                    Clause *r = make_theory_reason(s, reason, reason_n);
                    enqueue(s, ext2int(prop_lit), r);
                    s->theory_props++;
                    continue;
                }
            }
        }

        if (confl) {
            s->conflicts++;
            conflicts_here++;

            if (s->level_n == 0) { free(learnt); return SAT_UNSATISFIABLE; }

            int btlevel;
            int learnt_n = analyze(s, confl, learnt, &btlevel);
            backtrack_to(s, btlevel);

            if (learnt_n == 1) {
                enqueue(s, learnt[0], NULL);
            } else {
                Clause *c = clause_new(learnt_n, learnt, 1);
                cl_push(&s->learnts, c);
                cl_push(&s->watches[c->lits[0]], c);
                cl_push(&s->watches[c->lits[1]], c);
                s->learned_count++;
                enqueue(s, learnt[0], c);
            }
            decay_var_act(s);

            if (s->conflict_limit > 0 && s->conflicts >= s->conflict_limit) {
                free(learnt);
                return SAT_UNKNOWN;
            }
        } else {
            if (max_conflicts >= 0 && conflicts_here >= max_conflicts) {
                backtrack_to(s, 0);
                free(learnt);
                return SAT_UNKNOWN;
            }
            Lit p = pick_branch_lit(s);
            if (p == LIT_UNDEF) { free(learnt); return SAT_SATISFIABLE; }
            s->decisions++;
            new_decision_level(s);
            enqueue(s, p, NULL);
        }
    }
}

/* ====== public API ====== */
Sat *sat_new(int n_atoms) {
    if (n_atoms < 1) return NULL;
    Sat *s = calloc(1, sizeof(Sat));
    s->nvars = n_atoms;
    s->value    = malloc((size_t)n_atoms * sizeof(int));
    s->level    = calloc((size_t)n_atoms, sizeof(int));
    s->reason   = calloc((size_t)n_atoms, sizeof(Clause *));
    s->phase    = calloc((size_t)n_atoms, sizeof(int));
    s->activity = calloc((size_t)n_atoms, sizeof(double));
    s->seen     = calloc((size_t)n_atoms, sizeof(char));
    for (int i = 0; i < n_atoms; i++) s->value[i] = -1;

    s->trail_cap = n_atoms;
    s->trail = malloc((size_t)s->trail_cap * sizeof(Lit));
    s->trail_n = 0; s->qhead = 0;
    s->level_cap = 16;
    s->trail_lim = malloc((size_t)s->level_cap * sizeof(int));
    s->level_n = 0;

    s->watches = calloc((size_t)(2 * n_atoms), sizeof(CList));

    s->heap     = malloc((size_t)n_atoms * sizeof(Var));
    s->heap_pos = malloc((size_t)n_atoms * sizeof(int));
    s->heap_n = 0;
    for (int v = 0; v < n_atoms; v++) { s->heap_pos[v] = -1; heap_push(s, v); }

    s->var_inc = 1.0;
    s->var_decay = 0.95;
    s->restart_unit = 64;
    s->conflict_limit = -1;
    return s;
}

void sat_free(Sat *s) {
    if (!s) return;
    for (int i = 0; i < s->clauses.n; i++) free(s->clauses.buf[i]);
    for (int i = 0; i < s->learnts.n; i++) free(s->learnts.buf[i]);
    free(s->clauses.buf); free(s->learnts.buf);
    if (s->watches) {
        for (int i = 0; i < 2 * s->nvars; i++) free(s->watches[i].buf);
        free(s->watches);
    }
    free(s->value); free(s->level); free(s->reason); free(s->phase);
    free(s->activity); free(s->seen);
    free(s->trail); free(s->trail_lim);
    free(s->heap); free(s->heap_pos);
    free(s);
}

void sat_set_theory(Sat *s, const TheoryHooks *t) { s->theory = t; }

int sat_solve(Sat *s) {
    if (s->unsat) return SAT_UNSATISFIABLE;

    Clause *c = propagate(s);
    if (c && s->level_n == 0) return SAT_UNSATISFIABLE;

    s->luby_i = 0;
    while (1) {
        long budget = (long)s->restart_unit * luby(s->luby_i);
        int r = search(s, budget);
        if (r != SAT_UNKNOWN) return r;
        s->restarts++;
        s->luby_i++;
        if (s->conflict_limit > 0 && s->conflicts >= s->conflict_limit)
            return SAT_UNKNOWN;
    }
}

int sat_value(const Sat *s, int atom) {
    if (atom < 1 || atom > s->nvars) return SAT_UNDEF;
    int x = s->value[atom - 1];
    if (x < 0) return SAT_UNDEF;
    return x ? SAT_TRUE : SAT_FALSE;
}

void sat_set_conflict_limit(Sat *s, long l) { s->conflict_limit = l; }
void sat_set_verbose       (Sat *s, int l)  { s->verbose = l; }

long sat_decisions       (const Sat *s) { return s->decisions; }
long sat_conflicts       (const Sat *s) { return s->conflicts; }
long sat_theory_conflicts(const Sat *s) { return s->theory_conflicts; }
long sat_theory_props    (const Sat *s) { return s->theory_props; }
long sat_propagations    (const Sat *s) { return s->propagations; }
long sat_restarts        (const Sat *s) { return s->restarts; }
long sat_learned_clauses (const Sat *s) { return s->learned_count; }
