// Copyright (C) 2026 Kiyotsugu Arai
// SPDX-License-Identifier: LGPL-3.0-or-later

// linear_programming.hpp
// Linear programming (simplex method)
//
// Solving linear programming problems via the two-phase revised simplex method
// Reference: based on simplex.c from Okumura "Algorithm Dictionary in C",
//            fully reworked into a type-safe template implementation
//
// Improvements:
//   - Global variables → encapsulated in a struct
//   - Fixed-size arrays → dynamic Matrix<T>/Vector<T>
//   - scanf → problem received as function arguments
//   - Error handling: exit() → state returned via a result struct
//   - Full support for equality (=) and inequality (<=, >=) constraints

#ifndef SANGI_LINEAR_PROGRAMMING_HPP
#define SANGI_LINEAR_PROGRAMMING_HPP

#include "optimization_base.hpp"
#include <math/core/vector.hpp>
#include <math/core/matrix.hpp>

#include <vector>
#include <cmath>
#include <algorithm>
#include <limits>
#include <cassert>
#include <stdexcept>

namespace sangi {

/// Kind of constraint inequality
enum class ConstraintType {
    LessEqual,    // <=
    GreaterEqual, // >=
    Equal         // =
};

/// A single constraint row of a linear programming problem
template<typename T>
struct LinearConstraint {
    std::vector<T> coefficients;  // Left-hand-side coefficients (same length as the number of variables)
    ConstraintType type;          // Kind of inequality
    T rhs;                        // Right-hand-side constant

    LinearConstraint() : type(ConstraintType::LessEqual), rhs(T(0)) {}

    LinearConstraint(std::vector<T> coeffs, ConstraintType t, T r)
        : coefficients(std::move(coeffs)), type(t), rhs(r) {}
};

/// Definition of a linear programming problem
/// minimize c^T x   subject to  A_i x {<=, >=, =} b_i
template<typename T>
struct LinearProgram {
    std::vector<T> objective;                      // Objective function coefficients (minimization)
    T objectiveConstant = T(0);                    // Constant term of the objective function
    std::vector<LinearConstraint<T>> constraints;  // List of constraints

    size_t numVariables() const { return objective.size(); }
    size_t numConstraints() const { return constraints.size(); }
};

/// State of the LP solver
enum class LPStatus {
    Optimal,       // An optimal solution was found
    Infeasible,    // No feasible solution
    Unbounded,     // The objective function is unbounded
    MaxIterations  // Maximum number of iterations reached
};

/// Result of the LP solver
template<typename T>
struct LPResult {
    LPStatus status;
    T objectiveValue;              // Optimal objective function value
    std::vector<T> solution;       // Optimal solution (value of each variable)
    size_t iterations;             // Total number of pivots

    static LPResult optimal(T val, std::vector<T> sol, size_t iters) {
        return {LPStatus::Optimal, val, std::move(sol), iters};
    }
    static LPResult infeasible(size_t iters) {
        return {LPStatus::Infeasible, T(0), {}, iters};
    }
    static LPResult unbounded(size_t iters) {
        return {LPStatus::Unbounded, -std::numeric_limits<T>::infinity(), {}, iters};
    }
    static LPResult maxIterations(T val, std::vector<T> sol, size_t iters) {
        return {LPStatus::MaxIterations, val, std::move(sol), iters};
    }
};

namespace detail {

/// Internal implementation of the two-phase revised simplex method
template<typename T>
class SimplexSolver {
public:
    SimplexSolver(const LinearProgram<T>& lp, T eps, size_t maxIter)
        : eps_(eps), maxIter_(maxIter), totalPivots_(0)
    {
        setup(lp);
    }

    LPResult<T> solve() {
        // Phase 1: if artificial variables exist, find a feasible basic solution
        if (n3_ > n2_) {
            auto status = phase1();
            if (status == LPStatus::Infeasible) {
                return LPResult<T>::infeasible(totalPivots_);
            }
        }

        // Phase 2: optimize the objective function
        auto status = phase2();
        if (status == LPStatus::Unbounded) {
            return LPResult<T>::unbounded(totalPivots_);
        }
        if (status == LPStatus::MaxIterations) {
            return LPResult<T>::maxIterations(
                extractObjectiveValue(), extractSolution(), totalPivots_);
        }

        return LPResult<T>::optimal(
            extractObjectiveValue(), extractSolution(), totalPivots_);
    }

private:
    T eps_;
    size_t maxIter_;
    size_t totalPivots_;

