#include "solver.hpp"
#include <iostream>
#include <cassert>
#include <algorithm>

using std::make_pair;
using std::pair;
using std::cout;
using std::endl;

Expr operator+(const Expr& x, const Expr& y)
{
	Expr exp;
	exp.coeffs.resize(std::max(x.coeffs.size(), y.coeffs.size()), 0);
	for (int c = 0; c < exp.coeffs.size(); c++) {
		if (c < x.coeffs.size() && c < y.coeffs.size()) {
			exp.coeffs[c] = x.coeffs[c] + y.coeffs[c];
		}
		else {
			exp.coeffs[c]= c < x.coeffs.size()? x.coeffs[c] : y.coeffs[c];
		}
	}
	exp.constant = x.constant + y.constant;
	return exp;
}

Expr operator+(const Expr& x)
{
	return x;
}

Expr operator+(Var x, Var y)
{
	Expr exp;
	exp.coeffs.resize(std::max(x.get(IntAttr::ModelColumn) + 1, y.get(IntAttr::ModelColumn) + 1), 0);
	exp.coeffs[x.get(IntAttr::ModelColumn)] = x.get(DoubleAttr::Coeff);
	exp.coeffs[y.get(IntAttr::ModelColumn)] = y.get(DoubleAttr::Coeff);
	return exp;
}

Expr operator+(Var x, double a)
{
	Expr exp;
	exp.coeffs.resize(x.get(IntAttr::ModelColumn) + 1);
	exp.constant = a;
	return exp;
}

Expr operator+(double a, Var x)
{
	return x + a;
}

Expr operator-(const Expr& x, const Expr& y)
{
	Expr exp;
	exp.coeffs.resize(std::max(x.coeffs.size(), y.coeffs.size()), 0);
	for (int c = 0; c < exp.coeffs.size(); c++) {
		if (c < x.coeffs.size() && c < y.coeffs.size()) {
			exp.coeffs[c] = x.coeffs[c] - y.coeffs[c];
		}
		else {
			exp.coeffs[c] = c < x.coeffs.size() ? x.coeffs[c] : y.coeffs[c];
		}
	}
	exp.constant = x.constant + y.constant;
	return exp;
}

Expr operator-(const Expr& x)
{
	Expr expr(x);
	for (int c = 0; c < expr.coeffs.size(); c++) {
		expr.coeffs[c] = -expr.coeffs[c];
	}
	expr.constant = -expr.constant;
	return expr;
}

Expr operator-(Var x)
{
	return -Expr(x);
}

Expr operator-(Var x, Var y)
{
	Expr exp;
	exp.coeffs.resize(std::max(x.get(IntAttr::ModelColumn) + 1, y.get(IntAttr::ModelColumn) + 1), 0);
	exp.coeffs[x.get(IntAttr::ModelColumn)] = x.get(DoubleAttr::Coeff);
	exp.coeffs[y.get(IntAttr::ModelColumn)] = -y.get(DoubleAttr::Coeff);
	return exp;
}

Expr operator-(Var x, double a)
{
	return x - Var(a);
}

Expr operator-(double a, Var x)
{
	return x - a;
}

Expr operator*(double a, Var x)
{
	Expr exp;
	exp.coeffs.resize(x.get(IntAttr::ModelColumn) + 1, 0);
	exp.coeffs[x.get(IntAttr::ModelColumn)] = a * x.get(DoubleAttr::Coeff);
	return exp;
}

Expr operator*(Var x, double a)
{
	return a * x;
}

Expr operator*(const Expr& x, double a)
{
	Expr exp = x;
	for (int c = 0; c < exp.coeffs.size(); c++) {
		exp.coeffs[c] *= a;
	}
	exp.constant *= a;
	return exp;
}

Expr operator*(double a, const Expr& x)
{
	return x * a;
}

Expr operator/(Var x, double a)
{
	return Expr(x) / a;
}

Expr operator/(const Expr& x, double a)
{
	Expr exp = x;
	for (int c = 0; c < exp.coeffs.size(); c++) {
		exp.coeffs[c] /= a;
	}
	exp.constant /= a;
	return exp;
}

Model::Model():
	cn(0),
	bn(0),
	sense(0),
	vars(nullptr)
{

}

Var* Model::addVars(int num)
{
	Var* old = vars;
	vars = new Var[cn + num];
	for (int c = cn; c < cn + num; c++) {
		vars[c].col = c;
	}
	memcpy(vars, old, sizeof(Var) * cn);
	delete[] old;
	cn += num;
	return vars;
}

