/*
 * maxsat.c -- Phase 1 MaxSAT implementation: model-improving search.
 *
 *
 *   ALGORITHM
 *
 *   For each soft clause C_i we introduce a fresh "blocking" variable
 *   b_i and post the relaxed clause (C_i OR b_i). If the SAT engine
 *   wants to falsify C_i it must set b_i = true to "pay" for it. The
 *   cost of a model is the number of blocking variables set to true.
 *
 *   The model-improving loop is:
 *
 *       cost = +infinity                      (no model yet)
 *       loop:
 *           build a fresh SAT engine with:
 *             - all hard clauses
 *             - all relaxed soft clauses
 *             - if we have a model: at-most-(cost - 1) over the b_i's
 *           solve.
 *           if UNSAT:
 *               break.  the previous model is optimal.
 *           else:
 *               record the model; update cost to the new (lower) value.
 *
 *   The first iteration has no cardinality constraint --- we just
 *   check whether the hard clauses are satisfiable at all. Subsequent
 *   iterations force strict improvement.
 *
 *   The cardinality constraint sum(b_i) <= k uses Sinz's sequential
 *   counter encoding: O(n*k) clauses and O(n*k) auxiliary variables.
 *   It is sound, complete, and arc-consistent.
 *
 *
 *   LIMITATIONS OF THIS APPROACH
 *
 *   Each iteration rebuilds the SAT engine from scratch. The learnt
 *   clauses from previous iterations are discarded. For problems with
 *   many soft clauses this is wasteful; on Phase 2 we will replace
 *   it with core-guided search that reuses learning effectively.
 */
#define _POSIX_C_SOURCE 200809L
#include "maxsat.h"
#include "sat.h"

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

/* ====== buffered clause ====== */
typedef struct {
    int *lits;
    int  n;
} Clause;

struct MaxSatCtx {
    int n_vars;                /* highest user-allocated var */

    Clause *hards;
    int     n_hards, cap_hards;

    Clause *softs;
    int     n_softs, cap_softs;

    /* Solver configuration */
    MaxSatAlgorithm algorithm;

    /* Solved state */
    int  solved;
    int *best_model;           /* indexed 1..n_vars; 0 if false, 1 if true */
    long best_cost;
    long iterations;
    long total_decisions;
    long cores_found;          /* core-guided only */
};

/* ====== buffer helpers ====== */
static void clause_buf_push(Clause **arr, int *n, int *cap,
                            const int *lits, int nlits) {
    if (*n == *cap) {
        *cap = *cap ? *cap * 2 : 16;
        *arr = realloc(*arr, (size_t)*cap * sizeof(Clause));
    }
    Clause *c = &(*arr)[(*n)++];
    c->n = nlits;
    c->lits = malloc((size_t)nlits * sizeof(int));
    if (nlits > 0) memcpy(c->lits, lits, (size_t)nlits * sizeof(int));
}

/* ====== lifecycle ====== */
MaxSatCtx *maxsat_new(void) {
    MaxSatCtx *m = calloc(1, sizeof(MaxSatCtx));
    m->algorithm = MAXSAT_ALG_CORE_GUIDED;
    return m;
}

void maxsat_set_algorithm(MaxSatCtx *m, MaxSatAlgorithm alg) {
    m->algorithm = alg;
}

MaxSatAlgorithm maxsat_get_algorithm(const MaxSatCtx *m) {
    return m->algorithm;
}

void maxsat_free(MaxSatCtx *m) {
    if (!m) return;
    for (int i = 0; i < m->n_hards; i++) free(m->hards[i].lits);
    free(m->hards);
    for (int i = 0; i < m->n_softs; i++) free(m->softs[i].lits);
    free(m->softs);
    free(m->best_model);
    free(m);
}

int maxsat_new_var(MaxSatCtx *m) {
    return ++m->n_vars;
}

