/*
 * milp.c -- MILP solver implementation.
 *
 *   LP relaxation : two-phase Bland's-rule simplex on a dense tableau
 *   Integer part  : recursive depth-first branch-and-bound,
 *                   most-fractional variable selection,
 *                   bound-based pruning against the incumbent.
 *
 * Internal standard form (after shift y_j = x_j - lo_j):
 *
 *      max  c'y                (sense flipped if user said MIN)
 *      s.t. A y  =  b'         (slack/surplus/artificial added per row)
 *           y   >= 0
 *
 * Upper bounds y_j <= hi_j - lo_j are added as explicit <= rows.
 * Phase 1 drives all artificials to zero. Phase 2 optimises c'y
 * with artificials forbidden from re-entering the basis.
 */
#include "milp.h"
#include <stdlib.h>
#include <string.h>
#include <math.h>
#include <float.h>
#include <stdio.h>

static const double EPS_FEAS  = 1e-7;
static const double EPS_OPT   = 1e-7;
static const double EPS_PIVOT = 1e-9;
static const double EPS_INT   = 1e-6;
static const double INFTY     = 1e20;
static const int    MAX_PIVOTS = 200000;

struct MilpModel {
    int n, m;
    int sense;
    double *c;
    double *A;          /* m * n, row-major */
    double *b;          /* m */
    int    *op;         /* m */
    int    *type;       /* n */
    double *lo, *hi;    /* n */

    int verbose;
    int node_limit;

    int    nodes;
    int    lp_calls;
    double lp_relax;
};

/* ---------------------------------------------------------------------
 * Public model constructors / mutators
 * ------------------------------------------------------------------- */

MilpModel *milp_new(int n, int m) {
    MilpModel *p = calloc(1, sizeof(*p));
    p->n = n;
    p->m = m;
    p->sense = MILP_MIN;
    p->c    = calloc(n,     sizeof(double));
    p->A    = calloc((size_t)m * n, sizeof(double));
    p->b    = calloc(m,     sizeof(double));
    p->op   = calloc(m,     sizeof(int));
    p->type = calloc(n,     sizeof(int));
    p->lo   = calloc(n,     sizeof(double));
    p->hi   = malloc((size_t)n * sizeof(double));
    for (int j = 0; j < n; j++) p->hi[j] = INFTY;
    p->node_limit = 1 << 20;
    return p;
}

void milp_free(MilpModel *p) {
    if (!p) return;
    free(p->c);    free(p->A);  free(p->b);  free(p->op);
    free(p->type); free(p->lo); free(p->hi);
    free(p);
}

void milp_set_objective(MilpModel *m, int sense, const double *c) {
    m->sense = sense;
    memcpy(m->c, c, (size_t)m->n * sizeof(double));
}

void milp_set_row(MilpModel *m, int row, const double *a, int op, double rhs) {
    memcpy(m->A + (size_t)row * m->n, a, (size_t)m->n * sizeof(double));
    m->op[row] = op;
    m->b[row]  = rhs;
}

void milp_set_var_type(MilpModel *m, int var, int type) {
    m->type[var] = type;
    if (type == MILP_BINARY) {
        m->lo[var] = 0;
        m->hi[var] = 1;
    }
}

void milp_set_var_bounds(MilpModel *m, int var, double lo, double hi) {
    m->lo[var] = lo;
    m->hi[var] = hi;
}

void milp_set_verbose   (MilpModel *m, int level) { m->verbose = level; }
void milp_set_node_limit(MilpModel *m, int limit) { m->node_limit = limit; }

int    milp_nodes_explored(const MilpModel *m) { return m->nodes; }
int    milp_lps_solved    (const MilpModel *m) { return m->lp_calls; }
double milp_lp_relaxation (const MilpModel *m) { return m->lp_relax; }

/* ---------------------------------------------------------------------
 * Standard-form LP and simplex
 * ------------------------------------------------------------------- */

typedef struct {
    int rows;            /* constraint rows (excluding objective row 0) */
    int cols;            /* total columns excluding rhs */
    int n_orig;          /* original variable count */
    int art_start;       /* first artificial column index */
    int n_art;
    double *T;           /* (rows+1) * (cols+1), row-major */
    int    *basis;       /* rows */
} LP;

#define TIJ(L,i,j) ((L)->T[(size_t)(i) * ((L)->cols + 1) + (j)])

static void lp_free(LP *L) {
    if (!L) return;
    free(L->T); free(L->basis); free(L);
}