    size_t m_;    // Number of constraints
    size_t n_;    // Number of original variables
    size_t n1_;   // n + number of >= slack variables
    size_t n2_;   // n1 + number of <= slack variables
    size_t n3_;   // n2 + number of artificial variables
    size_t jmax_; // Rightmost column index of the tableau

    // Constraint coefficient matrix a[0..m][0..n] — row 0 is the objective function
    std::vector<std::vector<T>> a_;

    // Transformation matrix q[0..m][0..m]
    std::vector<std::vector<T>> q_;

    // Temporary buffer for the pivot column
    std::vector<T> pivotCol_;

    // col[i] = column index of the basic variable of row i (i=1..m)
    std::vector<int> col_;

    // row[j] = corresponding row index if column j is a basic variable, 0 if non-basic
    std::vector<int> row_;

    // nonzeroRow[j] = row index where slack/artificial variable j is nonzero
    std::vector<int> nonzeroRow_;

    // Inequality kind of each constraint
    std::vector<ConstraintType> ineq_;

    // Objective function coefficients c[0..n] (c[0]=-constant term, c[1..n]=each variable's coefficient)
    // Copied into a_[0] at the start of phase 2 (same structure as the original C code)
    std::vector<T> c_;

    // Constant term of the objective function
    T objConst_;

    /// Set up the problem in internal form
    void setup(const LinearProgram<T>& lp) {
        m_ = lp.numConstraints();
        n_ = lp.numVariables();
        objConst_ = lp.objectiveConstant;

        // Preprocessing to make the right-hand side non-negative
        ineq_.resize(m_ + 1);
        a_.assign(m_ + 1, std::vector<T>(n_ + 1, T(0)));

        // Save the objective function coefficients (a_[0] is set in phase 2)
        c_.assign(n_ + 1, T(0));
        for (size_t j = 0; j < n_; ++j) {
            c_[j + 1] = lp.objective[j];
        }
        c_[0] = -objConst_;

        // Store the constraints (flip the inequality if the right-hand side is negative)
        for (size_t i = 0; i < m_; ++i) {
            const auto& c = lp.constraints[i];
            assert(c.coefficients.size() == n_);

            ineq_[i + 1] = c.type;
            a_[i + 1][0] = c.rhs;

            for (size_t j = 0; j < n_; ++j) {
                a_[i + 1][j + 1] = c.coefficients[j];
            }

            // Normalize the right-hand side to be non-negative
            if (a_[i + 1][0] < T(0)) {
                flipInequality(i + 1);
                for (size_t j = 0; j <= n_; ++j) {
                    a_[i + 1][j] = -a_[i + 1][j];
                }
            } else if (a_[i + 1][0] == T(0) && ineq_[i + 1] == ConstraintType::GreaterEqual) {
                // 0 >= ... is flipped to <= (flip only the coefficients, the right-hand side stays 0)
                ineq_[i + 1] = ConstraintType::LessEqual;
                for (size_t j = 1; j <= n_; ++j) {
                    a_[i + 1][j] = -a_[i + 1][j];
                }
            }
        }

        // Prepare slack and artificial variables
        prepare();
    }

    /// Flip the inequality
    void flipInequality(size_t i) {
        if (ineq_[i] == ConstraintType::LessEqual) {
            ineq_[i] = ConstraintType::GreaterEqual;
        } else if (ineq_[i] == ConstraintType::GreaterEqual) {
            ineq_[i] = ConstraintType::LessEqual;
        }
        // Equal is left unchanged
    }