void maxsat_add_hard_clause(MaxSatCtx *m, const int *lits, int n) {
    clause_buf_push(&m->hards, &m->n_hards, &m->cap_hards, lits, n);
}

void maxsat_add_soft_clause(MaxSatCtx *m, const int *lits, int n) {
    clause_buf_push(&m->softs, &m->n_softs, &m->cap_softs, lits, n);
}

void maxsat_add_soft_clause_weighted(MaxSatCtx *m, const int *lits, int n,
                                     long weight) {
    if (weight <= 0) return;
    for (long w = 0; w < weight; w++) {
        clause_buf_push(&m->softs, &m->n_softs, &m->cap_softs, lits, n);
    }
}

void maxsat_add_at_most_one(MaxSatCtx *m, const int *lits, int n) {
    /* Pairwise: for every i<j, the clause (¬lits[i] ∨ ¬lits[j]). */
    int cl[2];
    for (int i = 0; i < n; i++) {
        for (int j = i + 1; j < n; j++) {
            cl[0] = -lits[i];
            cl[1] = -lits[j];
            maxsat_add_hard_clause(m, cl, 2);
        }
    }
}

/* ====== at-most-k via Sinz's sequential counter ======
 *
 * For at-most-k over `vars[0..n-1]`, we introduce auxiliary variables
 *
 *     s_{i,j}  for i in 1..n-1, j in 1..k,
 *
 * with the intended meaning "at least j of vars[0..i-1] are true".
 * The encoding uses (n-1)*k aux variables, allocated contiguously
 * starting at `aux_base + 1`.
 *
 * The clauses (numbered as in Sinz 2005) are:
 *
 *   (3) x_1 -> s_{1,1}
 *   (4) ~s_{1,j}                              for j in 2..k
 *   (5) x_i -> s_{i,1}                        for 1 < i < n
 *   (6) s_{i-1,j} -> s_{i,j}                  for 1 < i < n, j in 1..k
 *   (7) x_i AND s_{i-1,j-1} -> s_{i,j}        for 1 < i < n, j in 2..k
 *   (8) x_i -> ~s_{i-1,k}                     for 1 < i <= n
 *
 * Returns the new highest variable id (aux_base + (n-1)*k). */
static int add_at_most_k(Sat *s, const int *vars, int n, int k, int aux_base) {
    if (k >= n) return aux_base;       /* trivially satisfied */
    if (k == 0) {                       /* forbid every variable individually */
        for (int i = 0; i < n; i++) {
            int cl = -vars[i];
            sat_add_clause(s, &cl, 1);
        }
        return aux_base;
    }

    /* S(i, j) maps to a positive variable id. */
    #define S(i, j)  (aux_base + ((i) - 1) * k + (j))

    int cl[3];

    /* (3) x_1 -> s_{1,1} */
    cl[0] = -vars[0];
    cl[1] = S(1, 1);
    sat_add_clause(s, cl, 2);

    /* (4) ¬s_{1,j} for j > 1 */
    for (int j = 2; j <= k; j++) {
        cl[0] = -S(1, j);
        sat_add_clause(s, cl, 1);
    }

    for (int i = 2; i < n; i++) {
        /* (5) x_i -> s_{i,1} */
        cl[0] = -vars[i - 1];
        cl[1] = S(i, 1);
        sat_add_clause(s, cl, 2);

        /* (6) s_{i-1,j} -> s_{i,j} for j in 1..k */
        for (int j = 1; j <= k; j++) {
            cl[0] = -S(i - 1, j);
            cl[1] = S(i, j);
            sat_add_clause(s, cl, 2);
        }

        /* (7) x_i AND s_{i-1,j-1} -> s_{i,j} for j in 2..k */
        for (int j = 2; j <= k; j++) {
            cl[0] = -vars[i - 1];
            cl[1] = -S(i - 1, j - 1);
            cl[2] = S(i, j);
            sat_add_clause(s, cl, 3);
        }
    }

    /* (8) x_i -> ¬s_{i-1,k} for 1 < i <= n */
    for (int i = 2; i <= n; i++) {
        cl[0] = -vars[i - 1];
        cl[1] = -S(i - 1, k);
        sat_add_clause(s, cl, 2);
    }

    #undef S
    return aux_base + (n - 1) * k;
}