static LP *lp_build(const MilpModel *m) {
    int n = m->n, m_orig = m->m;

    /* Add explicit upper-bound rows (lower bounds handled via shift) */
    int n_hi = 0;
    for (int j = 0; j < n; j++)
        if (m->hi[j] < INFTY * 0.5) n_hi++;

    int rows = m_orig + n_hi;

    double *A  = calloc((size_t)rows * n, sizeof(double));
    double *b  = calloc(rows,             sizeof(double));
    int    *op = calloc(rows,             sizeof(int));

    /* Original constraints with shift y_j = x_j - lo_j */
    for (int i = 0; i < m_orig; i++) {
        double rhs = m->b[i];
        for (int j = 0; j < n; j++) {
            double aij = m->A[(size_t)i * n + j];
            A[(size_t)i * n + j] = aij;
            rhs -= aij * m->lo[j];
        }
        b[i]  = rhs;
        op[i] = m->op[i];
    }
    int row = m_orig;
    for (int j = 0; j < n; j++) {
        if (m->hi[j] < INFTY * 0.5) {
            A[(size_t)row * n + j] = 1.0;
            b[row]  = m->hi[j] - m->lo[j];
            op[row] = MILP_LE;
            row++;
        }
    }

    /* Force rhs >= 0 by row-flipping */
    for (int i = 0; i < rows; i++) {
        if (b[i] < 0) {
            for (int j = 0; j < n; j++)
                A[(size_t)i * n + j] = -A[(size_t)i * n + j];
            b[i] = -b[i];
            if      (op[i] == MILP_LE) op[i] = MILP_GE;
            else if (op[i] == MILP_GE) op[i] = MILP_LE;
        }
    }

    int n_slack = 0, n_surp = 0, n_art = 0;
    for (int i = 0; i < rows; i++) {
        if      (op[i] == MILP_LE) n_slack++;
        else if (op[i] == MILP_GE) { n_surp++; n_art++; }
        else                       { n_art++; }
    }
    int cols = n + n_slack + n_surp + n_art;

    LP *L = calloc(1, sizeof(*L));
    L->rows      = rows;
    L->cols      = cols;
    L->n_orig    = n;
    L->art_start = n + n_slack + n_surp;
    L->n_art     = n_art;
    L->T         = calloc((size_t)(rows + 1) * (cols + 1), sizeof(double));
    L->basis     = calloc(rows, sizeof(int));

    int slack_col = n;
    int surp_col  = n + n_slack;
    int art_col   = L->art_start;

    for (int i = 0; i < rows; i++) {
        for (int j = 0; j < n; j++)
            TIJ(L, i + 1, j) = A[(size_t)i * n + j];
        TIJ(L, i + 1, cols) = b[i];

        if (op[i] == MILP_LE) {
            TIJ(L, i + 1, slack_col) = 1.0;
            L->basis[i] = slack_col++;
        } else if (op[i] == MILP_GE) {
            TIJ(L, i + 1, surp_col++)  = -1.0;
            TIJ(L, i + 1, art_col)     =  1.0;
            L->basis[i] = art_col++;
        } else { /* EQ */
            TIJ(L, i + 1, art_col) = 1.0;
            L->basis[i] = art_col++;
        }
    }

    free(A); free(b); free(op);
    return L;
}

/* Bland's-rule simplex on the tableau in L.
 * forbid_art: if non-zero, artificial columns may not enter the basis.
 * Returns 0 = optimal, 1 = unbounded, 2 = iter limit. */
static int simplex_run(LP *L, int forbid_art) {
    for (int iter = 0; iter < MAX_PIVOTS; iter++) {
        /* Entering: smallest-index column with z_j < 0 */
        int enter = -1;
        for (int j = 0; j < L->cols; j++) {
            if (forbid_art && j >= L->art_start) continue;
            if (TIJ(L, 0, j) < -EPS_OPT) { enter = j; break; }
        }
        if (enter == -1) return 0;

        /* Leaving: min ratio with smallest basic-index tiebreak */
        int leave = -1;
        double best = DBL_MAX;
        for (int i = 1; i <= L->rows; i++) {
            double a = TIJ(L, i, enter);
            if (a > EPS_PIVOT) {
                double r = TIJ(L, i, L->cols) / a;
                if (r < best - EPS_OPT) {
                    best = r; leave = i;
                } else if (r < best + EPS_OPT && leave >= 1 &&
                           L->basis[i - 1] < L->basis[leave - 1]) {
                    leave = i;
                }
            }
        }
        if (leave == -1) return 1; /* unbounded */

        /* Pivot */
        double piv = TIJ(L, leave, enter);
        for (int j = 0; j <= L->cols; j++) TIJ(L, leave, j) /= piv;
        for (int i = 0; i <= L->rows; i++) {
            if (i == leave) continue;
            double f = TIJ(L, i, enter);
            if (fabs(f) > EPS_PIVOT) {
                for (int j = 0; j <= L->cols; j++)
                    TIJ(L, i, j) -= f * TIJ(L, leave, j);
            }
        }
        L->basis[leave - 1] = enter;
    }
    return 2;
}