void Model::addConstr(const Expr& expr, ConstrOper sense, double rhs)
{
	
	if (sense == ConstrOper::LESS_EQUAL) {
		matrix.push_back(vector<double>(1, rhs - expr.constant));
		matrix.back().insert(matrix.back().end(), expr.coeffs.begin(), expr.coeffs.end());
	}
	else if (sense == ConstrOper::GREATER_EQUAL) {
		matrix.push_back(vector<double>(1, expr.constant - rhs));
		for (int coeff : expr.coeffs) {
			matrix.back().push_back(-coeff);
		}
	}
	else {
		addConstr(expr, ConstrOper::LESS_EQUAL, rhs);
		addConstr(expr, ConstrOper::GREATER_EQUAL, rhs);
	}
	
	for (int c = matrix.back().size(); c <= cn; c++) {
		matrix.back().push_back(0);
	}
}

void Model::setObjective(Expr obje, int _sense)
{
	assert(_sense == 1 || _sense == -1);
	if (sense == 0) {
		matrix.insert(matrix.begin(), obje.coeffs);
		matrix.front().insert(matrix.front().begin(), -obje.constant);
		for (int c = obje.coeffs.size() + 1; c <= cn; c++) {
			matrix.front().push_back(0);
		}
	}
	else {
		matrix.front().front() = -obje.constant;
		for (int c = 0; c < obje.coeffs.size(); c++) {
			matrix.front()[c + 1] = obje.coeffs[c];
		}
		for (int c = obje.coeffs.size() + 1; c <= cn; c++) {
			matrix.front()[c] = 0;
		}

	}
	sense = _sense;
}

void Model::print()
{
	for (size_t i = 0; i < matrix.size(); i++) {
		for (size_t j = 0; j < matrix[0].size(); j++) {
			cout << matrix[i][j] << "\t";
		}
		cout << endl;
	}
}

Rtn Model::optimize()
{
	bn = matrix.size();
	for (int row = 1; row < bn; row++) {
		matrix.front().push_back(0);
		for (int col = 1; col < bn; col++) {
			matrix[row].push_back(col == row ? 1 : 0);
		}
	}
	cn = matrix.front().size();
	for (int i = 0; i < cn; i++) {
		matrix.front()[i] = sense * matrix.front()[i];
	}
	
	for (size_t i = 1; i < bn; i++) {
		basic.push_back(cn - bn + i);
	}

	// === 判断初始解是否为可行解 ===
	bool initial_feasible = true;
	for (int row = 1; row < bn; row++) {
		if (matrix[row].front() < 0) {
			initial_feasible = false;
			break;
		}
	}

	// === 构造初始可行解 === 
	if (!initial_feasible) {
		vector<double> coeff = matrix.front();
		matrix.front() = vector<double>(cn, .0);
		matrix.front().push_back(1);
		pair<size_t, size_t> t = { -1 ,cn };

		for (int row = 1; row < bn; row++) {
			matrix[row].push_back(-1);
			if (t.first == -1 || matrix[row].front() < matrix[t.first].front()) {
				t.first = row;
			}
		}

		_gaussian(t);
		if (fabs(_simplex()) > 1e-10) {
			rtn = Rtn::INFEASIBLE;
			return rtn;
		}

		// if the x0 in B, we should pivot it.
		auto iter = find(basic.begin(), basic.end(), cn);
		if (iter != basic.end()) {
			for (int col = 1; col < matrix.front().size(); col++) {
				if (fabs(matrix.front()[col]) > 1e-10) {
					t = make_pair(iter - basic.begin() + 1, col);
					_gaussian(t);
					break;
				}
			}
		}

		for (int row = 0; row < bn; row++) {
			matrix[row].pop_back();
		}

		// recover the coefficient line
		for (int col = 0; col < cn; col++) {
			matrix.front()[col] = coeff[col];
		}
		for (int row = 1; row <= basic.size(); row++) {
			int norm = matrix.front()[basic[row - 1]];
			for (int col = 0; col < cn; col++) {
				matrix.front()[col] -= norm * matrix[row][col];
			}
		}
	}

	res = -sense * _simplex();
	for (int row = 1; row < bn; row++) {
		if (basic[row - 1] - 1 < cn - bn) {
			vars[basic[row - 1] - 1].val = matrix[row].front();
		}
	}
	return rtn;
}

double Model::get(DoubleAttr attr)
{
	return res;
}