/* ====== one iteration of model-improving search ======
 *
 * Build a fresh SAT engine with all hard clauses, all soft clauses
 * relaxed with their blocking variables, and (if `cardinality_k >= 0`)
 * the constraint sum(b_i) <= cardinality_k.
 *
 * On SAT: cost is written to *out_cost, the model to model[1..orig_n_vars],
 *         and 1 is returned.
 * On UNSAT: returns 0.
 *
 * Also accumulates SAT-decision stats into *out_decisions. */
static int run_iteration(const MaxSatCtx *m,
                         const int *blocking, int n_blocking,
                         long cardinality_k,
                         int *model, int orig_n_vars,
                         long *out_cost, long *out_decisions)
{
    /* Worst-case aux variable count for at-most-k. */
    int aux_max = 0;
    if (cardinality_k >= 0 && cardinality_k < n_blocking) {
        aux_max = (n_blocking - 1) * (int)cardinality_k;
    }
    int total_vars = m->n_vars + aux_max;
    if (total_vars < m->n_vars) total_vars = m->n_vars;

    Sat *s = sat_new(total_vars > 0 ? total_vars : 1);

    for (int i = 0; i < m->n_hards; i++)
        sat_add_clause(s, m->hards[i].lits, m->hards[i].n);

    /* Relaxed soft clauses: (C_i OR b_i). */
    int *ext = NULL;
    int  ext_cap = 0;
    for (int i = 0; i < m->n_softs; i++) {
        Clause *c = &m->softs[i];
        if (c->n + 1 > ext_cap) {
            ext_cap = c->n + 1;
            ext = realloc(ext, (size_t)ext_cap * sizeof(int));
        }
        if (c->n > 0) memcpy(ext, c->lits, (size_t)c->n * sizeof(int));
        ext[c->n] = blocking[i];
        sat_add_clause(s, ext, c->n + 1);
    }
    free(ext);

    /* Cardinality bound. */
    if (cardinality_k >= 0) {
        add_at_most_k(s, blocking, n_blocking, (int)cardinality_k, m->n_vars);
    }

    int r = sat_solve(s);
    *out_decisions += sat_decisions(s);

    if (r == SAT_UNSATISFIABLE) {
        sat_free(s);
        return 0;
    }

    long cost = 0;
    for (int i = 0; i < n_blocking; i++)
        if (sat_value(s, blocking[i]) == 1) cost++;

    for (int v = 1; v <= orig_n_vars; v++)
        model[v] = (sat_value(s, v) == 1) ? 1 : 0;

    *out_cost = cost;
    sat_free(s);
    return 1;
}

static long solve_linear(MaxSatCtx *m) {
    int orig_n_vars = m->n_vars;
    int *model = calloc((size_t)(orig_n_vars + 1), sizeof(int));

    /* No soft clauses: feasibility only. */
    if (m->n_softs == 0) {
        long unused = 0;
        int  ok = run_iteration(m, NULL, 0, -1, model, orig_n_vars,
                                &unused, &m->total_decisions);
        m->iterations = 1;
        if (!ok) {
            free(model);
            m->best_cost = MAXSAT_HARD_UNSAT;
            return MAXSAT_HARD_UNSAT;
        }
        m->best_model = model;
        m->best_cost = 0;
        return 0;
    }

    /* Allocate blocking vars in the MaxSatCtx var space. */
    int *blocking = malloc((size_t)m->n_softs * sizeof(int));
    for (int i = 0; i < m->n_softs; i++)
        blocking[i] = ++m->n_vars;

    long current_cost = -1;
    while (1) {
        m->iterations++;
        long this_cost = 0;
        long card_k = (current_cost < 0) ? -1 : current_cost - 1;
        int  ok = run_iteration(m, blocking, m->n_softs, card_k,
                                model, orig_n_vars,
                                &this_cost, &m->total_decisions);
        if (!ok) {
            if (current_cost < 0) {
                /* The very first call failed: hard clauses alone are
                 * infeasible. */
                free(blocking);
                free(model);
                m->best_cost = MAXSAT_HARD_UNSAT;
                return MAXSAT_HARD_UNSAT;
            }
            break;       /* previous current_cost is the optimum */
        }
        current_cost = this_cost;
        if (current_cost == 0) break;
    }

    free(blocking);
    m->best_model = model;
    m->best_cost  = current_cost;
    return current_cost;
}

