/*
 * sat.c -- minimal CDCL SAT solver
 *
 * Implements:
 *
 *   - Two-watched-literal propagation
 *   - First-UIP conflict analysis with clause learning
 *   - Non-chronological backjumping
 *   - VSIDS variable activity heuristic on a binary heap
 *   - Phase saving
 *   - Luby restart sequence (unit = 64 conflicts)
 *
 * Not implemented (tutorial scope):
 *
 *   - Learned-clause minimization (self-subsuming resolution)
 *   - Clause database reduction
 *   - Preprocessing (variable elimination, subsumption)
 *   - Incremental / assumption-based solving
 *
 * Internal literal encoding: variable v in [0..n_vars-1], literal is
 * (v << 1) | sign with sign=0 for positive, sign=1 for negative.
 * Negation is just lit ^ 1.
 *
 * DIMACS external lit +k maps to (k-1)<<1, -k maps to ((k-1)<<1)|1.
 */

#include "sat.h"

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

/* ===================================================================
 * Literal / variable 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)

static Lit ext2int(int e) {
    int v = (e > 0 ? e : -e) - 1;
    return MK_LIT(v, e < 0 ? 1 : 0);
}

/* ===================================================================
 * Clause
 * =================================================================== */

typedef struct Clause {
    int  size;
    int  learnt;
    Lit  lits[1];   /* flexible: real size given by `size` */
} Clause;

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

/* Dynamic array of Clause* (used for clause storage and watch lists). */
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;          /* trivially UNSAT? */
    int verbose;

    /* Per-variable arrays (size nvars). */
    int     *value;     /* 1, 0, or -1 (undef) */
    int     *level;     /* decision level when assigned (0 if undef) */
    Clause **reason;    /* propagating clause, or NULL for decision / undef */
    int     *phase;     /* last assigned value, used for branching */
    double  *activity;
    char    *seen;      /* scratch for analyze() */

    /* Trail (assignment stack). */
    Lit *trail;
    int  trail_n;
    int  trail_cap;
    int  qhead;          /* propagation queue head into trail */

    /* Decision levels: trail_lim[d] = first trail index at level d+1. */
    int *trail_lim;
    int  level_n;        /* current decision level */
    int  level_cap;

    /* Watch lists: one CList per literal (so 2*nvars lists). */
    CList *watches;

    /* Clause storage. */
    CList clauses;
    CList learnts;

    /* VSIDS activity heap (binary max-heap on activity). */
    Var    *heap;        /* heap[0..heap_n-1] */
    int    *heap_pos;    /* heap_pos[v] = position in heap, or -1 */
    int     heap_n;
    double  var_inc;
    double  var_decay;

    /* Restart. */
    int   luby_i;
    long  restart_unit;

    /* Statistics. */
    long decisions;
    long propagations;
    long conflicts;
    long restarts;
    long learned_count;

    /* Limits. */
    long conflict_limit;
};

/* ===================================================================
 * VSIDS binary heap (max-heap by activity)
 * =================================================================== */

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]);
}

/* ===================================================================
 * Activity bookkeeping
 * =================================================================== */

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 / trail / backtracking
 * =================================================================== */

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

static void enqueue(Sat *s, Lit l, Clause *reason) {
    Var v = VAR(l);
    int val = !SIGN(l);                          /* +l -> 1, -l -> 0 */
    s->value[v]  = val;
    s->level[v]  = s->level_n;
    s->reason[v] = reason;
    s->phase[v]  = val;
    s->trail[s->trail_n++] = 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--) {
        Var v = VAR(s->trail[i]);
        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;
}