int Model::get(IntAttr attr)
{
	return cn - bn;
}

double Model::_simplex()
{
	pair<size_t, size_t> t;
	while (1) {
		rtn = _pivot(t);
		if (rtn == Rtn::OPTIMAL || rtn == Rtn::UNBOUNDED) {
			break;
		}
		_gaussian(t);
	}
	return matrix.front().front();
}

Rtn Model::_pivot(pair<size_t, size_t>& p)
{
	p = make_pair(0, 0);
	double cmin = INT_MAX;
	vector<double> coef = matrix.front();

	// === 非主轴元素中找最小值 ===
	for (size_t i = 1; i < coef.size(); i++) {
		if (cmin > coef[i] && find(basic.begin(), basic.end(), i) == basic.end()) {
			cmin = coef[i];
			p.second = i;
		}
	}
	if (cmin >= 0) {
		return Rtn::OPTIMAL;
	}
	double bmin = INT_MAX;
	for (size_t row = 1; row < bn; row++) {
		double tmp = matrix[row].front() / matrix[row][p.second];
		if (matrix[row][p.second] > 0 && bmin > tmp) {
			bmin = tmp;
			p.first = row;
		}
	}

	if (abs(bmin - INT_MAX) < 1e-10) {
		return Rtn::UNBOUNDED;
	}

	for (auto iter = basic.begin(); iter != basic.end(); iter++) {
		if (matrix[p.first][*iter] != 0) {
			*iter = p.second;
			break;
		}
	}
	assert(basic[p.first - 1] == p.second);
	return Rtn::PIVOT;
}

void Model::_gaussian(pair<size_t, size_t> p)
{
	size_t x = p.first, y = p.second;

	// === 主行归一化 ===
	double norm = matrix[x][y];
	for (size_t col = 0; col < matrix[x].size(); col++) {
		matrix[x][col] /= norm;
	}

	// === 其余行变换 ===
	for (size_t row = 0; row < bn; row++) {
		if (row == x) {
			continue;
		}
		if (matrix[row][y] != 0) {
			double norm = matrix[row][y];
			for (size_t col = 0; col < matrix[x].size(); col++) {
				matrix[row][col] = matrix[row][col] - norm * matrix[x][col];
			}
		}
	}
	basic[x - 1] = y;	// 换元
}


Expr::Expr(double constant)
	:constant(constant)
{
}

Expr::Expr(Var var, double coeff)
{
	this->coeffs.resize(var.col + 1);
	this->coeffs[var.col] = coeff;
	this->constant = 0;
}

Expr Expr::operator=(const Expr& rhs)
{
	return *this;
}

void Expr::operator+=(const Expr& expr)
{
	coeffs.resize(std::max(coeffs.size(), expr.coeffs.size()));
	for (int c = 0; c < expr.coeffs.size(); c++) {
		coeffs[c] += expr.coeffs[c];
	}
	constant += expr.constant;
}

void Expr::operator-=(const Expr& expr)
{
	coeffs.resize(std::max(coeffs.size(), expr.coeffs.size()));
	for (int c = 0; c < expr.coeffs.size(); c++) {
		coeffs[c] -= expr.coeffs[c];
	}
	constant -= expr.constant;
}

void Expr::operator*=(double mult)
{
	for (int c = 0; c < coeffs.size(); c++) {
		coeffs[c] *= mult;
	}
	constant *= mult;
}

void Expr::operator/=(double a)
{
	for (int c = 0; c < coeffs.size(); c++) {
		coeffs[c] /= a;
	}
	constant /= a;
}

Expr Expr::operator+(const Expr& rhs)
{
	coeffs.resize(std::max(coeffs.size(), rhs.coeffs.size()));
	for (int c = 0; c < rhs.coeffs.size(); c++) {
		coeffs[c] += rhs.coeffs[c];
	}
	constant += rhs.constant;
	return *this;
}

Expr Expr::operator-(const Expr& rhs)
{
	coeffs.resize(std::max(coeffs.size(), rhs.coeffs.size()));
	for (int c = 0; c < rhs.coeffs.size(); c++) {
		coeffs[c] -= rhs.coeffs[c];
	}
	constant -= rhs.constant;
	return *this;
}

double Var::get(DoubleAttr attr)
{
	switch (attr) {
	case DoubleAttr::Coeff:
		return coeffs;
	case DoubleAttr::X:
		return val;
	}
	return val;
}

int Var::get(IntAttr attr)
{
	return col;
}