/* ====== Phase 2: core-guided MaxSAT (Fu-Malik) ======
 *
 * Each soft clause C_i has a chain of selector variables. Initially
 * one: s_i^0. The relaxed hard clause is (C_i OR s_i^0).
 *
 * Each iteration:
 *
 *   1. Build a SAT formula with:
 *        - all hard clauses
 *        - for each soft i: (C_i OR all_selectors[i])
 *        - all accumulated at-most-one constraints over previous cores
 *        - unit clauses ¬cur_selector[i] for every soft (the "assumption"
 *          that this soft must hold)
 *
 *   2. Solve.
 *      - If SAT: done. The cost is the number of cores found so far,
 *        because each core's at-most-one forces at most one of its
 *        old selectors to flip true, and at least one must (the core
 *        is unsat without the relaxation).
 *      - If UNSAT: a subset of the soft clauses is jointly infeasible
 *        under "all must hold." Extract the core.
 *
 *   3. Relax the core: for each soft i in the core, allocate a fresh
 *      selector and append it to the soft's chain. The OLD selector
 *      becomes a "release indicator." Then add at-most-one over the
 *      old selectors of the core.
 *
 *   4. cost := cost + 1, loop.
 *
 * Core extraction (Phase 2): deletion-based. Re-solve the formula
 * with each assumption removed in turn; if the formula is still UNSAT
 * without it, the assumption was not in the core. This costs O(n_softs)
 * SAT calls per core. A real solver would use proper SAT-engine
 * assumptions and read the core off the conflict clause for free.
 */

/* Build a fresh SAT engine populated with hard clauses, all relaxed
 * soft clauses (with their full selector chain), and all accumulated
 * core constraints. The caller is responsible for sat_free. */
static Sat *build_engine(const MaxSatCtx *m,
                         int **selectors_list, const int *selectors_n,
                         const Clause *core_constraints, int n_core_constraints,
                         int total_vars)
{
    if (total_vars < 1) total_vars = 1;
    Sat *s = sat_new(total_vars);

    for (int i = 0; i < m->n_hards; i++)
        sat_add_clause(s, m->hards[i].lits, m->hards[i].n);

    int  ext_cap = 0;
    int *ext = NULL;
    for (int i = 0; i < m->n_softs; i++) {
        int n_lits = m->softs[i].n + selectors_n[i];
        if (n_lits > ext_cap) {
            ext_cap = n_lits;
            ext = realloc(ext, (size_t)ext_cap * sizeof(int));
        }
        if (m->softs[i].n > 0)
            memcpy(ext, m->softs[i].lits, (size_t)m->softs[i].n * sizeof(int));
        memcpy(ext + m->softs[i].n, selectors_list[i],
               (size_t)selectors_n[i] * sizeof(int));
        sat_add_clause(s, ext, n_lits);
    }
    free(ext);

    for (int i = 0; i < n_core_constraints; i++)
        sat_add_clause(s, core_constraints[i].lits, core_constraints[i].n);

    return s;
}

/* Solve the formula under the "all-softs-hold" assumption, restricted
 * to a subset of softs given by `active_indices[0..n_active-1]`. Returns
 * 1 on SAT (and writes the model and per-soft selector values), 0 on UNSAT.
 *
 * Assumptions are encoded as unit clauses on the fly. */