    /// Introduce slack and artificial variables
    void prepare() {
        size_t totalCols = n_ + 2 * m_ + 1;

        col_.assign(m_ + 1, 0);
        row_.assign(totalCols, 0);
        nonzeroRow_.assign(totalCols, 0);

        // Add a -1 slack variable to each >= constraint
        n1_ = n_;
        for (size_t i = 1; i <= m_; ++i) {
            if (ineq_[i] == ConstraintType::GreaterEqual) {
                ++n1_;
                nonzeroRow_[n1_] = static_cast<int>(i);
            }
        }

        // Add a +1 slack variable to each <= constraint (used as the initial basic variable)
        n2_ = n1_;
        for (size_t i = 1; i <= m_; ++i) {
            if (ineq_[i] == ConstraintType::LessEqual) {
                ++n2_;
                col_[i] = static_cast<int>(n2_);
                nonzeroRow_[n2_] = static_cast<int>(i);
                row_[n2_] = static_cast<int>(i);
            }
        }

        // Add an artificial variable to each >= and = constraint
        n3_ = n2_;
        for (size_t i = 1; i <= m_; ++i) {
            if (ineq_[i] != ConstraintType::LessEqual) {
                ++n3_;
                col_[i] = static_cast<int>(n3_);
                nonzeroRow_[n3_] = static_cast<int>(i);
                row_[n3_] = static_cast<int>(i);
            }
        }

        // Initialize the transformation matrix Q to the identity matrix
        q_.assign(m_ + 1, std::vector<T>(m_ + 1, T(0)));
        for (size_t i = 0; i <= m_; ++i) {
            q_[i][i] = T(1);
        }

        pivotCol_.resize(m_ + 1);
    }

    /// Compute the tableau element (i, j)
    /// Revised simplex method: computed dynamically as Q*A without keeping the full tableau
    T tableau(size_t i, size_t j) const {
        if (col_[i] < 0) return T(0); // Deleted row

        if (j <= n_) {
            // Original variable column: dot product of row i of Q and column j of A
            T s = T(0);
            for (size_t k = 0; k <= m_; ++k) {
                s += q_[i][k] * a_[k][j];
            }
            return s;
        }

        // Slack/artificial variable column
        T s = q_[i][nonzeroRow_[j]];
        if (j <= n1_) return -s;       // >= slack variable (coefficient -1)
        if (j <= n2_ || i != 0) return s;  // <= slack variable (coefficient +1)
        return s + T(1);               // Artificial variable (+1 correction on the objective function row)
    }

    /// Pivot operation: replace the basic variable of row ip with the variable of column jp
    /// Note: column 0 of the Q matrix is not touched (protects the objective function's constant term)
    void pivot(size_t ip, size_t jp) {
        T u = pivotCol_[ip];

        // Normalize the pivot row (columns 1..m only)
        for (size_t j = 1; j <= m_; ++j) {
            q_[ip][j] /= u;
        }

        // Remove the pivot row's contribution from the other rows (columns 1..m only)
        for (size_t i = 0; i <= m_; ++i) {
            if (i != ip) {
                u = pivotCol_[i];
                for (size_t j = 1; j <= m_; ++j) {
                    q_[i][j] -= q_[ip][j] * u;
                }
            }
        }

        // Update the basic variable indices
        row_[col_[ip]] = 0;
        col_[ip] = static_cast<int>(jp);
        row_[jp] = static_cast<int>(ip);
    }

    /// Minimization loop: pick a non-basic variable with negative reduced cost and pivot
    LPStatus minimize() {
        for (;;) {
            if (totalPivots_ >= maxIter_) {
                return LPStatus::MaxIterations;
            }

            // Pivot column selection: the non-basic variable with the most negative reduced cost
            size_t jp = 0;
            T bestCost = -eps_;
            for (size_t j = 1; j <= jmax_; ++j) {
                if (row_[j] == 0) {
                    T cost = tableau(0, j);
                    if (cost < bestCost) {
                        bestCost = cost;
                        jp = j;
                    }
                }
            }
            if (jp == 0) break; // All reduced costs non-negative → optimal

            // Compute the pivot column and select the pivot row (minimum ratio test)
            for (size_t i = 0; i <= m_; ++i) {
                pivotCol_[i] = tableau(i, jp);
            }

            size_t ip = 0;
            T minRatio = std::numeric_limits<T>::max();
            for (size_t i = 1; i <= m_; ++i) {
                if (pivotCol_[i] > eps_) {
                    T ratio = tableau(i, 0) / pivotCol_[i];
                    if (ratio < minRatio) {
                        minRatio = ratio;
                        ip = i;
                    }
                }
            }

            if (ip == 0) {
                // All elements of the pivot column are non-positive → the objective function is unbounded
                return LPStatus::Unbounded;
            }

            pivot(ip, jp);
            ++totalPivots_;
        }

        return LPStatus::Optimal;
    }