/* ===================================================================
 * Unit propagation with two watched literals
 *
 * Invariant maintained between propagations: for every clause of size >= 2,
 * c->lits[0] and c->lits[1] are either both non-false, or at least one is
 * true (which keeps the clause satisfied without needing to look at the rest).
 *
 * When a literal becomes assigned (so its negation becomes false), we walk
 * the list of clauses watching that negation and try to find a replacement
 * watch elsewhere in the clause. If we can't, the clause is unit on the
 * other watched literal -- propagate it -- or it's a conflict.
 * =================================================================== */

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];

            /* Canonicalize: make sure c->lits[1] is the now-false watch. */
            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 the other watch is already true, the clause is satisfied;
             * leave the watch where it is. */
            if (lit_value(s, other) == 1) {
                buf[j++] = c;
                continue;
            }

            /* Look for a new watch among c->lits[2..size-1]. */
            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;

            /* No replacement: clause is unit on `other`, or false. */
            buf[j++] = c;

            if (lit_value(s, other) == 0) {
                /* Conflict. Preserve the rest of this watch list and abort. */
                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
 *
 * Walks the implication graph backwards from the conflict, resolving with
 * each propagating reason clause, marking variables we've seen. We stop
 * when exactly one literal at the current decision level remains -- that
 * is the First Unique Implication Point. Its negation is the asserting
 * literal; the rest of the marked literals (all from lower levels) form
 * the body of the learned clause.
 *
 * Returns the size of the learned clause. The clause is in `out` with
 * out[0] = asserting literal. The backjump level is written to *btlevel.
 * =================================================================== */

static int analyze(Sat *s, Clause *confl, Lit *out, int *btlevel) {
    int counter = 0;       /* unresolved literals at current level */
    int trail_idx = s->trail_n - 1;
    int learnt_n = 1;      /* reserve out[0] for the asserting literal */
    int cur_level = s->level_n;
    Lit p = LIT_UNDEF;

    do {
        /* Walk the literals of `confl`, skipping lits[0] when this isn't
         * the first iteration (because lits[0] of a reason clause is the
         * literal that was propagated -- which is what we just resolved on). */
        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;
            }
        }

        /* Find next seen variable, walking back along the trail. */
        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);

    /* p is the 1UIP at the current level. */
    out[0] = NEG(p);

    /* Backjump level = max level among out[1..]. */
    *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; }
    }
    /* Move highest-level lit to out[1] so the new clause has its two
     * watched literals at out[0] (current level, will be propagated after
     * backjump) and out[1] (the backjump level). */
    if (learnt_n > 1 && max_i != 1) {
        Lit t = out[1]; out[1] = out[max_i]; out[max_i] = t;
    }

    /* Clear seen flags on lower-level lits (current-level seen flags were
     * cleared while popping the trail). */
    for (int i = 1; i < learnt_n; i++) s->seen[VAR(out[i])] = 0;

    return learnt_n;
}

/* ===================================================================
 * Branch literal selection (VSIDS + phase saving)
 * =================================================================== */

static Lit pick_branch_lit(Sat *s) {
    while (s->heap_n > 0) {
        Var v = heap_pop(s);
        if (s->value[v] == -1) {
            /* Saved phase: 1 -> positive lit, 0 -> negative lit. */
            return MK_LIT(v, s->phase[v] ? 0 : 1);
        }
    }
    return LIT_UNDEF;  /* every variable assigned */
}

/* ===================================================================
 * Luby restart sequence (0-indexed: 1,1,2,1,1,2,4,1,1,2,1,1,2,4,8,...)
 * =================================================================== */

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;
}

/* ===================================================================
 * Adding clauses
 * =================================================================== */

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;
    }

    /* Sort to detect duplicates and tautologies in O(n log n). */
    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;          /* dup */
        if (m > 0 && lits[m - 1] == NEG(lits[i])) {             /* taut */
            free(lits); return 0;
        }
        lits[m++] = lits[i];
    }
    n = m;

    /* Check against any existing level-0 assignments. */
    int k = 0;
    for (int i = 0; i < n; i++) {
        int v = lit_value(s, lits[i]);
        if (v == 1) { free(lits); return 0; }                   /* clause sat */
        if (v == 0) continue;                                   /* drop false */
        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);
        /* Propagate immediately to detect early infeasibility. */
        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 core: one "stretch" between restarts.
 * Returns SAT_*, or SAT_UNKNOWN if the conflict budget was exhausted.
 * =================================================================== */

static int search(Sat *s, long max_conflicts) {
    long conflicts_here = 0;

    /* Scratch buffer for learnt clauses (reused across iterations). */
    Lit *learnt = malloc((size_t)(s->nvars + 1) * sizeof(Lit));

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

        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_vars) {
    if (n_vars < 1) return NULL;

    Sat *s = calloc(1, sizeof(Sat));
    s->nvars = n_vars;

    s->value    = malloc((size_t)n_vars * sizeof(int));
    s->level    = calloc((size_t)n_vars, sizeof(int));
    s->reason   = calloc((size_t)n_vars, sizeof(Clause *));
    s->phase    = calloc((size_t)n_vars, sizeof(int));
    s->activity = calloc((size_t)n_vars, sizeof(double));
    s->seen     = calloc((size_t)n_vars, sizeof(char));
    for (int i = 0; i < n_vars; i++) s->value[i] = -1;

    s->trail_cap = n_vars;
    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_vars), sizeof(CList));

    s->heap     = malloc((size_t)n_vars * sizeof(Var));
    s->heap_pos = malloc((size_t)n_vars * sizeof(int));
    s->heap_n = 0;
    for (int v = 0; v < n_vars; 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;
    s->verbose = 0;
    s->unsat = 0;
    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);
}

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

    /* Initial propagation in case there are pending unit clauses. */
    Clause *c = propagate(s);
    if (c) { if (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->verbose >= 2) {
            fprintf(stderr,
                "restart %ld  conflicts=%ld  learned=%ld  next_budget=%ld\n",
                s->restarts, s->conflicts, s->learned_count,
                (long)s->restart_unit * luby(s->luby_i));
        }
        if (s->conflict_limit > 0 && s->conflicts >= s->conflict_limit)
            return SAT_UNKNOWN;
    }
}

int sat_value(const Sat *s, int v) {
    if (v < 1 || v > s->nvars) return SAT_UNDEF;
    int x = s->value[v - 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_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; }