static int solve_with_active_set(const MaxSatCtx *m,
                                 int **selectors_list, const int *selectors_n,
                                 const int *current_selector,
                                 const Clause *core_constraints,
                                 int n_core_constraints,
                                 int total_vars,
                                 const int *active_indices, int n_active,
                                 int *model_out, int orig_n_vars,
                                 long *decisions_inout)
{
    Sat *s = build_engine(m, selectors_list, selectors_n,
                          core_constraints, n_core_constraints, total_vars);
    for (int k = 0; k < n_active; k++) {
        int i = active_indices[k];
        int neg_sel = -current_selector[i];
        sat_add_clause(s, &neg_sel, 1);
    }
    int r = sat_solve(s);
    *decisions_inout += sat_decisions(s);
    int sat = (r == SAT_SATISFIABLE);
    if (sat && model_out) {
        for (int v = 1; v <= orig_n_vars; v++)
            model_out[v] = (sat_value(s, v) == 1) ? 1 : 0;
    }
    sat_free(s);
    return sat;
}

static long solve_core_guided(MaxSatCtx *m) {
    int orig_n_vars = m->n_vars;
    int *model = calloc((size_t)(orig_n_vars + 1), sizeof(int));

    /* Edge case: no softs, just feasibility. */
    if (m->n_softs == 0) {
        long unused = 0;
        int ok = run_iteration(m, NULL, 0, -1, model, orig_n_vars,
                               &unused, &m->total_decisions);
        m->iterations = 1;
        if (!ok) { free(model); m->best_cost = MAXSAT_HARD_UNSAT; return MAXSAT_HARD_UNSAT; }
        m->best_model = model; m->best_cost = 0; return 0;
    }

    /* Per-soft selector chain. */
    int **selectors_list  = calloc((size_t)m->n_softs, sizeof(int *));
    int  *selectors_n     = calloc((size_t)m->n_softs, sizeof(int));
    int  *selectors_cap   = calloc((size_t)m->n_softs, sizeof(int));
    int  *current_selector = calloc((size_t)m->n_softs, sizeof(int));
    for (int i = 0; i < m->n_softs; i++) {
        selectors_cap[i] = 4;
        selectors_list[i] = malloc(sizeof(int) * 4);
        int s = ++m->n_vars;
        selectors_list[i][0] = s;
        selectors_n[i] = 1;
        current_selector[i] = s;
    }

    /* Accumulated at-most-one constraints over old selectors. Each
     * core's relaxation adds (n_core choose 2) binary clauses here. */
    Clause *core_constraints = NULL;
    int n_core_constraints = 0, cap_core_constraints = 0;

    /* Active set: indices of softs whose assumption we still post.
     * Initially all of them; deletion-based extraction shrinks it
     * temporarily during a core-extraction pass but resets each
     * outer iteration. */
    int *active = malloc((size_t)m->n_softs * sizeof(int));

    long cost = 0;
    while (1) {
        m->iterations++;

        /* Outer solve under "all softs must hold." */
        for (int i = 0; i < m->n_softs; i++) active[i] = i;
        int got_sat = solve_with_active_set(m, selectors_list, selectors_n,
                                            current_selector,
                                            core_constraints, n_core_constraints,
                                            m->n_vars,
                                            active, m->n_softs,
                                            model, orig_n_vars,
                                            &m->total_decisions);
        if (got_sat) break;

        /* UNSAT under all assumptions: extract a core by deletion. */
        int *candidate = malloc((size_t)m->n_softs * sizeof(int));
        int  n_cand = m->n_softs;
        for (int i = 0; i < m->n_softs; i++) candidate[i] = i;

        int idx = 0;
        while (idx < n_cand) {
            /* Try removing candidate[idx]: solve without that assumption. */
            int *trial = malloc((size_t)(n_cand - 1) * sizeof(int));
            int  n_trial = 0;
            for (int k = 0; k < n_cand; k++)
                if (k != idx) trial[n_trial++] = candidate[k];

            int still_unsat = !solve_with_active_set(m, selectors_list, selectors_n,
                                                    current_selector,
                                                    core_constraints, n_core_constraints,
                                                    m->n_vars,
                                                    trial, n_trial,
                                                    NULL, orig_n_vars,
                                                    &m->total_decisions);
            free(trial);

            if (still_unsat) {
                /* This assumption isn't required; drop it. */
                for (int k = idx; k < n_cand - 1; k++)
                    candidate[k] = candidate[k + 1];
                n_cand--;
            } else {
                /* Necessary; keep and advance. */
                idx++;
            }
        }

        if (n_cand == 0) {
            /* No assumption was responsible: hard clauses alone are infeasible. */
            free(candidate); free(active); free(model);
            for (int i = 0; i < m->n_softs; i++) free(selectors_list[i]);
            free(selectors_list); free(selectors_n); free(selectors_cap); free(current_selector);
            for (int i = 0; i < n_core_constraints; i++) free(core_constraints[i].lits);
            free(core_constraints);
            m->best_cost = MAXSAT_HARD_UNSAT;
            return MAXSAT_HARD_UNSAT;
        }

        m->cores_found++;
        cost++;

        /* Relax: for each soft in the core, allocate a fresh selector,
         * keep the old one (now "freed"), and add at-most-one over
         * the old ones. */
        int *old_selectors = malloc((size_t)n_cand * sizeof(int));
        for (int k = 0; k < n_cand; k++) {
            int i = candidate[k];
            old_selectors[k] = current_selector[i];

            int s_new = ++m->n_vars;
            if (selectors_n[i] == selectors_cap[i]) {
                selectors_cap[i] *= 2;
                selectors_list[i] = realloc(selectors_list[i],
                                            (size_t)selectors_cap[i] * sizeof(int));
            }
            selectors_list[i][selectors_n[i]++] = s_new;
            current_selector[i] = s_new;
        }

        /* Pairwise at-most-one over old selectors. */
        for (int k = 0; k < n_cand; k++) {
            for (int l = k + 1; l < n_cand; l++) {
                if (n_core_constraints == cap_core_constraints) {
                    cap_core_constraints = cap_core_constraints ? cap_core_constraints * 2 : 16;
                    core_constraints = realloc(core_constraints,
                                               (size_t)cap_core_constraints * sizeof(Clause));
                }
                Clause *c = &core_constraints[n_core_constraints++];
                c->n = 2;
                c->lits = malloc(2 * sizeof(int));
                c->lits[0] = -old_selectors[k];
                c->lits[1] = -old_selectors[l];
            }
        }
        free(old_selectors);
        free(candidate);
    }

    free(active);
    for (int i = 0; i < m->n_softs; i++) free(selectors_list[i]);
    free(selectors_list); free(selectors_n); free(selectors_cap); free(current_selector);
    for (int i = 0; i < n_core_constraints; i++) free(core_constraints[i].lits);
    free(core_constraints);

    m->best_model = model;
    m->best_cost  = cost;
    return cost;
}

/* ====== dispatcher ====== */
long maxsat_solve(MaxSatCtx *m) {
    if (m->solved) return m->best_cost;
    m->solved = 1;
    if (m->algorithm == MAXSAT_ALG_LINEAR) return solve_linear(m);
    return solve_core_guided(m);
}

int maxsat_value(const MaxSatCtx *m, int var) {
    if (!m->best_model || var < 1 || var > m->n_vars) return 0;
    return m->best_model[var];
}

long maxsat_iterations     (const MaxSatCtx *m) { return m->iterations; }
long maxsat_total_decisions(const MaxSatCtx *m) { return m->total_decisions; }
long maxsat_cores_found    (const MaxSatCtx *m) { return m->cores_found; }