    /// Phase 1: drive the artificial variables out of the basis and find a feasible basic solution
    LPStatus phase1() {
        jmax_ = n3_;

        // Set up an objective function that minimizes the sum of the artificial variables
        // Modify row 0 of Q: subtract 1 from each row that has an artificial variable
        for (size_t i = 0; i <= m_; ++i) {
            if (col_[i] > static_cast<int>(n2_)) {
                q_[0][i] -= T(1);
            }
        }

        auto status = minimize();
        if (status == LPStatus::Unbounded) {
            return LPStatus::Infeasible;
        }

        // If the artificial objective value is positive, the problem is infeasible
        T artObj = tableau(0, 0);
        if (artObj < -eps_) {
            return LPStatus::Infeasible;
        }

        // Delete the rows of artificial variables remaining in the basis as redundant
        for (size_t i = 1; i <= m_; ++i) {
            if (col_[i] > static_cast<int>(n2_)) {
                col_[i] = -1; // Mark the row as deleted
            }
        }

        // Restore the original objective function
        q_[0].assign(m_ + 1, T(0));
        q_[0][0] = T(1);

        // Cancel out the reduced costs of original variables that are in the basis
        for (size_t i = 1; i <= m_; ++i) {
            int j = col_[i];
            if (j > 0 && j <= static_cast<int>(n_)) {
                T u = c_[j]; // Coefficient of column j of the objective function
                if (u != T(0)) {
                    for (size_t k = 1; k <= m_; ++k) {
                        q_[0][k] -= q_[i][k] * u;
                    }
                }
            }
        }

        return LPStatus::Optimal;
    }

    /// Phase 2: minimize the original objective function
    LPStatus phase2() {
        jmax_ = n2_; // Exclude the artificial variable columns
        // Copy the objective function coefficients into a_[0] (same as phase2() in the original C code)
        for (size_t j = 0; j <= n_; ++j) {
            a_[0][j] = c_[j];
        }
        return minimize();
    }

    /// Obtain the optimal objective function value
    T extractObjectiveValue() const {
        return -tableau(0, 0) + objConst_;
    }

    /// Obtain the variable values of the optimal solution
    std::vector<T> extractSolution() const {
        std::vector<T> sol(n_, T(0));
        for (size_t j = 1; j <= n_; ++j) {
            int i = row_[j];
            if (i != 0) {
                sol[j - 1] = tableau(i, 0);
            }
        }
        return sol;
    }
};

} // namespace detail

/// Solve a linear programming problem with the simplex method
///
/// @tparam T Numeric type (double, float, etc.)
/// @param lp  Definition of the linear programming problem (minimization problem)
/// @param eps Tolerance (threshold for pivot selection)
/// @param maxIterations Maximum number of pivots
/// @return Result (status, objective function value, solution vector, number of iterations)
///
/// Example usage:
/// @code
///   // minimize -x - 2y  subject to  x + y <= 4,  x <= 3,  y <= 3
///   LinearProgram<double> lp;
///   lp.objective = {-1.0, -2.0};
///   lp.constraints = {
///       {{1.0, 1.0}, ConstraintType::LessEqual, 4.0},
///       {{1.0, 0.0}, ConstraintType::LessEqual, 3.0},
///       {{0.0, 1.0}, ConstraintType::LessEqual, 3.0}
///   };
///   auto result = simplex(lp);
///   // result.status == LPStatus::Optimal
///   // result.objectiveValue == -7.0
///   // result.solution == {1.0, 3.0}
/// @endcode
template<typename T>
LPResult<T> simplex(
    const LinearProgram<T>& lp,
    T eps = T(1e-10),
    size_t maxIterations = 10000);

/// Wrapper for maximization problems: maximize c^T x → minimize -c^T x
template<typename T>
LPResult<T> simplexMaximize(
    const LinearProgram<T>& lp,
    T eps = T(1e-10),
    size_t maxIterations = 10000);

// The implementation is separated into linear_programming_impl.hpp.

} // namespace sangi

#endif // SANGI_LINEAR_PROGRAMMING_HPP