/* Phase 1: minimise sum of artificials = maximise -sum(art). */
static int simplex_phase1(LP *L) {
    for (int j = 0; j <= L->cols; j++) TIJ(L, 0, j) = 0;
    /* z-row for max(-sum(art)) is +1 in artificial columns. */
    for (int j = L->art_start; j < L->art_start + L->n_art; j++)
        TIJ(L, 0, j) = 1.0;
    /* Eliminate basic variables (which include all artificials) from row 0 */
    for (int i = 1; i <= L->rows; i++) {
        int b = L->basis[i - 1];
        double zb = TIJ(L, 0, b);
        if (fabs(zb) > EPS_PIVOT)
            for (int j = 0; j <= L->cols; j++)
                TIJ(L, 0, j) -= zb * TIJ(L, i, j);
    }
    return simplex_run(L, 0);
}

/* After phase 1, some artificials may still be basic at rhs = 0 (degenerate).
 * Drive each one out by pivoting in any non-artificial column with a non-zero
 * coefficient in that row. If no such column exists, the row is redundant
 * (all zeros except in artificial columns); leave it alone -- with forbid_art
 * in phase 2 the artificial will never enter elsewhere and the row stays inert.
 *
 * Why this matters: a degenerate basic artificial at value 0 can be lifted to
 * a positive value during phase 2 if its row has negative coefficients in
 * subsequent entering columns. That silently destroys feasibility. */
static void drive_out_artificials(LP *L) {
    for (int i = 1; i <= L->rows; i++) {
        int b = L->basis[i - 1];
        if (b < L->art_start) continue;          /* not artificial */
        /* rhs must be ~0 here (phase 1 succeeded) */
        int pivot_col = -1;
        for (int j = 0; j < L->art_start; j++) { /* non-artificial cols only */
            if (fabs(TIJ(L, i, j)) > EPS_PIVOT) { pivot_col = j; break; }
        }
        if (pivot_col < 0) continue;             /* row redundant -- leave it */

        double piv = TIJ(L, i, pivot_col);
        for (int j = 0; j <= L->cols; j++) TIJ(L, i, j) /= piv;
        for (int k = 0; k <= L->rows; k++) {
            if (k == i) continue;
            double f = TIJ(L, k, pivot_col);
            if (fabs(f) > EPS_PIVOT) {
                for (int j = 0; j <= L->cols; j++)
                    TIJ(L, k, j) -= f * TIJ(L, i, j);
            }
        }
        L->basis[i - 1] = pivot_col;
    }
}

/* Phase 2: optimise the original objective, with artificials forbidden. */
static int simplex_phase2(const MilpModel *m, LP *L) {
    double sign = (m->sense == MILP_MAX) ? 1.0 : -1.0;
    for (int j = 0; j <= L->cols; j++) TIJ(L, 0, j) = 0;
    for (int j = 0; j < m->n; j++)
        TIJ(L, 0, j) = -sign * m->c[j];
    for (int i = 1; i <= L->rows; i++) {
        int b = L->basis[i - 1];
        double zb = TIJ(L, 0, b);
        if (fabs(zb) > EPS_PIVOT)
            for (int j = 0; j <= L->cols; j++)
                TIJ(L, 0, j) -= zb * TIJ(L, i, j);
    }
    return simplex_run(L, 1);
}

static void lp_extract(const MilpModel *m, LP *L,
                       double *x_out, double *obj_out) {
    int n = m->n;
    for (int j = 0; j < n; j++) x_out[j] = m->lo[j];      /* shift back */
    for (int i = 1; i <= L->rows; i++) {
        int b = L->basis[i - 1];
        if (b < n) x_out[b] += TIJ(L, i, L->cols);
    }
    double obj = 0;
    for (int j = 0; j < n; j++) obj += m->c[j] * x_out[j];
    *obj_out = obj;
}

static int solve_lp_internal(MilpModel *m, double *x_out, double *obj_out) {
    LP *L = lp_build(m);
    m->lp_calls++;

    int s1 = simplex_phase1(L);
    if (s1 != 0) { lp_free(L); return MILP_ERROR; }

    /* Feasibility check */
    double sum_art = 0;
    for (int i = 1; i <= L->rows; i++)
        if (L->basis[i - 1] >= L->art_start)
            sum_art += TIJ(L, i, L->cols);
    if (sum_art > EPS_FEAS) { lp_free(L); return MILP_INFEASIBLE; }

    drive_out_artificials(L);

    int s2 = simplex_phase2(m, L);
    if (s2 == 1) { lp_free(L); return MILP_UNBOUNDED; }
    if (s2 == 2) { lp_free(L); return MILP_ERROR; }

    lp_extract(m, L, x_out, obj_out);
    lp_free(L);
    return MILP_OPTIMAL;
}

int milp_solve_lp(MilpModel *m, double *x_out, double *obj_out) {
    m->nodes = 0; m->lp_calls = 0;
    int s = solve_lp_internal(m, x_out, obj_out);
    m->lp_relax = (s == MILP_OPTIMAL) ? *obj_out : 0;
    return s;
}

/* ---------------------------------------------------------------------
 * Branch-and-bound
 * ------------------------------------------------------------------- */

static int find_branch_var(const MilpModel *m, const double *x) {
    int branch = -1;
    double best_priority = -1.0;
    for (int j = 0; j < m->n; j++) {
        if (m->type[j] == MILP_CONTINUOUS) continue;
        double xj = x[j];
        double frac = xj - floor(xj);
        if (frac > EPS_INT && frac < 1 - EPS_INT) {
            /* Higher = closer to 0.5, the most informative branch. */
            double priority = 0.5 - fabs(frac - 0.5);
            if (priority > best_priority) {
                best_priority = priority;
                branch = j;
            }
        }
    }
    return branch;
}

static int strictly_better(int sense, double a, double b) {
    return (sense == MILP_MIN) ? (a < b - EPS_OPT) : (a > b + EPS_OPT);
}

static int cannot_improve(int sense, double bound, double incumbent) {
    return (sense == MILP_MIN) ? (bound >= incumbent - EPS_OPT)
                               : (bound <= incumbent + EPS_OPT);
}

static void bb_recurse(MilpModel *m, double *x_buf,
                       double *best_x, double *best_obj, int *found) {
    if (m->nodes >= m->node_limit) return;
    m->nodes++;

    double obj;
    int s = solve_lp_internal(m, x_buf, &obj);
    if (s != MILP_OPTIMAL) return;          /* infeasible / unbounded subproblem */

    if (m->verbose >= 2)
        printf("  node %4d: LP %.6f%s\n",
               m->nodes, obj, *found ? "" : " (no incumbent yet)");

    if (*found && cannot_improve(m->sense, obj, *best_obj)) return;

    int br = find_branch_var(m, x_buf);
    if (br == -1) {
        if (!*found || strictly_better(m->sense, obj, *best_obj)) {
            memcpy(best_x, x_buf, (size_t)m->n * sizeof(double));
            *best_obj = obj;
            *found = 1;
            if (m->verbose)
                printf("  incumbent: %.6f at node %d\n", obj, m->nodes);
        }
        return;
    }

    double xv     = x_buf[br];
    double down   = floor(xv);
    double up     = ceil(xv);
    double old_lo = m->lo[br];
    double old_hi = m->hi[br];

    if (down >= old_lo - EPS_OPT) {
        m->hi[br] = down;
        bb_recurse(m, x_buf, best_x, best_obj, found);
        m->hi[br] = old_hi;
    }
    if (up <= old_hi + EPS_OPT) {
        m->lo[br] = up;
        bb_recurse(m, x_buf, best_x, best_obj, found);
        m->lo[br] = old_lo;
    }
}

int milp_solve(MilpModel *m, double *x_out, double *obj_out) {
    m->nodes = 0;
    m->lp_calls = 0;

    int s = solve_lp_internal(m, x_out, obj_out);
    m->lp_relax = (s == MILP_OPTIMAL) ? *obj_out : 0;
    if (s != MILP_OPTIMAL) return s;

    if (m->verbose)
        printf("Root LP: %.6f (%s)\n",
               *obj_out, m->sense == MILP_MIN ? "min" : "max");

    if (find_branch_var(m, x_out) == -1) {
        m->nodes = 1;
        if (m->verbose) printf("Root LP integer feasible.\n");
        return MILP_OPTIMAL;
    }

    int found = 0;
    double best_obj = (m->sense == MILP_MIN) ? INFTY : -INFTY;
    double *best_x = calloc(m->n, sizeof(double));
    double *x_buf  = calloc(m->n, sizeof(double));

    bb_recurse(m, x_buf, best_x, &best_obj, &found);

    int ret;
    if (!found) {
        ret = MILP_INFEASIBLE;
    } else {
        memcpy(x_out, best_x, (size_t)m->n * sizeof(double));
        *obj_out = best_obj;
        ret = MILP_OPTIMAL;
    }
    free(best_x); free(x_buf);

    if (m->verbose && ret == MILP_OPTIMAL)
        printf("MIP optimal: %.6f, nodes=%d, LPs=%d, root gap=%.4f\n",
               *obj_out, m->nodes, m->lp_calls,
               fabs(m->lp_relax - *obj_out));
    return ret;
}
