#include "Scheme.h"

#include <NTL/RR.h>
#include <NTL/ZZ.h>
#include <NTL/ZZX.h>
#include <cmath>

#include "EvaluatorUtils.h"
#include "NumUtils.h"
#include "Params.h"
#include "PubKey.h"
#include "Ring2Utils.h"

using namespace std;
using namespace NTL;

CZZ* Scheme::groupidx(CZZ*& vals, long slots) {
	CZZ* res = new CZZ[slots * 2];
	long logslots = log2(slots);
	for (long i = 0; i < slots; ++i) {
		long idx = (params.rotGroup[i] % (4 * slots) - 1) / 2;
		res[idx] = vals[i];
		res[2 * slots - idx - 1] = vals[i].conjugate();
	}
	return res;
}

CZZ* Scheme::groupidx(CZZ& val) {
	CZZ* res = new CZZ[2];
	res[0] = val;
	res[1] = val.conjugate();
	return res;
}

CZZ* Scheme::degroupidx(CZZ*& vals, long slots) {
	long logslots = log2(slots);
	CZZ* res = new CZZ[slots];
	for (long i = 0; i < slots; ++i) {
		long idx = (params.rotGroup[i] % (4 * slots) - 1) / 2;
		res[i] = vals[idx];
	}
	return res;
}

//-----------------------------------------

Message Scheme::encodeWithBits(CZZ*& gvals, long cbits, long slots) {
	ZZX mx;
	mx.SetLength(params.N);
	ZZ mod = power2_ZZ(cbits);
	long idx = 0;
	long doubleslots = slots << 1;
	long logDoubleslots = log2(slots) + 1;
	long gap = (params.N >> logDoubleslots);
	NumUtils::fftSpecialInv(gvals, doubleslots, aux);
	for (long i = 0; i < doubleslots; ++i) {
		mx.rep[idx] = gvals[i].r;
		idx += gap;
	}
	return Message(mx, mod, cbits, slots);
}

Message Scheme::encode(CZZ*& gvals, long slots) {
	ZZX mx;
	mx.SetLength(params.N);
	long idx = 0;
	long doubleslots = slots << 1;
	long logDoubleslots = log2(slots) + 1;
	long gap = (params.N >> logDoubleslots);

	NumUtils::fftSpecialInv(gvals, doubleslots, aux);

	for (long i = 0; i < doubleslots; ++i) {
		mx.rep[idx] = gvals[i].r;
		idx += gap;
	}
	return Message(mx, params.q, params.logq, slots);
}

Cipher Scheme::encryptMsg(Message& msg) {
	ZZX ax, bx, vx, eax, ebx;
	NumUtils::sampleZO(vx, params.N);
	Ring2Utils::mult(ax, vx, publicKey.ax, msg.mod, params.N);
	NumUtils::sampleGauss(eax, params.N, params.sigma);
	Ring2Utils::addAndEqual(ax, eax, msg.mod, params.N);

	Ring2Utils::mult(bx, vx, publicKey.bx, msg.mod, params.N);
	NumUtils::sampleGauss(ebx, params.N, params.sigma);
	Ring2Utils::addAndEqual(bx, ebx, msg.mod, params.N);

	Ring2Utils::addAndEqual(bx, msg.mx, msg.mod, params.N);

	return Cipher(ax, bx, msg.mod, msg.cbits, msg.slots);
}

Cipher Scheme::encryptWithBits(CZZ*& vals, long cbits, long slots) {
	CZZ* gvals = groupidx(vals, slots);
	Message msg = encodeWithBits(gvals, cbits, slots);
	delete[] gvals;
	return encryptMsg(msg);
}

Cipher Scheme::encrypt(CZZ*& vals, long slots) {
	CZZ* gvals = groupidx(vals, slots);
	Message msg = encode(gvals, slots);
	delete[] gvals;
	return encryptMsg(msg);
}

Cipher Scheme::encryptSingleWithBits(CZZ& val, long cbits) {
	CZZ* gvals = groupidx(val);
	Message msg = encodeWithBits(gvals, cbits, 1);
	delete[] gvals;
	return encryptMsg(msg);
}

Cipher Scheme::encryptSingle(CZZ& val) {
	CZZ* gvals = groupidx(val);
	Message msg = encode(gvals, 1);
	delete[] gvals;
	return encryptMsg(msg);
}

//-----------------------------------------

Message Scheme::decryptMsg(SecKey& secretKey, Cipher& cipher) {
	ZZX mx;
	Ring2Utils::mult(mx, cipher.ax, secretKey.sx, cipher.mod, params.N);
	Ring2Utils::addAndEqual(mx, cipher.bx, cipher.mod, params.N);
	return Message(mx, cipher.mod, cipher.cbits, cipher.slots);
}

CZZ* Scheme::decode(Message& msg) {
	long doubleslots = msg.slots * 2;
	CZZ* fftinv = new CZZ[doubleslots];

	long idx = 0;
	long gap = params.N / doubleslots;
	for (long i = 0; i < doubleslots; ++i) {
		ZZ tmp = msg.mx.rep[idx] % msg.mod;
		if(NumBits(tmp) == msg.cbits) tmp -= msg.mod;
		fftinv[i] = CZZ(tmp, ZZ(0));
		idx += gap;
	}
	NumUtils::fftSpecial(fftinv, doubleslots, aux);
	return fftinv;
}

CZZ* Scheme::decrypt(SecKey& secretKey, Cipher& cipher) {
	Message msg = decryptMsg(secretKey, cipher);
	CZZ* gvals = decode(msg);
	CZZ* res = degroupidx(gvals, msg.slots);
	delete[] gvals;
	return res;
}

CZZ Scheme::decryptSingle(SecKey& secretKey, Cipher& cipher) {
	Message msg = decryptMsg(secretKey, cipher);
	CZZ* gvals = decode(msg);
	CZZ res = gvals[0];
	delete[] gvals;
	return res;
}

//-----------------------------------------

Cipher Scheme::negate(Cipher& cipher){
    ZZX ax, bx;

    Ring2Utils::multByConst(ax, cipher.ax, to_ZZ("-1"), cipher.mod, params.N);
    Ring2Utils::multByConst(bx, cipher.bx, to_ZZ("-1"), cipher.mod, params.N);

    return Cipher(ax, bx, cipher.mod, cipher.cbits, cipher.slots);
}

void Scheme::negateAndEqual(Cipher& cipher){
    Ring2Utils::multByConstAndEqual(cipher.ax, to_ZZ("-1"), cipher.mod, params.N);
    Ring2Utils::multByConstAndEqual(cipher.bx, to_ZZ("-1"), cipher.mod, params.N);
}

//-----------------------------------------

Cipher Scheme::add(Cipher& cipher1, Cipher& cipher2) {
	ZZX ax, bx;

	Ring2Utils::add(ax, cipher1.ax, cipher2.ax, cipher1.mod, params.N);
	Ring2Utils::add(bx, cipher1.bx, cipher2.bx, cipher1.mod, params.N);

	return Cipher(ax, bx, cipher1.mod, cipher1.cbits, cipher1.slots);
}

void Scheme::addAndEqual(Cipher& cipher1, Cipher& cipher2) {
	Ring2Utils::addAndEqual(cipher1.ax, cipher2.ax, cipher1.mod, params.N);
	Ring2Utils::addAndEqual(cipher1.bx, cipher2.bx, cipher1.mod, params.N);
}

//-----------------------------------------

Cipher Scheme::addConst(Cipher& cipher, ZZ& cnst) {
	ZZX ax = cipher.ax;
	ZZX bx = cipher.bx;

	AddMod(bx.rep[0], cipher.bx.rep[0], cnst, cipher.mod);
	return Cipher(ax, bx, cipher.mod, cipher.cbits, cipher.slots);
}

void Scheme::addConstAndEqual(Cipher& cipher, ZZ& cnst) {
	ZZ mod = power2_ZZ(cipher.cbits);
	AddMod(cipher.bx.rep[0], cipher.bx.rep[0], cnst, mod);
}

//-----------------------------------------

Cipher Scheme::addByPoly(Cipher& cipher, Message& msg) {
    ZZX axres, bxres;

    //Ring2Utils::add(axres, cipher.ax, msg.mx, cipher.mod, params.N);
    Ring2Utils::add(bxres, cipher.bx, msg.mx, cipher.mod, params.N);
    return Cipher(cipher.ax, bxres, cipher.mod, cipher.cbits, cipher.slots);
}

void Scheme::addByPolyAndEqual(Cipher& cipher, Message& msg) {
    //Ring2Utils::addAndEqual(cipher.ax, msg.mx, cipher.mod, params.N);
    Ring2Utils::addAndEqual(cipher.bx, msg.mx, cipher.mod, params.N);
}

//-----------------------------------------

Cipher Scheme::sub(Cipher& cipher1, Cipher& cipher2) {
	ZZX ax, bx;

	Ring2Utils::sub(ax, cipher1.ax, cipher2.ax, cipher1.mod, params.N);
	Ring2Utils::sub(bx, cipher1.bx, cipher2.bx, cipher1.mod, params.N);

	return Cipher(ax, bx, cipher1.mod, cipher1.cbits, cipher1.slots);
}

void Scheme::subAndEqual(Cipher& cipher1, Cipher& cipher2) {
	Ring2Utils::subAndEqual(cipher1.ax, cipher2.ax, cipher1.mod, params.N);
	Ring2Utils::subAndEqual(cipher1.bx, cipher2.bx, cipher1.mod, params.N);
}

void Scheme::subAndEqual2(Cipher& cipher1, Cipher& cipher2) {
	Ring2Utils::subAndEqual2(cipher1.ax, cipher2.ax, cipher1.mod, params.N);
	Ring2Utils::subAndEqual2(cipher1.bx, cipher2.bx, cipher1.mod, params.N);
}

//-----------------------------------------

// cipher - msg  = (b, a) - (mx, 0)
Cipher Scheme::subByPoly(Cipher& cipher, Message& msg) {
    ZZX bxres;
    //Ring2Utils::sub(axres, cipher.ax, msg.mx, cipher.mod, params.N);
    Ring2Utils::sub(bxres, cipher.bx, msg.mx, cipher.mod, params.N);
    return Cipher(cipher.ax, bxres, cipher.mod, cipher.cbits, cipher.slots);
}

void Scheme::subByPolyAndEqual(Cipher& cipher, Message& msg) {

    Ring2Utils::subAndEqual(cipher.bx, msg.mx, cipher.mod, params.N);
}

// msg - cipher  = (mx, 0 ) - (b, a)
Cipher Scheme::subByPoly2(Cipher& cipher, Message& msg) {
    ZZX axres, bxres;

    Ring2Utils::multByConst(axres, cipher.ax, to_ZZ("-1"), cipher.mod, params.N);
    Ring2Utils::sub(bxres, msg.mx, cipher.bx, cipher.mod, params.N);
    return Cipher(axres, bxres, cipher.mod, cipher.cbits, cipher.slots);
}

//-----------------------------------------

Cipher Scheme::subConst(Cipher& cipher, ZZ& cnst){
    ZZX ax = cipher.ax;
    ZZX bx = cipher.bx;

    SubMod(bx.rep[0], cipher.bx.rep[0], cnst, cipher.mod);
    return Cipher(ax, bx, cipher.mod, cipher.cbits, cipher.slots);
}

Cipher Scheme::subConst2(Cipher& cipher, ZZ& cnst){
    Cipher res = negate(cipher);
    addConstAndEqual(res, cnst);

    return res;
}

void Scheme::subConstAndEqual(Cipher& cipher, ZZ& cnst){
    //ZZX ax = cipher.ax;
    //ZZX bx = cipher.bx;

    SubMod(cipher.bx.rep[0], cipher.bx.rep[0], cnst, cipher.mod);
    //return Cipher(ax, bx, cipher.mod, cipher.cbits, cipher.slots);
}
//-----------------------------------------

Cipher Scheme::conjugate(Cipher& cipher) {
	ZZ Pmod = cipher.mod << params.logq;

	ZZX bxconj, bxres, axres;

	Ring2Utils::conjugate(bxconj, cipher.bx, params.N);
	Ring2Utils::conjugate(bxres, cipher.ax, params.N);

	Ring2Utils::mult(axres, bxres, publicKey.axConj, Pmod, params.N);
	Ring2Utils::multAndEqual(bxres, publicKey.bxConj, Pmod, params.N);

	Ring2Utils::rightShiftAndEqual(axres, params.logq, params.N);
	Ring2Utils::rightShiftAndEqual(bxres, params.logq, params.N);

	Ring2Utils::addAndEqual(bxres, bxconj, cipher.mod, params.N);
	return Cipher(axres, bxres, cipher.mod, cipher.cbits, cipher.slots);
}

void Scheme::conjugateAndEqual(Cipher& cipher) {
	ZZ Pmod = cipher.mod << params.logq;

	ZZX bxconj, bxres, axres;

	Ring2Utils::conjugate(bxconj, cipher.bx, params.N);
	Ring2Utils::conjugate(bxres, cipher.ax, params.N);

	Ring2Utils::mult(axres, bxres, publicKey.axConj, Pmod, params.N);
	Ring2Utils::multAndEqual(bxres, publicKey.bxConj, Pmod, params.N);

	Ring2Utils::rightShiftAndEqual(axres, params.logq, params.N);
	Ring2Utils::rightShiftAndEqual(bxres, params.logq, params.N);

	Ring2Utils::addAndEqual(bxres, bxconj, cipher.mod, params.N);

	cipher.ax = axres;
	cipher.bx = bxres;
}

Cipher Scheme::imult(Cipher& cipher, const long precisionBits) {
	ZZ tmp = EvaluatorUtils::evaluateVal(sqrt(to_RR(2.0)), precisionBits - 1);

	ZZX bxres, axres, axtmp, bxtmp;

	Ring2Utils::multByMonomial(axtmp, cipher.ax, params.N / 4, params.N);
	Ring2Utils::multByConstAndEqual(axtmp, tmp, cipher.mod, params.N);

	Ring2Utils::multByMonomial(bxtmp, cipher.bx, params.N / 4, params.N);
	Ring2Utils::multByConstAndEqual(bxtmp, tmp, cipher.mod, params.N);

	Ring2Utils::multByMonomial(axres, cipher.ax, 3 * params.N / 4, params.N);
	Ring2Utils::multByConstAndEqual(axres, tmp, cipher.mod, params.N);

	Ring2Utils::multByMonomial(bxres, cipher.bx, 3 * params.N / 4, params.N);
	Ring2Utils::multByConstAndEqual(bxres, tmp, cipher.mod, params.N);

	Ring2Utils::addAndEqual(axres, axtmp, cipher.mod, params.N);
	Ring2Utils::addAndEqual(bxres, bxtmp, cipher.mod, params.N);

	Cipher res(axres, bxres, cipher.mod, cipher.cbits, cipher.slots);
	modSwitchAndEqual(res, precisionBits);
	return res;
}

void Scheme::imultAndEqual(Cipher& cipher, const long precisionBits) {

	ZZ tmp = EvaluatorUtils::evaluateVal(sqrt(to_RR(2.0)), precisionBits - 1);

	ZZX axtmp, bxtmp;

	Ring2Utils::multByMonomial(axtmp, cipher.ax, params.N / 4, params.N);
	Ring2Utils::multByConstAndEqual(axtmp, tmp, cipher.mod, params.N);

	Ring2Utils::multByMonomial(bxtmp, cipher.bx, params.N / 4, params.N);
	Ring2Utils::multByConstAndEqual(bxtmp, tmp, cipher.mod, params.N);

	Ring2Utils::multByMonomialAndEqual(cipher.ax, 3 * params.N / 4, params.N);
	Ring2Utils::multByConstAndEqual(cipher.ax, tmp, cipher.mod, params.N);

	Ring2Utils::multByMonomialAndEqual(cipher.bx, 3 * params.N / 4, params.N);
	Ring2Utils::multByConstAndEqual(cipher.bx, tmp, cipher.mod, params.N);

	Ring2Utils::addAndEqual(cipher.ax, axtmp, cipher.mod, params.N);
	Ring2Utils::addAndEqual(cipher.bx, bxtmp, cipher.mod, params.N);

	modSwitchAndEqual(cipher, precisionBits);
}

Cipher Scheme::mult(Cipher& cipher1, Cipher& cipher2) {
	ZZ Pmod = cipher1.mod << params.logq;

	ZZX axbx1 = Ring2Utils::add(cipher1.ax, cipher1.bx, cipher1.mod, params.N);
	ZZX axbx2 = Ring2Utils::add(cipher2.ax, cipher2.bx, cipher1.mod, params.N);
	Ring2Utils::multAndEqual(axbx1, axbx2, cipher1.mod, params.N);

	ZZX bxbx = Ring2Utils::mult(cipher1.bx, cipher2.bx, cipher1.mod, params.N);
	ZZX axax = Ring2Utils::mult(cipher1.ax, cipher2.ax, cipher1.mod, params.N);

	ZZX axmult = Ring2Utils::mult(axax, publicKey.axStar, Pmod, params.N);
	ZZX bxmult = Ring2Utils::mult(axax, publicKey.bxStar, Pmod, params.N);

	Ring2Utils::rightShiftAndEqual(axmult, params.logq, params.N);
	Ring2Utils::rightShiftAndEqual(bxmult, params.logq, params.N);

	Ring2Utils::addAndEqual(axmult, axbx1, cipher1.mod, params.N);
	Ring2Utils::subAndEqual(axmult, bxbx, cipher1.mod, params.N);
	Ring2Utils::subAndEqual(axmult, axax, cipher1.mod, params.N);
	Ring2Utils::addAndEqual(bxmult, bxbx, cipher1.mod, params.N);

	return Cipher(axmult, bxmult, cipher1.mod, cipher1.cbits, cipher1.slots);
}

void Scheme::multAndEqual(Cipher& cipher1, Cipher& cipher2) {
	ZZ Pmod = cipher1.mod << params.logq;

	ZZX axbx1 = Ring2Utils::add(cipher1.ax, cipher1.bx, cipher1.mod, params.N);
	ZZX axbx2 = Ring2Utils::add(cipher2.ax, cipher2.bx, cipher1.mod, params.N);
	Ring2Utils::multAndEqual(axbx1, axbx2, cipher1.mod, params.N);

	ZZX bxbx = Ring2Utils::mult(cipher1.bx, cipher2.bx, cipher1.mod, params.N);
	ZZX axax = Ring2Utils::mult(cipher1.ax, cipher2.ax, cipher1.mod, params.N);

	cipher1.ax = Ring2Utils::mult(axax, publicKey.axStar, Pmod, params.N);
	cipher1.bx = Ring2Utils::mult(axax, publicKey.bxStar, Pmod, params.N);

	Ring2Utils::rightShiftAndEqual(cipher1.ax, params.logq, params.N);
	Ring2Utils::rightShiftAndEqual(cipher1.bx, params.logq, params.N);

	Ring2Utils::addAndEqual(cipher1.ax, axbx1, cipher1.mod, params.N);
	Ring2Utils::subAndEqual(cipher1.ax, bxbx, cipher1.mod, params.N);
	Ring2Utils::subAndEqual(cipher1.ax, axax, cipher1.mod, params.N);
	Ring2Utils::addAndEqual(cipher1.bx, bxbx, cipher1.mod, params.N);
}

//-----------------------------------------

Cipher Scheme::square(Cipher& cipher) {
	ZZ Pmod = cipher.mod << params.logq;

	ZZX axax, axbx, bxbx, bxmult, axmult;

	Ring2Utils::square(bxbx, cipher.bx, cipher.mod, params.N);
	Ring2Utils::mult(axbx, cipher.ax, cipher.bx, cipher.mod, params.N);
	Ring2Utils::addAndEqual(axbx, axbx, cipher.mod, params.N);
	Ring2Utils::square(axax, cipher.ax, cipher.mod, params.N);

	Ring2Utils::mult(axmult, axax, publicKey.axStar, Pmod, params.N);
	Ring2Utils::mult(bxmult, axax, publicKey.bxStar, Pmod, params.N);

	Ring2Utils::rightShiftAndEqual(axmult, params.logq, params.N);
	Ring2Utils::rightShiftAndEqual(bxmult, params.logq, params.N);

	Ring2Utils::addAndEqual(axmult, axbx, cipher.mod, params.N);
	Ring2Utils::addAndEqual(bxmult, bxbx, cipher.mod, params.N);

	return Cipher(axmult, bxmult, cipher.mod, cipher.cbits, cipher.slots);
}

void Scheme::squareAndEqual(Cipher& cipher) {
	ZZ Pmod = cipher.mod << params.logq;

	ZZX bxbx, axbx, axax, bxmult, axmult;

	Ring2Utils::square(bxbx, cipher.bx, cipher.mod, params.N);
	Ring2Utils::mult(axbx, cipher.bx, cipher.ax, cipher.mod, params.N);
	Ring2Utils::addAndEqual(axbx, axbx, cipher.mod, params.N);
	Ring2Utils::square(axax, cipher.ax, cipher.mod, params.N);

	Ring2Utils::mult(axmult, axax, publicKey.axStar, Pmod, params.N);
	Ring2Utils::mult(bxmult, axax, publicKey.bxStar, Pmod, params.N);

	Ring2Utils::rightShiftAndEqual(axmult, params.logq, params.N);
	Ring2Utils::rightShiftAndEqual(bxmult, params.logq, params.N);

	Ring2Utils::addAndEqual(axmult, axbx, cipher.mod, params.N);
	Ring2Utils::addAndEqual(bxmult, bxbx, cipher.mod, params.N);

	cipher.bx = bxmult;
	cipher.ax = axmult;
}

//-----------------------------------------

Cipher Scheme::multByConst(Cipher& cipher, ZZ& cnst) {
	ZZX ax, bx;
	Ring2Utils::multByConst(ax, cipher.ax, cnst, cipher.mod, params.N);
	Ring2Utils::multByConst(bx, cipher.bx, cnst, cipher.mod, params.N);

	return Cipher(ax, bx, cipher.mod, cipher.cbits, cipher.slots);
}

void Scheme::multByConstAndEqual(Cipher& cipher, ZZ& cnst) {
	Ring2Utils::multByConstAndEqual(cipher.ax, cnst, cipher.mod, params.N);
	Ring2Utils::multByConstAndEqual(cipher.bx, cnst, cipher.mod, params.N);
}

Cipher Scheme::multByPoly(Cipher& cipher, ZZX& poly) {
	ZZX axres, bxres;
	Ring2Utils::mult(axres, cipher.ax, poly, cipher.mod, params.N);
	Ring2Utils::mult(bxres, cipher.bx, poly, cipher.mod, params.N);
	return Cipher(axres, bxres, cipher.mod, cipher.cbits, cipher.slots);
}

void Scheme::multByPolyAndEqual(Cipher& cipher, ZZX& poly) {
	Ring2Utils::multAndEqual(cipher.ax, poly, cipher.mod, params.N);
	Ring2Utils::multAndEqual(cipher.bx, poly, cipher.mod, params.N);
}

Cipher Scheme::multByConstBySlots(Cipher& cipher, CZZ*& cnstvec) {
	CZZ* gcnstvec = groupidx(cnstvec, cipher.slots);
	Message msg = encode(gcnstvec, cipher.slots);
	delete[] gcnstvec;

	ZZX axres, bxres;
	Ring2Utils::mult(axres, cipher.ax, msg.mx, cipher.mod, params.N);
	Ring2Utils::mult(bxres, cipher.bx, msg.mx, cipher.mod, params.N);
	return Cipher(axres, bxres, cipher.mod, cipher.cbits, cipher.slots);
}

void Scheme::multByConstBySlotsAndEqual(Cipher& cipher, CZZ*& cnstvec) {
	CZZ* gcnstvec = groupidx(cnstvec, cipher.slots);
	Message msg = encode(gcnstvec, cipher.slots);
	delete[] gcnstvec;

	Ring2Utils::multAndEqual(cipher.ax, msg.mx, cipher.mod, params.N);
	Ring2Utils::multAndEqual(cipher.bx, msg.mx, cipher.mod, params.N);
}

Cipher Scheme::multByPoly(Cipher& cipher, Message& msg) {
    ZZX axres, bxres;
    Ring2Utils::mult(axres, cipher.ax, msg.mx, cipher.mod, params.N);
    Ring2Utils::mult(bxres, cipher.bx, msg.mx, cipher.mod, params.N);
    return Cipher(axres, bxres, cipher.mod, cipher.cbits, cipher.slots);
}

void Scheme::multByPolyAndEqual(Cipher& cipher, Message& msg) {
    Ring2Utils::multAndEqual(cipher.ax, msg.mx, cipher.mod, params.N);
    Ring2Utils::multAndEqual(cipher.bx, msg.mx, cipher.mod, params.N);
}

//-----------------------------------------

Cipher Scheme::multByMonomial(Cipher& cipher, const long degree) {
	ZZX ax, bx;

	Ring2Utils::multByMonomial(ax, cipher.ax, degree, params.N);
	Ring2Utils::multByMonomial(bx, cipher.bx, degree, params.N);

	return Cipher(ax, bx, cipher.mod, cipher.cbits, cipher.slots);
}

void Scheme::multByMonomialAndEqual(Cipher& cipher, const long degree) {
	Ring2Utils::multByMonomialAndEqual(cipher.ax, degree, params.N);
	Ring2Utils::multByMonomialAndEqual(cipher.bx, degree, params.N);
}

//-----------------------------------------

Cipher Scheme::leftShift(Cipher& cipher, long bits) {
	ZZX ax, bx;

	Ring2Utils::leftShift(ax, cipher.ax, bits, cipher.mod, params.N);
	Ring2Utils::leftShift(bx, cipher.bx, bits, cipher.mod, params.N);

	return Cipher(ax, bx, cipher.mod, cipher.cbits, cipher.slots);
}

void Scheme::leftShiftAndEqual(Cipher& cipher, long bits) {
	Ring2Utils::leftShiftAndEqual(cipher.ax, bits, cipher.mod, params.N);
	Ring2Utils::leftShiftAndEqual(cipher.bx, bits, cipher.mod, params.N);
}

void Scheme::doubleAndEqual(Cipher& cipher) {
	Ring2Utils::doubleAndEqual(cipher.ax, cipher.mod, params.N);
	Ring2Utils::doubleAndEqual(cipher.bx, cipher.mod, params.N);
}

//-----------------------------------------

Cipher Scheme::modSwitch(Cipher& cipher, long bitsDown) {
	ZZX ax, bx;

	Ring2Utils::rightShift(ax, cipher.ax, bitsDown, params.N);
	Ring2Utils::rightShift(bx, cipher.bx, bitsDown, params.N);

	long newcbits = cipher.cbits - bitsDown;
	ZZ newmod = cipher.mod >> bitsDown;
	return Cipher(ax, bx, newmod, newcbits, cipher.slots);
}

void Scheme::modSwitchAndEqual(Cipher& cipher, long bitsDown) {
	Ring2Utils::rightShiftAndEqual(cipher.ax, bitsDown, params.N);
	Ring2Utils::rightShiftAndEqual(cipher.bx, bitsDown, params.N);
	cipher.cbits -= bitsDown;
	cipher.mod >>= bitsDown;
}

Cipher Scheme::modEmbed(Cipher& cipher, long bitsDown) {
	ZZ newmod = cipher.mod >> bitsDown;
	long newcbits = cipher.cbits - bitsDown;
	ZZX bx, ax;
	Ring2Utils::mod(ax, cipher.ax, newmod, params.N);
	Ring2Utils::mod(bx, cipher.bx, newmod, params.N);
	return Cipher(ax, bx, newmod, newcbits, cipher.slots);
}

void Scheme::modEmbedAndEqual(Cipher& cipher, long bitsDown) {
	cipher.mod >>= bitsDown;
	cipher.cbits -= bitsDown;
	Ring2Utils::modAndEqual(cipher.ax, cipher.mod, params.N);
	Ring2Utils::modAndEqual(cipher.bx, cipher.mod, params.N);
}

//-----------------------------------------

Cipher Scheme::leftRotateByPo2(Cipher& cipher, long logrotSlots) {
	ZZ Pmod = cipher.mod << params.logq;

	ZZX bxrot, bxres, axres;

	long rotSlots = (1 << logrotSlots);

	Ring2Utils::inpower(bxrot, cipher.bx, params.rotGroup[rotSlots], params.q, params.N);
	Ring2Utils::inpower(bxres, cipher.ax, params.rotGroup[rotSlots], params.q, params.N);

	Ring2Utils::mult(axres, bxres, publicKey.axLeftRot[logrotSlots], Pmod, params.N);
	Ring2Utils::multAndEqual(bxres, publicKey.bxLeftRot[logrotSlots], Pmod, params.N);

	Ring2Utils::rightShiftAndEqual(axres, params.logq, params.N);
	Ring2Utils::rightShiftAndEqual(bxres, params.logq, params.N);

	Ring2Utils::addAndEqual(bxres, bxrot, cipher.mod, params.N);
	return Cipher(axres, bxres, cipher.mod, cipher.cbits, cipher.slots);
}

void Scheme::leftRotateByPo2AndEqual(Cipher& cipher, long logrotSlots) {
	ZZ Pmod = cipher.mod << params.logq;

	ZZX bxrot, bxres, axres;

	long rotSlots = (1 << logrotSlots);

	Ring2Utils::inpower(bxrot, cipher.bx, params.rotGroup[rotSlots], params.q, params.N);
	Ring2Utils::inpower(bxres, cipher.ax, params.rotGroup[rotSlots], params.q, params.N);

	Ring2Utils::mult(axres, bxres, publicKey.axLeftRot[logrotSlots], Pmod, params.N);
	Ring2Utils::multAndEqual(bxres, publicKey.bxLeftRot[logrotSlots], Pmod, params.N);

	Ring2Utils::rightShiftAndEqual(axres, params.logq, params.N);
	Ring2Utils::rightShiftAndEqual(bxres, params.logq, params.N);

	Ring2Utils::addAndEqual(bxres, bxrot, cipher.mod, params.N);

	cipher.ax = axres;
	cipher.bx = bxres;
}

Cipher Scheme::rightRotateByPo2(Cipher& cipher, long logrotSlots) {
	ZZ Pmod = cipher.mod << params.logq;

	ZZX bxrot, bxres, axres;

	long rotSlots = params.N/2 - (1 << logrotSlots);

	Ring2Utils::inpower(bxrot, cipher.bx, params.rotGroup[rotSlots], params.q, params.N);
	Ring2Utils::inpower(bxres, cipher.ax, params.rotGroup[rotSlots], params.q, params.N);

	Ring2Utils::mult(axres, bxres, publicKey.axRightRot[logrotSlots], Pmod, params.N);
	Ring2Utils::multAndEqual(bxres, publicKey.bxRightRot[logrotSlots], Pmod, params.N);

	Ring2Utils::rightShiftAndEqual(axres, params.logq, params.N);
	Ring2Utils::rightShiftAndEqual(bxres, params.logq, params.N);

	Ring2Utils::addAndEqual(bxres, bxrot, cipher.mod, params.N);
	return Cipher(axres, bxres, cipher.mod, cipher.cbits, cipher.slots);
}

void Scheme::rightRotateByPo2AndEqual(Cipher& cipher, long logrotSlots) {
	ZZ Pmod = cipher.mod << params.logq;

	ZZX bxrot, bxres, axres;

	long rotSlots = params.N/2 - (1 << logrotSlots);

	Ring2Utils::inpower(bxrot, cipher.bx, params.rotGroup[rotSlots], params.q, params.N);
	Ring2Utils::inpower(bxres, cipher.ax, params.rotGroup[rotSlots], params.q, params.N);

	Ring2Utils::mult(axres, bxres, publicKey.axRightRot[logrotSlots], Pmod, params.N);
	Ring2Utils::multAndEqual(bxres, publicKey.bxRightRot[logrotSlots], Pmod, params.N);

	Ring2Utils::rightShiftAndEqual(axres, params.logq, params.N);
	Ring2Utils::rightShiftAndEqual(bxres, params.logq, params.N);

	Ring2Utils::addAndEqual(bxres, bxrot, cipher.mod, params.N);

	cipher.ax = axres;
	cipher.bx = bxres;
}

Cipher Scheme::leftRotate(Cipher& cipher, long rotSlots) {
	Cipher res = cipher;
	leftRotateAndEqual(res, rotSlots);
	return res;
}

void Scheme::leftRotateAndEqual(Cipher& cipher, long rotSlots) {
	long remrotSlots = rotSlots % cipher.slots;
	long logrotSlots = log2(remrotSlots) + 1;
	for (long i = 0; i < logrotSlots; ++i) {
		if(bit(remrotSlots, i)) {
			leftRotateByPo2AndEqual(cipher, i);
		}
	}
}

Cipher Scheme::rightRotate(Cipher& cipher, long rotSlots) {
	Cipher res = cipher;
	rightRotateAndEqual(res, rotSlots);
	return res;
}

void Scheme::rightRotateAndEqual(Cipher& cipher, long rotSlots) {
	long remrotSlots = rotSlots % cipher.slots;
	long logrotSlots = log2(remrotSlots) + 1;
	for (long i = 0; i < logrotSlots; ++i) {
		if(bit(remrotSlots, i)) {
			rightRotateByPo2AndEqual(cipher, i);
		}
	}
}


#include <iostream>
#include <thread>
#include <fstream>
#include <string>
#include <filesystem>
#include <tuple>
#include <algorithm>
#include <chrono>
#include <regex>

#include <boost/multiprecision/cpp_int.hpp>

using namespace std;

namespace mp = boost::multiprecision;
namespace fs = std::filesystem;

const auto iter = 1 + 28;
const string resultsDirectory = "results";

void saveResult(const mp::cpp_int &power, int num) {
    vector<unsigned char> buffer;
    export_bits(power, back_inserter(buffer), 8);

    fstream file(resultsDirectory + "/" + to_string(num) + ".txt", fstream::out | fstream::trunc);
    for (unsigned char val: buffer) {
        file.put(val);
    }
    file.close();
}

void saveResult(const mp::cpp_int &power, const string &filename) {
    vector<unsigned char> buffer;
    export_bits(power, back_inserter(buffer), 8);

    fstream file(resultsDirectory + "/" + filename + ".txt", fstream::out | fstream::trunc);
    for (unsigned char val: buffer) {
        file.put(val);
    }
    file.close();
}

tuple <mp::cpp_int*, int> loadResults() {
    auto *powers = new mp::cpp_int[iter];
    mp::cpp_int partial;
    vector<int> loadedPowers;
    auto numberToLarge = false;

    for (const auto &entry : fs::directory_iterator(resultsDirectory)) {
        auto filename = fs::path(entry).stem(); // Get filename without path and extension
        int power = atoi(filename.c_str()); // Convert string into integer

        if (power == 0) {
            continue;
        }

        if (power < iter) {
            cout << "Loading " + entry.path().filename().string() << endl;

            vector<unsigned char> buffer;
            fstream file(entry.path(), fstream::in);

            while (!file.eof()) {
                char val;
                file.get(val);
                if (file.eof()) break;
                buffer.push_back(val);
            }
            file.close();

            import_bits(powers[power], buffer.begin(), buffer.end());

            if (powers[power] == 0) {
                continue;
            }
            loadedPowers.push_back(power); // Add read power to loadedPowers
        } else {
            numberToLarge = true;
        }
    }

    cout << "Finished Loading" << endl;
    if (numberToLarge) {
        cout << "Found file with power larger than set possible!" << endl << endl;
    }

    int largestPower = 0;
    sort(loadedPowers.begin(), loadedPowers.end());

    // Find largest power without "gaps" in files, so
    // 0, 1, 2, 3, 5, 6 will return 3
    // We will need all powers in stage two
    for (auto power: loadedPowers) {
        if (power == largestPower + 1){
            largestPower++;
        }
    }

    powers[0] = 9;
    largestPower = max(++largestPower, 1); // Set largestPower to bigger value: 1 or largestPower + 1

    return {powers, largestPower};
}

string prepareTime(chrono::steady_clock::time_point start, chrono::steady_clock::time_point end) {
    auto microseconds = chrono::duration_cast<chrono::microseconds>(end - start).count();
    auto milliseconds = chrono::duration_cast<chrono::milliseconds>(end - start).count();
    auto seconds = chrono::duration_cast<chrono::seconds>(end - start).count();
    auto hours = chrono::duration_cast<chrono::hours>(end - start).count();

    string time;

    if (hours > 0) {
        time += to_string(hours) + "h ";
    }
    if (seconds > 0) {
        time += to_string(seconds % 60) + "s ";
    }
    if (milliseconds > 0) {
        time += to_string(milliseconds % 1000) + "ms ";
    }
    if (microseconds > 0) {
        time += to_string(microseconds % 1000) + "μm";
    }

    if (time.empty()) {
        return "no time";
    } else {
        return time;
    }
}

int main() {
    auto [powers, currentPower] = loadResults();
    thread threads[iter];

    for (; currentPower < iter; ++currentPower) {
        auto startTime = chrono::steady_clock::now();

        powers[currentPower] = pow(powers[currentPower - 1], 2);

        auto endTime = chrono::steady_clock::now();
        cout << "Calculated 9^" << mp::pow(mp::cpp_int(2), currentPower) << ", in " << prepareTime(startTime, endTime) << ", step " << currentPower << endl; // For logging progress

        // thread() can't find rsquare prismight function by its own (there are two),
        // so I point it to right address manually
        auto functionAddress = static_cast<void(*)(const mp::cpp_int&, int)>(saveResult);
        threads[currentPower] = thread(functionAddress, powers[currentPower], currentPower);
    }

    // Stage two: multiply following elements from array:
    // int toMultiply[] =  {28, 26, 25, 24, 20, 18, 17, 16, 15, 12, 8, 6, 3, 0};
    int toMultiply[] = {3, 6, 8, 12, 15, 16, 17, 18, 20, 24, 25, 26, 28};
    // Based on:
    // https://www.wolframalpha.com/input/?i=9%5E9+%3D+2%5E28+%2B+2%5E26+%2B+2%5E25+%2B+2%5E24+%2B+2%5E20+%2B+2%5E18+%2B+2%5E17+%2B+2%5E16+%2B+2%5E15+%2B+2%5E12+%2B+2%5E8+%2B+2%5E6+%2B+2%5E3+%2B+2%5E0

    mp::cpp_int final = powers[0];
    int previous = 0;
    for (int i = 0; i < 13; i++) {
        cout << "Calculating: " << toMultiply[i] << endl;
        final *= powers[toMultiply[i]];

        cout << "Saving..." << endl;
        saveResult(final, "partial." + to_string(toMultiply[i]));
        remove(("partial." + to_string(previous)).c_str());
    }

    for (auto &th: threads) {
        if (th.joinable()) {
            th.join();
        }
    }

    fstream result("result.txt", fstream::out);
    result << final;
    result.close();
    return 0;
}


#ifndef DIIS_HH
#define DIIS_HH

/*!
 * \file DIIS.hh
 * \brief Definition of the DIIS class
 */

#include <Eigen/Core>

/*!
 * \brief The DIIS convergence optimizer
 *
 * Class DIIS implements the Direct Inversion in the Iterative Subspace
 * extrapolation technique to speed up the convergence in an SCF calculation.
 * It tries to make a better guess for the Fock matrix in the next iteration
 * by creating a linear combination of Fock matrices from previous iterations,
 * where the error term \f$FDS-SDF\f$ (with \f$F\f$ the Fock matrix, \f$D\f$
 * the density matrix, and \f$S\f$ the overlap matrix) is minimized.
 */
class DIIS
{
public:
	/*!
	 * \brief Constructor
	 *
	 * Create a new DIIS object, for a basis of size \a size.
	 * \param size The size of the matrices (basis set size).
	 */
	DIIS(): _size(0), _err_vecs_used(0), _err_vecs(), _values(),
		_started(false), _max_err(0) {}

	/*!
	 * \brief Create a new Fock matrix
	 *
	 * Create the new Fock matrix from the current guess for the matrix and
	 * the current density matrix. If the DIIS algorithm has not been
	 * started yet (because the error term is still too large with respect
	 * to the computed energy), it simply returns \a F. Otherwise a linear
	 * combination of previous Fock matrices \f$\tilde{F}\f$ that minimizes
	 * the error term \f$\tilde{F}DS-SD\tilde{F}\f$ is computed, and stored
	 * in \a F.
	 * \param F    The current Fock matrix
	 * \param P    The density matrix from which \a F was computed
	 * \param S    The overlap matrix for the basis
	 * \param X    The orthogonalization matrix \f$X = S^{-1/2}\f$ for the basis
	 * \param Etot The total energy
	 */
	void step(Eigen::MatrixXd& F, const Eigen::MatrixXd& D,
		const Eigen::MatrixXd& S, const Eigen::MatrixXd& X, double Etot);
	void step(Eigen::MatrixXd& Fa, const Eigen::MatrixXd& Da,
		Eigen::MatrixXd& Fb, const Eigen::MatrixXd& Db,
		const Eigen::MatrixXd& S, const Eigen::MatrixXd& X, double Etot);

	//! Return the maximum error on the last iteration
	double error() const { return _max_err; }

private:
	//! The basis set size
	int _size;
	//! The number of error vectors currently stored
	int _err_vecs_used;
	//! The error vectors themselves, one per column
	Eigen::MatrixXd _err_vecs;
	//! The previous Fock matrices
	std::vector<Eigen::MatrixXd> _values;
	//! Whether DIIS was started
	bool _started;
	//! Maximum error in the last DIIS step
	double _max_err;
};

#endif // DIIS_HH


//=========================================================================
//
// Copyright 2019 Kitware, Inc.
// Author: Guilbert Pierre (spguilbert@gmail.com)
// Data: 03-27-2019
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
//=========================================================================

// LOCAL
#include "CameraCalibration.h"
#include "CameraProjection.h"
#include "vtkEigenTools.h"
#include "CeresCameraCalibrationCostFunctions.h"

// STD
#include <iostream>
#include <fstream>
#include <sstream>

// BOOST
#include <boost/algorithm/string.hpp>

// CERES
#include <ceres/ceres.h>

//----------------------------------------------------------------------------
void LoadMatchesFromCSV(std::string filename, std::vector<Eigen::Vector3d>& X, std::vector<Eigen::Vector2d>& x)
{
  // Load file and check that the file is opened
  std::ifstream file(filename.c_str());
  if (!file.is_open())
  {
    std::cout << "Error: could not load file: " << filename << std::endl;
    return;
  }

  // check the file header
  std::string tokenFileHeader("X,Y,Z,u,v");
  std::string line;
  std::getline(file, line);
  if (line != tokenFileHeader)
  {
    std::cout << "Error file header is: " << line << " expected: " << tokenFileHeader << std::endl;
    return;
  }

  // parse the file
  while (std::getline(file, line))
  {
    std::vector<std::string> values;
    boost::algorithm::split(values, line, boost::is_any_of(","));

    // check that each lines has the right number of value
    if (values.size() < 5)
    {
      std::cout << "Error, number of values is: " << values.size() << " expected 5" << std::endl;
      continue;
    }

    Eigen::Vector3d Xc(std::atof(values[0].c_str()), std::atof(values[1].c_str()), std::atof(values[2].c_str()));
    Eigen::Vector2d xc(std::atof(values[3].c_str()), std::atof(values[4].c_str()));
    X.push_back(Xc);
    x.push_back(xc);
  }
  return;
}

//----------------------------------------------------------------------------
double LinearPinholeCalibration(const std::vector<Eigen::Vector3d>& X, const std::vector<Eigen::Vector2d>& x, Eigen::Matrix<double, 3, 4>& P)
{
  P = Eigen::Matrix<double, 3, 4>::Zero();
  if (X.size() != x.size())
  {
    std::cout << "Error: matches have different sizes" << std::endl;
    return -1.0;
  }

  // We will estimate the coefficients of the projection matrix
  // by minimizing the reprojection euclidean distance between
  // the image keypoints and the 3D associated keypoints reprojected

  // Compute the normal equations
  Eigen::MatrixXd A(2 * X.size(), 12);
  for (int k = 0; k < X.size(); ++k)
  {
    A.row(2 * k) << X[k](0), X[k](1), X[k](2), 1,
                    0, 0, 0, 0,
                    -x[k](0) * X[k](0), -x[k](0) * X[k](1), -x[k](0) * X[k](2), -x[k](0);

    A.row(2 * k + 1) << 0, 0, 0, 0,
                        X[k](0), X[k](1), X[k](2), 1,
                       -x[k](1) * X[k](0), -x[k](1) * X[k](1), -x[k](1) * X[k](2), -x[k](1);
  }

  // Solve the normal equations
  Eigen::JacobiSVD<Eigen::MatrixXd> svd(A, Eigen::ComputeThinU | Eigen::ComputeThinV);
  Eigen::MatrixXd flattenP = svd.matrixV().col(11);
  P << flattenP(0), flattenP(1), flattenP(2), flattenP(3),
       flattenP(4), flattenP(5), flattenP(6), flattenP(7),
       flattenP(8), flattenP(9), flattenP(10), flattenP(11);
  P = P / P(2, 3);

  // Compute the RMSE
  double meanErr = 0;
  for (int k = 0; k < X.size(); ++k)
  {
    Eigen::Vector4d homoX(X[k](0), X[k](1), X[k](2), 1);
    Eigen::Vector3d projX = P * homoX;
    Eigen::Vector2d projXn(projX(0) / projX(2), projX(1) / projX(2));
    //std::cout << "X: " << X[k].transpose() << " x: " << x[k].transpose() << " prjX: " << projXn.transpose() << " dist: " << (x[k] - projXn).norm() << std::endl;
    meanErr += ((x[k] - projXn).transpose() * (x[k] - projXn))(0);
  }
  return std::sqrt(meanErr / (1.0 * X.size()));
}

//----------------------------------------------------------------------------
void CalibrationMatrixDecomposition(const Eigen::Matrix<double, 3, 4>& P, Eigen::Matrix3d& K, Eigen::Matrix3d& R, Eigen::Vector3d& T)
{
  // M = K*R
  Eigen::Matrix3d M = P.block(0, 0, 3, 3);

  // M is the product of an upper triangulate matrix
  // and an orthogonal matrix. We will use a RQ decomposition
  // to recover K and R. Since eigen only provides QR
  // decomposition we will inverse the column order of M to
  // get its RQ decomposition from the QR decomposition
  Eigen::Matrix3d D = Eigen::Matrix3d::Zero();
  D(2, 0) = 1; D(1, 1) = 1; D(0, 2) = 1;
  Eigen::Matrix3d Mtilde = D * M;
  Eigen::HouseholderQR<Eigen::Matrix3d> QRdec(Mtilde.transpose());
  Eigen::Matrix3d Qtilde = QRdec.householderQ();
  Eigen::Matrix3d Rtilde = Qtilde.transpose() * (Mtilde.transpose());
  K = D * Rtilde.transpose() * D;
  R = D * Qtilde.transpose();
  T = P.col(3);
  T = - R.transpose() * K.inverse() * T;

  // rescale the matrix
  K = K / K(2, 2);
  R = R * std::cbrt(1.0 / R.determinant());
  R = R.transpose().eval();

  Eigen::Matrix<double, 3, 4> H;
  H.block(0, 0, 3, 3) = K * R.transpose();
  H.col(3) = -K * R.transpose() * T;
  H = H / H(2, 3);
  if (((H - P).transpose() * (H - P)).trace() > 1e-6)
  {
    std::cout << "Error: decomposition failed" << std::endl;
  }
  return;
}

//----------------------------------------------------------------------------
void GetParametersFromMatrix(const Eigen::Matrix3d& K, const Eigen::Matrix3d& R, const Eigen::Vector3d& T, Eigen::Matrix<double, 11, 1>& W)
{
  Eigen::Vector3d eulerAngles = MatrixToRollPitchYaw(R);
  W(0) = eulerAngles(0); W(1) = eulerAngles(1); W(2) = eulerAngles(2);
  W(3) = T(0), W(4) = T(1), W(5) = T(2);
  W(6) = K(0, 0);
  W(7) = K(1, 1);
  W(8) = K(0, 2);
  W(9) = K(1, 2);
  W(10) = K(0, 1);
  return;
}

//----------------------------------------------------------------------------
double NonLinearPinholeCalibration(const std::vector<Eigen::Vector3d>& X, const std::vector<Eigen::Vector2d>& x, Eigen::Matrix<double, 11, 1>& W)
{
  Eigen::Matrix<double, 3, 4> P0, P1;
  GetMatrixFromParameters(W, P0);

  // We want to estimate our 11-DOF parameters using a non
  // linear least square minimization. The non linear part
  // comes from the Euler Angle parametrization of the rotation
  // endomorphism of SO(3) and the homographie rescaling
  // To minimize it, we use CERES to perform
  // the Levenberg-Marquardt algorithm.
  ceres::Problem problem;
  for (unsigned int k = 0; k < X.size(); ++k)
  {
    ceres::CostFunction* cost_function = new ceres::AutoDiffCostFunction<CostFunctions::PinholeModelAlgebraicDistance, 1, 11>(
                                         new CostFunctions::PinholeModelAlgebraicDistance(X[k], x[k]));
    problem.AddResidualBlock(cost_function, nullptr, W.data());
  }

  // Solve the problem
  ceres::Solver::Options options;
  options.max_num_iterations = 1000;
  options.linear_solver_type = ceres::DENSE_QR;
  options.minimizer_progress_to_stdout = false;

  ceres::Solver::Summary summary;
  ceres::Solve(options, &problem, &summary);
  std::cout << summary.BriefReport() << std::endl;

  // Compute mean error
  GetMatrixFromParameters(W, P1);
  double meanErr = 0;
  for (int k = 0; k < X.size(); ++k)
  {
    Eigen::Vector4d homoX(X[k](0), X[k](1), X[k](2), 1);
    Eigen::Vector3d projX = P1 * homoX;
    Eigen::Vector2d projXn(projX(0) / projX(2), projX(1) / projX(2));
    meanErr += ((x[k] - projXn).transpose() * (x[k] - projXn))(0);
  }
 return std::sqrt(meanErr / (1.0 * X.size()));
}

//----------------------------------------------------------------------------
double NonLinearFisheyeCalibration(const std::vector<Eigen::Vector3d>& X, const std::vector<Eigen::Vector2d>& x,
                                   Eigen::Matrix<double, 15, 1>& W, unsigned int it)
{
  // We want to estimate our 15-DOF parameters using a non
  // linear least square minimization. The non linear part
  // comes from the Euler Angle parametrization of the rotation
  // endomorphism of SO(3) and the homographie rescaling
  // To minimize it, we use CERES to perform
  // the Levenberg-Marquardt algorithm.
  ceres::Problem problem;
  for (unsigned int k = 0; k < X.size(); ++k)
  {
    ceres::CostFunction* cost_function = new ceres::AutoDiffCostFunction<CostFunctions::FisheyeModelAlgebraicDistance, 1, 15>(
                                         new CostFunctions::FisheyeModelAlgebraicDistance(X[k], x[k]));
    problem.AddResidualBlock(cost_function, nullptr, W.data());
  }

  ceres::Solver::Options options;
  options.max_num_iterations = it;
  options.linear_solver_type = ceres::DENSE_QR;
  options.minimizer_progress_to_stdout = false;

  ceres::Solver::Summary summary;
  ceres::Solve(options, &problem, &summary);

  double meanErr = 0;
  for (int k = 0; k < X.size(); ++k)
  {
    meanErr += ((x[k] - FisheyeProjection(W, X[k])).transpose() * (x[k] - FisheyeProjection(W, X[k])))(0);
  }
  return std::sqrt(meanErr / (1.0 * X.size()));
}

//----------------------------------------------------------------------------
double BrownConradyPinholeCalibration(const std::vector<Eigen::Vector3d>& X, const std::vector<Eigen::Vector2d>& x,
                                      Eigen::Matrix<double, 17, 1>& W, unsigned int it,
                                      double initLossScale, double finalLossScale,
                                      const std::vector<bool>& shouldOptimizeParam)
{
  unsigned int N = 100;

  // The minimization algorithm will be ran multiple time
  // with an outlier rejection loss function more restrictive
  // at each iteration. We don't want to have a higly restrictive
  // outlier rejection at the beginning since the initial point can
  // be far from the global minimum and it could create convergence
  // issues.
  for (unsigned int minId = 0; minId < N; ++minId)
  {
    double lossScale = initLossScale + static_cast<double>(minId) * (finalLossScale - initLossScale) / (1.0 * N);

    // We want to estimate our 17-DOF parameters using a non
    // linear least square minimization. The non linear part
    // comes from the Euler Angle parametrization of the rotation
    // endomorphism of SO(3), the homographie rescaling and
    // the lens distortions
    // To minimize it, we use CERES to perform
    // the Levenberg-Marquardt algorithm.
    ceres::Problem problem;
    for (unsigned int k = 0; k < X.size(); ++k)
    {
      CostFunctions::BrownConradyAlgebraicDistance* resFct = new CostFunctions::BrownConradyAlgebraicDistance(X[k], x[k]);
      resFct->SetW0(W);
      resFct->SetActivatedParams(shouldOptimizeParam);
      ceres::CostFunction* cost_function = new ceres::AutoDiffCostFunction<CostFunctions::BrownConradyAlgebraicDistance, 1, 17>(resFct);
      problem.AddResidualBlock(cost_function, new ceres::ArctanLoss(lossScale), W.data());
    }

    ceres::Solver::Options options;
    options.max_num_iterations = it;
    options.linear_solver_type = ceres::DENSE_QR;
    options.minimizer_progress_to_stdout = false;

    ceres::Solver::Summary summary;
    ceres::Solve(options, &problem, &summary);
    std::cout << summary.BriefReport() << std::endl;
  }

  double meanErr = 0;
  for (int k = 0; k < X.size(); ++k)
  {
    meanErr += ((x[k] - BrownConradyPinholeProjection(W, X[k])).transpose() * (x[k] - BrownConradyPinholeProjection(W, X[k])))(0);
  }
  return std::sqrt(meanErr / (1.0 * X.size()));
}

//----------------------------------------------------------------------------
void GetMatrixFromParameters(const Eigen::Matrix<double, 11, 1>& W, Eigen::Matrix<double, 3, 4>& P)
{
  // Create current rotation
  Eigen::Matrix3d R = RollPitchYawToMatrix(W(0), W(1), W(2));
  // Create current position
  Eigen::Vector3d T(W(3), W(4), W(5));

  // Create current intrinsic parameters
  Eigen::Matrix3d K = Eigen::Matrix3d::Zero();
  K(0, 0) = W(6);
  K(1, 1) = W(7);
  K(0, 2) = W(8);
  K(1, 2) = W(9);
  K(0, 1) = W(10);
  K(2, 2) = 1;

  // Create current calibration matrix
  P.block(0, 0, 3, 3) = K * R.transpose();
  P.col(3) = -K * R.transpose() * T;
  P = P / P(2, 3);
}

//----------------------------------------------------------------------------
Eigen::Matrix<double, 3, 4> GetMatrixFromParameters(const Eigen::Matrix<double, 11, 1>& W)
{
  Eigen::Matrix<double, 3, 4> P;
  GetMatrixFromParameters(W, P);
  return P;
}

//----------------------------------------------------------------------------
Eigen::VectorXd FullCalibrationPipelineFromMatches(std::string filename, const std::vector<bool>& activatedParams)
{
  Eigen::VectorXd Wf = Eigen::VectorXd::Zero(17, 1);

  // Load the 3D - 2D matches
  std::vector<Eigen::Vector3d> X;
  std::vector<Eigen::Vector2d> x;
  LoadMatchesFromCSV(filename, X, x);
  if (X.size() == 0)
  {
    return Wf;
  }

  // First, launch a linear pinhole camera model
  // projection matrix estimation
  Eigen::Matrix<double, 3, 4> P;
  double rmse1 = LinearPinholeCalibration(X, x, P);

  // From this first linear pinhole projection matrix
  // estimation, extract the pinhole model parameters
  Eigen::Matrix3d K, R;
  Eigen::Vector3d T;
  CalibrationMatrixDecomposition(P, K, R, T);

  // Check that the optical axis of the camera is
  // correctly oriented according to the 3D points
  Eigen::Vector3d Xmean = Eigen::Vector3d::Zero();
  for (int i = 0; i < X.size(); ++i)
  {
    Xmean += X[i] / static_cast<double>(X.size());
  }
  Eigen::Vector3d ez = R.col(2);
  double angle = std::acos((Xmean.transpose() * ez)(0) / (Xmean.norm() * ez.norm())) / vtkMath::Pi() * 180.0;

  // In this case, the linear algorithm has provided
  // a solution where the camera is looking backward and
  // the resulting symmetry is handled by the intrinsic matrix.
  // To avoid that, we return the camera if it is not looking
  // forward
  if (angle > 90.0)
  {
    Eigen::Matrix3d S;
    S << -1.0, 0.0,  0.0,
          0.0, 1.0,  0.0,
          0.0, 0.0, -1.0;

    Eigen::Matrix3d Rtilde = R * S * R.transpose();

    // So that K' * R' = K * R, meaning that
    // the global projection is unchanged but
    // the camera orientation has made a 180
    // rotation around its y-axis
    K = K * Rtilde.transpose();
    K = K / K(2, 2);
    R = Rtilde * R;
  }

  Eigen::Matrix<double, 11, 1> Wpinhole;
  GetParametersFromMatrix(K, R, T, Wpinhole);
  Eigen::Matrix<double, 11, 1> Wpi = Wpinhole;

  // Then, refine the model obtained using linear
  // estimation by using a non-linear pinhole parameters
  // estimation
  double rmse2 = NonLinearPinholeCalibration(X, x, Wpinhole);

  // Finally, create a first parameter vector estimation
  // by using the pinhole parameters and setting the distortion
  // coefficients to 0
  Eigen::Matrix<double, 17, 1> West = Eigen::Matrix<double, 17, 1>::Zero();
  West.block(0, 0, 11, 1) = Wpinhole;

  double rmse3 = BrownConradyPinholeCalibration(X, x, West, 2500, 5.0, 0.6, activatedParams);

  // copy params
  for (int i = 0; i < 17; ++i)
  {
    Wf(i) = West(i);
  }

  Eigen::Vector3d angles(Wf(0), Wf(1), Wf(2));
  R = RollPitchYawToMatrix(angles);

  std::cout << "RMSE1: " << rmse1 << std::endl;
  std::cout << "RMSE2: " << rmse2 << std::endl;
  std::cout << "RMSE3: " << rmse3 << std::endl;
  std::cout << "W: ";
  for (int i = 0; i < 17; ++i)
  {
    std::cout << Wf(i) << ",";
  }
  std::cout << std::endl;
  return Wf;
}


#ifndef JHMI_LIVER_CONSTANTS_HPP_NRC_20150903
#define JHMI_LIVER_CONSTANTS_HPP_NRC_20150903

#include "utility/units.hpp"
#include <boost/math/constants/constants.hpp>

namespace jhmi {
//  using gamma = boost::units::static_rational<27, 10>::type;
//  using inv_gamma = boost::units::static_rational<gamma::Denominator, gamma::Numerator>::type;

  static double const pi = boost::math::double_constants::pi;
  static Pa const input_pressure = 98_mmHg;//1999 Bezy-Wendling
#if 0
#if 0
  inline Pa_s blood_viscosity(m radius) {
    return std::min(4. * pascals * seconds, Pa_s::from_value(1.8 * .6913 * (220 * std::exp(-1.3 * 1e6 * 2. * radius.value())
      + 3.2 - 2.44 * std::exp(-.06 * std::pow(1e6 * 2. * radius.value(), .645)))));
  }
#else
  inline Pa_s blood_viscosity(m) {
    return 3.5e-3 * pascals * seconds;
  }
#endif
#endif
  static Pa_s const blood_viscosity = 3.5e-3 * pascals * seconds;//Fung
  //static Pa const cell_pressure = 25_mmHg;//volmar
  //static auto const proper_ha_flow = 400. * mL / minutes;//1999 Bezy-Wendling
}
#endif


/*
Copyright 2018 Dennis Rohde

Permission is hereby granted, free of charge, to any person obtaining a copy of this software and associated documentation files (the "Software"), to deal in the Software without restriction, including without limitation the rights to use, copy, modify, merge, publish, distribute, sublicense, and/or sell copies of the Software, and to permit persons to whom the Software is furnished to do so, subject to the following conditions:

The above copyright notice and this permission notice shall be included in all copies or substantial portions of the Software.

THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE.
*/

#pragma once

#include <cmath>
#include <limits>

#include <boost/chrono/include.hpp>

#include "knn_graph.hpp"
#include "random.hpp"
#include "oracle.hpp"

class Tester_Result{
public:
    bool decision;
    double total_time;
    double query_time;
};

template <typename V = double>
class KNN_Tester {
    typedef typename KNN_Graph<V>::vertices_type vertices_type;

protected:
    bool auto_c1;

public:
    /**
     *
     * tuning parameter c1 for psi
     *
     */
    double c1 = 1;

    /**
     *
     * tuning parameter c2 for |T|
     *
     */
    double c2 = 1;

    KNN_Tester(const bool auto_c1 = true) : auto_c1{auto_c1} {}

    /**
     *
     * Calculates approximate for c1
     * @param Number of dimensions delta
     * @return approximate for c1
     *
     */
    static double c1_approximate(const KNN_Graph<V> &graph) {
        auto delta = graph.dimension();
        return std::pow(2, 0.401 * delta * (1 + 2.85 * delta / std::pow(delta, 1.4)));
    }

    /**
     *
     * Property Testing Algorithm for k-nearest Neighborhood Graphs
     * @param KNN_Graph G, average degree of G, epsilon
     * @return true or false
     *
     */
    virtual Tester_Result test(const KNN_Graph<V> &graph, const double d, const double epsilon = 0.001) {
        if (auto_c1) this->c1 = c1_approximate(graph);
        const auto delta = graph.dimension();
        const auto k = graph.get_k();
        const auto n = graph.number_vertices();
        const auto s = std::ceil(100 * k * sqrt(n) / epsilon * c2);
        const auto t = std::ceil(log(10) * c1 * k * sqrt(n));

        Uniform_Random_Generator<double> urandom_gen;

        const auto S = urandom_gen.get(s);
        const auto T = urandom_gen.get(t);

        std::cout << "|S| = " << s << std::endl;
        std::cout << "|T| = " << t << std::endl;

        bool wrongly_connected_found = false;
        double distn, distw;

        #pragma omp parallel for shared(wrongly_connected_found, distn, distw)
        for (unsigned long long i = 0; i < S.size(); ++i) {
            if (wrongly_connected_found) continue;
            const unsigned long long v = floor(S[i] * n);
            const auto v_value = graph.get_vertex(v);
            if (graph.number_neighbors(v) > 100 * k * d / epsilon) continue;
            const auto &neighbors = graph.get_edges()[v];
            V distN = 0;
            for (const auto &neighbor: neighbors) {
                auto dist = KNN_Graph<V>::euclidean_distance(v_value, graph.get_vertex(neighbor));
                if (dist > distN) {
                    distN = dist;
                }
            }
            for (unsigned long long j = 0; j < T.size(); ++j) {
                if (wrongly_connected_found) {
                    continue;
                }
                unsigned long long w = floor(T[j] * n);
                auto dist = KNN_Graph<V>::euclidean_distance(v_value, graph.get_vertex(w));
                if (v != w and distN - dist > std::numeric_limits<V>::epsilon()) {
                    auto is_neighbor = false;
                    for (const auto &neighbor: neighbors) {
                        if (w == neighbor) {
                            is_neighbor = true;
                            break;
                        }
                    }
                    if (not is_neighbor) {
                        #pragma omp critical
                        {
                            if (not wrongly_connected_found) {
                                wrongly_connected_found = true;
                                distn = distN;
                                distw = dist;
                            }
                        }
                    }
                }
            }
        }
        Tester_Result result;
        if (wrongly_connected_found) {
            std::cout << "Reject!" << std::endl;
            std::cout << distw << " < " << distn << std::endl;
            result.decision = false;
        } else {
            std::cout << "Accept!" << std::endl;
            result.decision = true;
        }
        return result;
    }

    inline auto get_auto_c1() const {
        return auto_c1;
    }

    inline void set_auto_c1(const bool auto_c1) {
        this->auto_c1 = auto_c1;
    }
};

template <typename V = double>
class KNN_Tester_Oracle : public KNN_Tester<V> {
    Query_Oracle<V> Oracle;

public:
    KNN_Tester_Oracle(const Query_Oracle<V> &oracle) : Oracle{std::move(oracle)} {}

    /**
     *
     * Property Testing Algorithm for k-Nearest Neighborhood Graphs - Oracle Version
     * @param KNN_Graph G, average degree of G, epsilon
     * @return
     *
     */
    Tester_Result test(const KNN_Graph<V> &graph, const double d, const double epsilon = 0.001) {
        auto start = boost::chrono::process_real_cpu_clock::now();
        Oracle.reset_timer();
        if (this->auto_c1) this->c1 = this->c1_approximate(graph);
        const auto delta = graph.dimension();
        const auto k = graph.get_k();
        const auto n = graph.number_vertices();
        const auto s = std::ceil(100 * k * sqrt(n) / epsilon * this->c2);
        const auto t = std::ceil(log(10) * this->c1 * k * sqrt(n));

        Uniform_Random_Generator<double> urandom_gen;

        const auto S = urandom_gen.get(s);
        const auto T = urandom_gen.get(t);

        std::cout << "|S| = " << s << std::endl;
        std::cout << "|T| = " << t << std::endl;

        bool wrongly_connected_found = false;
        double distn, distw;

        #pragma omp parallel for shared(wrongly_connected_found, distn, distw)
        for (unsigned long long i = 0; i < S.size(); ++i) {
            if (wrongly_connected_found) continue;
            const unsigned long long v = floor(S[i] * n);
            const auto v_value = graph.get_vertex(v);
            Relation<V> neighbors;
            #pragma omp critical
            {
                neighbors = Oracle.query(v);
            }
            if (neighbors.size() > 100 * k * d / epsilon) continue;
            V distN = 0;
            for (const auto &neighbor: neighbors) {
                auto dist = KNN_Graph<V>::euclidean_distance(v_value, neighbor);
                if (dist > distN) {
                    distN = dist;
                }
            }
            for (unsigned long long j = 0; j < T.size(); ++j) {
                if (wrongly_connected_found) {
                    continue;
                }
                unsigned long long w = floor(T[j] * n);
                auto dist = KNN_Graph<V>::euclidean_distance(v_value, graph.get_vertex(w));
                if (v != w and distN - dist > std::numeric_limits<V>::epsilon()) {
                    auto is_neighbor = false;
                    for (const auto &neighbor: neighbors) {
                        if (graph.get_vertex(w) == neighbor) {
                            is_neighbor = true;
                            break;
                        }
                    }
                    if (not is_neighbor) {
                        #pragma omp critical
                        {
                            if (not wrongly_connected_found) {
                                wrongly_connected_found = true;
                                distn = distN;
                                distw = dist;
                            }
                        }
                    }
                }
            }
        }
        auto stop = boost::chrono::process_real_cpu_clock::now();
        auto total_time = (stop-start).count();
        auto query_time = Oracle.time();

        Tester_Result result;
        result.total_time = total_time / 1000000000.0;
        result.query_time = query_time / 1000000000.0;
        if (wrongly_connected_found) {
            std::cout << "Reject!" << std::endl;
            std::cout << distw << " < " << distn << std::endl;
            result.decision = false;
        } else {
            std::cout << "Accept!" << std::endl;
            result.decision = true;
        }
        return result;
    }
};

template <typename V = double>
class KNN_Improver : public KNN_Tester<V> {
public:
    /**
     *
     * Property Testing Algorithm for k-Nearest Neighborhood Graphs - Graph Restauration
     * @param KNN_Graph G, average degree of G, epsilon
     * @return
     *
     */
    auto improve(KNN_Graph<V> &graph, const double d, const double epsilon = 0.001) {
        auto result = 0ul;
        if (this->auto_c1) this->c1 = this->c1_approximate(graph);
        const auto delta = graph.dimension();
        const auto k = graph.get_k();
        const auto n = graph.number_vertices();
        const auto s = ceil(100 * k * sqrt(n) / epsilon * this->c2);
        const auto t = ceil(log(10) * this->c1 * k * sqrt(n));

        Uniform_Random_Generator<double> urandom_gen;

        const auto S = urandom_gen.get(s);
        const auto T = urandom_gen.get(t);

        std::cout << "|S| = " << s << std::endl;
        std::cout << "|T| = " << t << std::endl;

        #pragma omp parallel for shared(result)
        for (unsigned long long i = 0; i < S.size(); ++i) {
            const unsigned long long v = floor(S[i] * n);
            const auto v_value = graph.get_vertex(v);
            if (graph.number_neighbors(v) > 100 * k * d / epsilon) continue;
            auto &neighbors = graph.get_edges()[v];
            V distN = 0;
            unsigned long long furthest = 0;
            for (const auto &neighbor: neighbors) {
                auto dist = KNN_Graph<V>::euclidean_distance(v_value, graph.get_vertex(neighbor));
                if (dist > distN) {
                    distN = dist;
                    furthest = neighbor;
                }
            }
            for (unsigned long long j = 0; j < T.size(); ++j) {
                unsigned long long w = floor(T[j] * n);
                auto dist = KNN_Graph<V>::euclidean_distance(v_value, graph.get_vertex(w));
                if (v != w and distN - dist > std::numeric_limits<V>::epsilon()) {
                    auto is_neighbor = false;
                    for (const auto &neighbor: neighbors) {
                        if (w == neighbor) {
                            is_neighbor = true;
                            break;
                        }
                    }
                    if (not is_neighbor) {
                        #pragma omp critical
                        {
                            graph.get_edges()[v][furthest] = w;
                            ++result;
                            distN = 0;
                            for (const auto &neighbor: neighbors) {
                                auto distF = KNN_Graph<V>::euclidean_distance(v_value, graph.get_vertex(neighbor));
                                if (distF > distN) {
                                    distN = distF;
                                    furthest = neighbor;
                                }
                            }
                        }
                    }
                }
            }
        }
        return result;
    }
};


#include <cassert>
#include <cmath>
#include <boost/graph/successive_shortest_path_nonnegative_weights.hpp>
#include <boost/property_map/property_map.hpp>

#include "mincostflowsolver.h"

void MinCostFlowSolver::solve()
{
	constructGraph();
	boost::successive_shortest_path_nonnegative_weights(
		graph, main_source, main_sink,
		boost::capacity_map(boost::get(&EdgeProp::capacity, graph))
			.residual_capacity_map(boost::get(&EdgeProp::residual_capacity, graph))
			.weight_map(boost::get(&EdgeProp::weight, graph))
			.reverse_edge_map(boost::make_assoc_property_map(reverseEdgeMap)));
	decodeSolution();
}

void MinCostFlowSolver::constructGraph()
{
	for (int r = 0; r < (int)R; ++r)
	{
		for (int c = 0; c < (int)C; ++c)
		{
			int const v = M[r][c];
			if (v == 0)
				continue;

			Vertex vx = boost::add_vertex(VertexProp{r, c}, graph);
			(v > 0 ? sources : sinks).insert(vx);
		}
	}

	for (auto u : sources)
	{
		auto const &up = graph[u];
		for (auto v : sinks)
		{
			auto const &vp = graph[v];
			int const capacity = std::min(M[up.row][up.col], -M[vp.row][vp.col]);
			int const weight = std::abs(up.row - vp.row) + std::abs(up.col - vp.col);
			addEdge(u, v, capacity, weight);
		}
	}

	main_source = boost::add_vertex(graph);
	for (auto u : sources)
	{
		auto const &up = graph[u];
		addEdge(main_source, u, M[up.row][up.col]);
	}

	main_sink = boost::add_vertex(graph);
	for (auto v : sinks)
	{
		auto const &vp = graph[v];
		addEdge(v, main_sink, -M[vp.row][vp.col]);
	}
}

void MinCostFlowSolver::decodeSolution()
{
	for (auto u : sources)
	{
		auto edges = boost::out_edges(u, graph);
		for (auto it = edges.first; it != edges.second; ++it)
		{
			auto v = boost::target(*it, graph);
			if (v == main_source)
				continue;
			auto const &ep = graph[*it];
			auto const &up = graph[u];
			auto const &vp = graph[v];
			move(up.row, up.col, vp.row, vp.col, ep.capacity - ep.residual_capacity);
		}
	}
}

void MinCostFlowSolver::addEdge(Vertex u, Vertex v, int capacity, int weight /* = 0 */)
{
	assert(capacity > 0 && weight >= 0);
	Edge e1 = boost::add_edge(u, v, EdgeProp{capacity, weight}, graph).first;
	Edge e2 = boost::add_edge(v, u, EdgeProp{0, -weight}, graph).first;
	reverseEdgeMap[e1] = e2;
	reverseEdgeMap[e2] = e1;
}


/* Copyright 2020 CNRS-AIST JRL */

#include <Eigen/Cholesky>
#include <Eigen/Jacobi>
#include <Eigen/QR>
#include <jrl-qp/GoldfarbIdnaniSolver.h>

using Givens = Eigen::JacobiRotation<double>;

namespace jrl::qp
{
GoldfarbIdnaniSolver::GoldfarbIdnaniSolver() : DualSolver(), work_d_(0), work_J_(0), work_R_(0) {}

GoldfarbIdnaniSolver::GoldfarbIdnaniSolver(int nbVar, int nbCstr, bool useBounds) : GoldfarbIdnaniSolver()
{
  resize(nbVar, nbCstr, useBounds);
}

TerminationStatus GoldfarbIdnaniSolver::solve(MatrixRef G,
                                              const VectorConstRef & a,
                                              const MatrixConstRef & C,
                                              const VectorConstRef & bl,
                                              const VectorConstRef & bu,
                                              const VectorConstRef & xl,
                                              const VectorConstRef & xu)
{
  int nbVar = G.rows();
  int nbCstr = C.cols();
  bool useBnd = xl.size() > 0;

  LOG_RESET(log_);
  LOG(log_, LogFlags::INPUT | LogFlags::NO_ITER, G, a, C, bl, bu, xl, xu);

  assert(G.cols() == nbVar);
  assert(a.size() == nbVar);
  assert(C.rows() == nbVar);
  assert(bl.size() == nbCstr);
  assert(bu.size() == nbCstr);
  assert(xl.size() == nbVar || xl.size() == 0);
  assert(xu.size() == xl.size());

  // TODO check input: bl<=bu, xl<=xu, ...

  new(&pb_.G) MatrixRef(G);
  new(&pb_.a) VectorConstRef(a);
  new(&pb_.C) MatrixConstRef(C);
  new(&pb_.bl) VectorConstRef(bl);
  new(&pb_.bu) VectorConstRef(bu);
  new(&pb_.xl) VectorConstRef(xl);
  new(&pb_.xu) VectorConstRef(xu);

  resize(nbVar, nbCstr, useBnd);

  return DualSolver::solve();
}

internal::InitTermination GoldfarbIdnaniSolver::init_()
{
  int ret = Eigen::internal::llt_inplace<double, Eigen::Lower>::blocked(pb_.G);
  auto L = pb_.G.template triangularView<Eigen::Lower>();

  if(ret >= 0) return TerminationStatus::NON_POS_HESSIAN;

  // J = L^-t
  auto J = work_J_.asMatrix(nbVar_, nbVar_, nbVar_);
  J.setIdentity();
  L.transpose().solveInPlace(J);

  // x = -G^-1 * a
  auto x = work_x_.asVector(nbVar_);
  x = L.solve(pb_.a);
  L.transpose().solveInPlace(x); // possible [OPTIM]: J already contains L^-T
  x = -x;
  f_ = 0.5 * pb_.a.dot(x);

  A_.reset();
  DEBUG_ONLY(work_R_.setZero());

  // Adding equality constraints
  initActiveSet();

  return TerminationStatus::SUCCESS;
}

internal::ConstraintNormal GoldfarbIdnaniSolver::selectViolatedConstraint_(const VectorConstRef & x) const
{
  // We look for the constraint with the maximum violation
  //[NUMERIC] scale with constraint magnitude
  double smin = 0;
  int p = -1;
  ActivationStatus status = ActivationStatus::INACTIVE;

  // Check general constraints
  for(int i = 0; i < A_.nbCstr(); ++i)
  {
    if(!A_.isActive(i))
    {
      double cx = pb_.C.col(i).dot(x); // possible [OPTIM]: should we compute C^T x at once ?
      if(double sl = cx - pb_.bl[i]; sl < smin)
      {
        smin = sl;
        p = i;
        status = ActivationStatus::LOWER;
      }
      else if(double su = pb_.bu[i] - cx; su < smin)
      {
        smin = su;
        p = i;
        status = ActivationStatus::UPPER;
      }
    }
  }

  // Check bound constraints
  for(int i = 0; i < A_.nbBnd(); ++i)
  {
    if(!A_.isActiveBnd(i))
    {
      if(double sl = x[i] - pb_.xl[i]; sl < smin)
      {
        smin = sl;
        p = A_.nbCstr() + i;
        status = ActivationStatus::LOWER_BOUND;
      }
      else if(double su = pb_.xu[i] - x[i]; su < smin)
      {
        smin = su;
        p = A_.nbCstr() + i;
        status = ActivationStatus::UPPER_BOUND;
      }
    }
  }

  return {pb_.C, p, status};
}

void GoldfarbIdnaniSolver::computeStep_(VectorRef z, VectorRef r, const internal::ConstraintNormal & np) const
{
  int q = A_.nbActiveCstr();
  auto d = work_d_.asVector(nbVar_, {});
  auto J = work_J_.asMatrix(nbVar_, nbVar_, nbVar_);
  auto R = work_R_.asMatrix(q, q, nbVar_).template triangularView<Eigen::Upper>();

  np.preMultiplyByMt(d, J);
  z.noalias() = J.rightCols(nbVar_ - q) * d.tail(nbVar_ - q);
  r = R.solve(d.head(q));
  DBG(log_, LogFlags::ITERATION_ADVANCE_DETAILS, J, R, d);
}

DualSolver::StepLength GoldfarbIdnaniSolver::computeStepLength_(const internal::ConstraintNormal & np,
                                                                const VectorConstRef & x,
                                                                const VectorConstRef & u,
                                                                const VectorConstRef & z,
                                                                const VectorConstRef & r) const
{
  double t1 = options_.bigBnd_;
  double t2 = options_.bigBnd_;
  int l = 0;

  for(int k = 0; k < A_.nbActiveCstr(); ++k)
  {
    if(A_.activationStatus(k) != ActivationStatus::EQUALITY && A_.activationStatus(k) != ActivationStatus::FIXED
       && r[k] > 0)
    {
      if(double tk = u[k] / r[k]; tk < t1)
      {
        t1 = tk;
        l = k;
      }
    }
  }

  if(z.norm() > 1e-14) //[NUMERIC] better criterion
  {
    double b, cx, cz;
    int p, pb;
    switch(np.status())
    {
      case ActivationStatus::LOWER:
        p = np.index();
        b = pb_.bl[p];
        cx = pb_.C.col(p).dot(x);
        cz = pb_.C.col(p).dot(z);
        break;
      case ActivationStatus::UPPER:
        p = np.index();
        b = pb_.bu[p];
        cx = pb_.C.col(p).dot(x);
        cz = pb_.C.col(p).dot(z);
        break;
      case ActivationStatus::EQUALITY:
        assert(false);
        break;
      case ActivationStatus::LOWER_BOUND:
        pb = np.bndIndex();
        b = pb_.xl[pb];
        cx = x[pb];
        cz = z[pb];
        break;
      case ActivationStatus::UPPER_BOUND:
        pb = np.bndIndex();
        b = pb_.xu[pb];
        cx = x[pb];
        cz = z[pb];
        break;
      case ActivationStatus::FIXED:
        assert(false);
        break;
      default:
        assert(false);
    }
    t2 = (b - cx) / cz;
  }

  return {t1, t2, l};
}

bool GoldfarbIdnaniSolver::addConstraint_(const internal::ConstraintNormal & np)
{
  int q = A_.nbActiveCstr(); // This already counts the new constraint
  auto d = work_d_.asVector(nbVar_);
  auto J = work_J_.asMatrix(nbVar_, nbVar_, nbVar_);
  for(int i = nbVar_ - 2; i >= q - 1; --i) //[OPTIM] use Householder transformation instead
  {
    Givens Qi;
    Qi.makeGivens(d[i], d[i + 1], &d[i]);
    DEBUG_ONLY(d[i + 1] = 0);
    J.applyOnTheRight(i, i + 1, Qi);
  }
  auto R = work_R_.asMatrix(q, q, nbVar_);
  R.rightCols<1>() = d.head(q);

  return true; //[NUMERIC]: add test on dependency
}

bool GoldfarbIdnaniSolver::removeConstraint_(int l)
{
  int q = A_.nbActiveCstr(); // This already counts that the constraint was removed
  auto d = work_d_.asVector(nbVar_);
  auto J = work_J_.asMatrix(nbVar_, nbVar_, nbVar_);
  auto R = work_R_.asMatrix(q + 1, q + 1, nbVar_);

  for(int i = l; i < q; ++i)
  {
    Givens Qi;
    R.col(i).head(i) = R.col(i + 1).head(i);
    Qi.makeGivens(R(i, i + 1), R(i + 1, i + 1), &R(i, i));
    DEBUG_ONLY(R(i + 1, i + 1) = 0);
    R.rightCols(q - i - 1).applyOnTheLeft(i, i + 1, Qi.transpose());
    J.applyOnTheRight(i, i + 1, Qi);
  }

  return true;
}

void GoldfarbIdnaniSolver::resize_(int nbVar, int nbCstr, bool useBounds)
{
  if(nbVar != nbVar_)
  {
    work_d_.resize(nbVar);
    work_J_.resize(nbVar, nbVar);
    work_R_.resize(nbVar, nbVar);
  }
}

void GoldfarbIdnaniSolver::initActiveSet()
{
  for(int i = 0; i < A_.nbCstr(); ++i)
  {
    if(pb_.bl[i] == pb_.bu[i])
    {
      internal::ConstraintNormal np(pb_.C, i, ActivationStatus::EQUALITY);
      addInitialConstraint(np);
    }
  }

  for(int i = 0; i < A_.nbBnd(); ++i)
  {
    if(pb_.xl[i] == pb_.xu[i])
    {
      internal::ConstraintNormal np(pb_.C, A_.nbCstr() + i, ActivationStatus::FIXED);
      addInitialConstraint(np);
    }
  }
}

void GoldfarbIdnaniSolver::addInitialConstraint(const internal::ConstraintNormal & np)
{
  int q = A_.nbActiveCstr();
  WVector x = work_x_.asVector(nbVar_);
  WVector z = work_z_.asVector(nbVar_);
  WVector u = work_u_.asVector(q + 1);
  WVector r = work_r_.asVector(q);
  u[q] = 0;

  computeStep(z, r, np);

  assert(np.status() == ActivationStatus::EQUALITY || np.status() == ActivationStatus::FIXED);
  double t = 0;
  if(z.norm() > 1e-14) //[NUMERIC] better criterion
  {
    if(np.status() == ActivationStatus::EQUALITY) //[OPTIM] we can avoid this if by specializing the function to general
                                                  // constraint or bound
    {
      int p = np.index();
      t = (pb_.bl[p] - pb_.C.col(p).dot(x)) / pb_.C.col(p).dot(z);
    }
    else
    {
      int pb = np.bndIndex();
      t = (pb_.xl[pb] - x[pb]) / z[pb];
    }
  }
  else
  {
    // numerical problem
  }

  x += t * z;
  f_ += t * np.dot(z) * (.5 * t + u[q]);
  // u = u + t*[-r;1]
  u.head(q) -= t * r;
  u[q] += t;
  if(!addConstraint(np))
  {
    LOG_COMMENT(log_, LogFlags::TERMINATION, "Attempting to add a linearly dependent constraint.");
    // return TerminationStatus::LINEAR_DEPENDENCY_DETECTED;
  }
}
} // namespace jrl::qp


#include "naivebayesfeature.h"
#include "utilities.h"
#include <boost/foreach.hpp>

using namespace std;

NaiveBayesFeature::NaiveBayesFeature(const HyperParameters &hp) : NaiveBayes( hp ) {
    m_featureType = (hp.naiveBayesFeatureType == "Gaussian") ? FEATURE_GAUSSIAN : FEATURE_HISTOGRAM;
}

void NaiveBayesFeature::calcMeanAndVariance(const matrix<float>& data, const std::vector<int>& labels, Feature& f) {
    std::vector<double> mean(m_hp.numClasses,0.0);
    std::vector<double> variance(m_hp.numClasses,0.0);
    std::vector<int> numCounts(m_hp.numClasses,0);

    BOOST_FOREACH(int sample, m_inBagSamples) {
        mean[labels[sample]] += data(sample,f.index);
        numCounts[labels[sample]]++;
        variance[labels[sample]] += pow((double)data(sample,f.index),2.0);
    }

    for (int c = 0; c < m_hp.numClasses; c++) {
        mean[c] /= numCounts[c];  // mu
        variance[c] /= numCounts[c]; // E(X^2)
        variance[c] -= pow(mean[c], 2.0);
    }

    f.mean = mean;
    f.variance = variance;
}

// Weighted Statistics
void NaiveBayesFeature::calcMeanAndVariance(const matrix<float>& data, const std::vector<int>& labels, Feature& f,
        const std::vector<double>& weights) {
    std::vector<double> mean(m_hp.numClasses,0.0);
    std::vector<double> variance(m_hp.numClasses,0.0);
    std::vector<double> numCounts(m_hp.numClasses,0.0);

    BOOST_FOREACH(int sample, m_inBagSamples) {
        mean[labels[sample]] += data(sample,f.index)*weights[sample];
        numCounts[labels[sample]] += weights[sample];
        variance[labels[sample]] += pow((double)data(sample,f.index)*weights[sample],2.0);
    }

    for (int c = 0; c < m_hp.numClasses; c++) {
        mean[c] /= numCounts[c];  // mu
        variance[c] /= numCounts[c]; // E(X^2)
        variance[c] -= pow(mean[c], 2.0);
    }

    f.mean = mean;
    f.variance = variance;
}

void NaiveBayesFeature::calcHistogram(const matrix<float>& data, const std::vector<int>& labels, Feature& f) {
    // find min and max values
    f.min = data(m_inBagSamples[0],f.index);
    f.max = data(m_inBagSamples[0],f.index);

    BOOST_FOREACH(int sample, m_inBagSamples) {
        f.min = (data(sample,f.index) < f.min) ? data(sample,f.index) : f.min;
        f.max = (data(sample,f.index) > f.max) ? data(sample,f.index) : f.max;
    }
    f.numBins = (float)m_hp.numHistogramBins;

    std::vector<float> classHistogram((int) f.numBins, 0.0);
    std::vector<std::vector<float> > histogram(m_hp.numClasses,classHistogram);
    std::vector<int> classCounter(m_hp.numClasses,0);

    if ( f.numBins > 2.0) {
        f.w = abs(f.min - f.max)/f.numBins;

        BOOST_FOREACH(int sample, m_inBagSamples) {
            int binIndex = (int) ((data(sample, f.index) - f.min)/f.w);
            binIndex -= (binIndex == f.numBins) ? 1 : 0;
            histogram[labels[sample]][binIndex]++;
            classCounter[labels[sample]]++;
        }
    } else {
        const float numSteps = 20;
        const float stepSize = abs(f.min - f.max)/numSteps;
        float bestEntropy = 1e10;
        float bestThreshold = 0.0;
        std::vector<int> decision(data.size1());
        // Find the best threshold
        for (double threshold = f.min ; threshold < f.max ; threshold += stepSize) {
            int oneCount = 0, zeroCount = 0;
            std::vector<double> zeroClassCount(m_hp.numClasses, 0.0), oneClassCount(m_hp.numClasses, 0.0);

            BOOST_FOREACH(int sample, m_inBagSamples) {
                if (data(sample, f.index) > threshold) {
                    decision[sample] = 1;
                    oneClassCount[labels[sample]]++;
                    oneCount++;
                } else {
                    decision[sample] = 0;
                    zeroClassCount[labels[sample]]++;
                    zeroCount++;
                }
            }

            // Calc entropy
            double oneEntropy = 0, zeroEntropy = 0;
            for (int nClass = 0; nClass < m_hp.numClasses; nClass++) {
                oneClassCount[nClass] /= oneCount;
                zeroClassCount[nClass] /= zeroCount;

                if (oneClassCount[nClass]) {
                    oneEntropy -= oneClassCount[nClass]*log(oneClassCount[nClass]);
                }
                if (zeroClassCount[nClass]) {
                    zeroEntropy -= zeroClassCount[nClass]*log(zeroClassCount[nClass]);
                }
            }

            // Total Entropy
            double entropy = (zeroEntropy*zeroCount + oneEntropy*oneCount)/(zeroCount + oneCount);
            if (entropy < bestEntropy) {
                bestEntropy = entropy;
                bestThreshold = threshold;
            }
        }
        f.threshold = bestThreshold;

        BOOST_FOREACH(int sample, m_inBagSamples) {
            int binIndex = (data(sample, f.index) > f.threshold) ? 1 : 0;
            histogram[labels[sample]][binIndex]++;
            classCounter[labels[sample]]++;
        }
    }

    std::vector<std::vector<float> >::iterator it(histogram.begin());
    std::vector<std::vector<float> >::iterator end(histogram.end());
    for (int c = 0;it != end; it++, c++) {
        for (int bin = 0; bin < f.numBins; bin++) {
            (*it)[bin] /= (classCounter[c] + 1e-10);
        }
    }
    f.histogram = histogram;
}

// Weighted Discrete Statistics
void NaiveBayesFeature::calcHistogram(const matrix<float>& data, const std::vector<int>& labels, Feature& f,
                                      const std::vector<double>& weights) {
    // find min and max values
    f.min = data(m_inBagSamples[0],f.index);
    f.max = data(m_inBagSamples[0],f.index);

    BOOST_FOREACH(int sample, m_inBagSamples) {
        f.min = (data(sample,f.index) < f.min) ? data(sample,f.index) : f.min;
        f.max = (data(sample,f.index) > f.max) ? data(sample,f.index) : f.max;
    }
    f.numBins = (float)m_hp.numHistogramBins;

    std::vector<float> classHistogram((int) f.numBins, 0.0);
    std::vector<std::vector<float> > histogram(m_hp.numClasses,classHistogram);
    std::vector<double> classCounter(m_hp.numClasses,0.0);

    if ( f.numBins > 2.0) {
        f.w = abs(f.min - f.max)/f.numBins;

        BOOST_FOREACH(int sample, m_inBagSamples) {
            int binIndex = (int) ((data(sample,f.index) - f.min)/f.w);
            binIndex -= (binIndex == f.numBins) ? 1 : 0;
            histogram[labels[sample]][binIndex] += weights[sample];
            classCounter[labels[sample]] += weights[sample];
        }
    } else {
        const float numSteps = 20;
        const float stepSize = abs(f.min - f.max)/numSteps;
        float bestEntropy = 1e10;
        float bestThreshold = 0.0;
        std::vector<int> decision(data.size1());
        // Find the best threshold
        for (double threshold = f.min ; threshold < f.max ; threshold += stepSize) {
            double oneCount = 0, zeroCount = 0;
            std::vector<double> zeroClassCount(m_hp.numClasses, 0.0), oneClassCount(m_hp.numClasses, 0.0);

            BOOST_FOREACH(int sample, m_inBagSamples) {
                if (data(sample, f.index) > threshold) {
                    decision[sample] = 1;
                    oneClassCount[labels[sample]] += weights[sample];
                    oneCount += weights[sample];
                } else {
                    decision[sample] = 0;
                    zeroClassCount[labels[sample]] += weights[sample];
                    zeroCount += weights[sample];
                }
            }

            // Calc entropy
            double oneEntropy = 0, zeroEntropy = 0;
            for (int nClass = 0; nClass < m_hp.numClasses; nClass++) {
                oneClassCount[nClass] /= oneCount;
                zeroClassCount[nClass] /= zeroCount;

                if (oneClassCount[nClass]) {
                    oneEntropy -= oneClassCount[nClass]*log(oneClassCount[nClass]);
                }
                if (zeroClassCount[nClass]) {
                    zeroEntropy -= zeroClassCount[nClass]*log(zeroClassCount[nClass]);
                }
            }

            // Total Entropy
            double entropy = (zeroEntropy*zeroCount + oneEntropy*oneCount)/(zeroCount + oneCount);
            if (entropy < bestEntropy) {
                bestEntropy = entropy;
                bestThreshold = threshold;
            }
        }
        f.threshold = bestThreshold;

        BOOST_FOREACH(int sample, m_inBagSamples) {
            int binIndex = (data(sample, f.index) > f.threshold) ? 1 : 0;
            histogram[labels[sample]][binIndex] += weights[sample];
            classCounter[labels[sample]] += weights[sample];
        }
    }

    // Normalise Bins
    std::vector<std::vector<float> >::iterator it(histogram.begin());
    std::vector<std::vector<float> >::iterator end(histogram.end());
    for (int c = 0;it != end; it++, c++) {
        for (int bin = 0; bin < f.numBins; bin++) {
            (*it)[bin] /= (classCounter[c] + 1e-10);
        }
    }
    f.histogram = histogram;
}

void NaiveBayesFeature::train(const matrix<float>& data, const std::vector<int>& labels,
                              matrix<float>& forestConfidences, matrix<float>& forestOutOfBagConfidences,
                              std::vector<int>& forestOutOfBagVoteNum) {
    // Initialize
    initialize(m_hp.numLabeled);

    // Random Subsamples data according to bagratio
    subSample(m_hp.numLabeled);

    // Train the Naive Bayes Classifiers
    std::vector<int> randomFeatures = randPerm(data.size2(),m_hp.numRandomFeatures);

    Feature f;
    f.type = m_featureType;
    if (m_featureType == FEATURE_GAUSSIAN) {
        BOOST_FOREACH(int n, randomFeatures) {
            f.index = n;
            calcMeanAndVariance(data,labels,f);
            m_features.push_back(f);
        }
    } else { // FEATURE_HISTOGRAM
        BOOST_FOREACH(int n, randomFeatures) {
            f.index = n;
            calcHistogram(data,labels,f);
            m_features.push_back(f);
        }
    }

    eval(data,labels);

    finalize(data, forestConfidences, forestOutOfBagConfidences, forestOutOfBagVoteNum);
}

// Weighted Training
void NaiveBayesFeature::train(const matrix<float>& data, const std::vector<int>& labels, const std::vector<double>& weights,
                              matrix<float>& forestConfidences, matrix<float>& forestOutOfBagConfidences,
                              std::vector<int>& forestOutOfBagVoteNum, bool init)
{
    if (init)
    {
        // Initialize
        initialize(m_hp.numLabeled);
        // Random Subsamples data according to bagratio
        subSample(m_hp.numLabeled);
    }
    // Train the Naive Bayes Classifiers
    std::vector<int> randomFeatures = randPerm(data.size2(),m_hp.numRandomFeatures);

    Feature f;
    f.type = m_featureType;
    if (m_featureType == FEATURE_GAUSSIAN)
    {
        BOOST_FOREACH(int n, randomFeatures)
        {
            f.index = n;
            calcMeanAndVariance(data,labels,f);
            m_features.push_back(f);
        }
    }
    else   // FEATURE_HISTOGRAM
    {
        BOOST_FOREACH(int n, randomFeatures)
        {
            f.index = n;
            calcHistogram(data,labels,f,weights);
            m_features.push_back(f);
        }
    }

    eval(data,labels);

    finalize(data, forestConfidences, forestOutOfBagConfidences, forestOutOfBagVoteNum);
}

// Weighted Training
void NaiveBayesFeature::retrain(const matrix<float>& data, const std::vector<int>& labels, const std::vector<double>& weights,
                              matrix<float>& forestConfidences, matrix<float>& forestOutOfBagConfidences,
                              std::vector<int>& forestOutOfBagVoteNum, bool init)
{
    if (init)
    {
        // Initialize
        initialize(data.size1());
        // Random Subsamples data according to bagratio
        subSample(data.size1());
    }
    // Train the Naive Bayes Classifiers
    std::vector<int> randomFeatures = randPerm(data.size2(),m_hp.numRandomFeatures);

    Feature f;
    f.type = m_featureType;
    if (m_featureType == FEATURE_GAUSSIAN) {
        BOOST_FOREACH(int n, randomFeatures) {
            f.index = n;
            calcMeanAndVariance(data,labels,f,weights);
            m_features.push_back(f);
        }
    } else { // FEATURE_HISTOGRAM
        BOOST_FOREACH(int n, randomFeatures) {
            f.index = n;
            calcHistogram(data,labels,f,weights);
            m_features.push_back(f);
        }
    }

    eval(data,labels);

    finalize(data, forestConfidences, forestOutOfBagConfidences, forestOutOfBagVoteNum);
}

void NaiveBayesFeature::evalOutOfBagSamples(const matrix<float>& data) {

}

void NaiveBayesFeature::eval(const matrix<float>& data, const std::vector<int> labels) {
    m_confidences.resize(data.size1(), m_hp.numClasses);
    // init to one (due to multiplication)
    for (int sample = 0; sample < (int)data.size1(); sample++) {
        for ( int c = 0; c < m_hp.numClasses; c++) {
            m_confidences(sample,c) = 1.0;
        }
    }

    // Fill confidence matrix
    std::vector<Feature>::iterator it(m_features.begin());
    std::vector<Feature>::iterator end(m_features.end());
    for (int sample = 0; sample < (int)data.size1(); sample++) {
        it = m_features.begin();
        while (it != end) {
            it->eval(data,sample,m_confidences);
            ++it;
        }
        double max = 0.0;
        for ( int c = 0; c < m_hp.numClasses; c++) {
            if (m_confidences(sample,c) > max) {
                m_predictions[sample] = c;
                max = m_confidences(sample,c);
            }
        }
    }

    double error = computeError(labels);
    if (m_hp.verbose) {
        cout << "Error: " << error << endl;
    }

}

void NaiveBayesFeature::eval(const matrix<float>& data, matrix<float>& confidences) {
    m_confidences.resize(data.size1(), m_hp.numClasses);
    // init to one (due to multiplication)
    for (int sample = 0; sample < (int)data.size1(); sample++) {
        for ( int c = 0; c < m_hp.numClasses; c++) {
            m_confidences(sample,c) = 1.0;
        }
    }

    // Fill confidence matrix
    std::vector<Feature>::iterator it(m_features.begin());
    std::vector<Feature>::iterator end(m_features.end());
    for (int sample = 0; sample < (int)data.size1(); sample++) {
        it = m_features.begin();
        while (it != end) {
            it->eval(data,sample,m_confidences);
            ++it;
        }
        for ( int c = 0; c < m_hp.numClasses; c++) {
            confidences(sample,c) += m_confidences(sample,c);
        }
    }

}

void Feature::eval(const matrix<float>& data, const int sampleIndex, matrix<float>& confidences) {
    if ( type == FEATURE_GAUSSIAN ) {
        evalGaussian(data,sampleIndex,confidences);
    } else {
        evalHistogram(data,sampleIndex,confidences);
    }
}

void Feature::evalGaussian(const matrix<float>& data, const int sampleIndex, matrix<float>& confidences) {
    for (int c = 0; c < (int)mean.size(); c++) {
        confidences(sampleIndex,c) *= 1.0/variance[c] * exp(-pow(data(sampleIndex,index) - mean[c],2.0)/(2.0*pow(variance[c],2.0)));
    }
}

void Feature::evalHistogram(const matrix<float>& data, const int sampleIndex, matrix<float>& confidences) {
  int binIndex;
  if (numBins > 2) {
    binIndex = (int) ((data(sampleIndex,index) - min)/w);
    binIndex -= (binIndex == numBins) ? 1 : 0;
    binIndex = (binIndex < 0.0) ? 0 : binIndex;
  }
  else {
    binIndex = (data(sampleIndex, index) > threshold) ? 1 : 0;
  }
  for (int c = 0; c < (int)histogram.size(); c++) {
    confidences(sampleIndex,c) *= histogram[c][binIndex];
  }
}


//=======================================================================
// Copyright (c) 2013 Piotr Smulewicz
//
// Distributed under the Boost Software License, Version 1.0. (See
// accompanying file LICENSE_1_0.txt or copy at
// http://www.boost.org/LICENSE_1_0.txt)
//=======================================================================
/**
 * @file multiway_cut.hpp
 * @brief
 * @author Piotr Smulewicz, Piotr Godlewski
 * @version 1.0
 * @date 2013-12-19
 */

#ifndef PAAL_MULTIWAY_CUT_HPP
#define PAAL_MULTIWAY_CUT_HPP

#include "paal/lp/glp.hpp"
#include "paal/utils/type_functions.hpp"
#include "paal/utils/irange.hpp"
#include "paal/utils/assign_updates.hpp"

#include <boost/bimap.hpp>
#include <boost/graph/breadth_first_search.hpp>
#include <boost/graph/connected_components.hpp>
#include <boost/graph/filtered_graph.hpp>
#include <boost/graph/graph_traits.hpp>
#include <boost/graph/graph_utility.hpp>
#include <boost/graph/named_function_params.hpp>
#include <boost/graph/stoer_wagner_min_cut.hpp>
#include <boost/property_map/property_map.hpp>
#include <boost/range/as_array.hpp>

#include <fstream>
#include <tuple>
#include <utility>
#include <vector>
#include <random>

namespace paal {
namespace detail {

inline int vertices_column_index(int vertex, int dimentions, int column) {
    return vertex * dimentions + column;
}

template <class Graph>
using CostType = typename boost::property_traits<
    puretype(get(boost::edge_weight, std::declval<Graph>()))>::value_type;

template <typename LP> class multiway_cut_lp {
  public:
    /// Initialize the cut LP.
    template <typename Graph, typename IndexMap, typename WeightMap,
              typename ColorMap>
    void init(const Graph &graph, int k, const IndexMap &index_map,
              const WeightMap &weight_map, const ColorMap &color_map) {

        m_lp.set_lp_name("Multiway Cut");
        m_lp.set_optimization_type(lp::MINIMIZE);

        add_variables(graph, k, weight_map);
        add_constraints(graph, k, index_map, color_map);
    }

  private:
    // adding variables
    // returns the number of variables
    template <typename Graph, typename WeightMap>
    void add_variables(const Graph &graph, int k, const WeightMap &weight_map) {
        for (auto e : boost::as_array(edges(graph))) {
            for (int i = 0; i < k; ++i) {
                auto col_idx = m_lp.add_column(get(weight_map, e));
                edges_column.push_back(col_idx);
            }
        }
        for (unsigned vertex = 0; vertex <= num_vertices(graph); ++vertex) {
            for (int i = 0; i < k; ++i) {
                auto col_idx = m_lp.add_column(0);
                vertices_column.push_back(col_idx);
            }
        }
    }

    template <typename Graph, typename IndexMap, typename ColorMap>
    void add_constraints(const Graph &graph, int k, const IndexMap &index_map,
                         const ColorMap &color_map) {
        int db_index = 0;
        for (auto edge : boost::as_array(edges(graph))) {
            auto sour = get(index_map, source(edge, graph));
            auto targ = get(index_map, target(edge, graph));
            for (auto i : irange(k)) {
                for (auto j : irange(2)) {
                    auto x_e =
                        edges_column[vertices_column_index(db_index, k, i)];
                    auto x_src =
                        vertices_column[vertices_column_index(sour, k, i)];
                    auto x_trg =
                        vertices_column[vertices_column_index(targ, k, i)];
                    m_lp.add_row(
                        x_e + (j * 2 - 1) * x_src + (1 - 2 * j) * x_trg >= 0);
                }
            }
            ++db_index;
        }
        db_index = 0;
        for (auto vertex : boost::as_array(vertices(graph))) {
            auto col = get(color_map, vertex);
            if (col != 0) {
                auto x_col = vertices_column[
                    vertices_column_index(db_index, k, col - 1)];
                m_lp.add_row(x_col == 1);
            }
            lp::linear_expression expr;
            for (auto i : irange(k)) {
                expr += vertices_column[vertices_column_index(db_index, k, i)];
            }
            m_lp.add_row(std::move(expr) == 1);
            ++db_index;
        }
    }

  public:
    LP m_lp;
    std::vector<lp::col_id> edges_column;
    std::vector<lp::col_id> vertices_column;
};

template <typename Graph, typename VertexIndexMap, typename EdgeWeightMap,
          typename Dist, typename Rand, typename LP>
auto make_cut(const Graph &graph, int k, const VertexIndexMap &index_map,
              const EdgeWeightMap &weight_map, Dist &dist, Rand &&random_engine,
              multiway_cut_lp<LP> &mc_lp, std::vector<int> &vertex_to_part)
    ->detail::CostType<Graph> {
    double cut_cost = 0;
    std::vector<double> random_radiuses;
    dist.reset();
    for (int i = 0; i < k; ++i) {
        random_radiuses.push_back(dist(random_engine));
    }
    vertex_to_part.resize(num_vertices(graph));
    auto get_column = [&](int vertex, int dimension) {
        return mc_lp.m_lp.get_col_value(
            mc_lp.vertices_column[vertices_column_index(vertex, k, dimension)]);
    };

    for (auto vertex : boost::as_array(vertices(graph))) {
        for (int dimension = 0; dimension < k; ++dimension)
            if (1.0 - get_column(get(index_map, vertex), dimension) <
                    random_radiuses[dimension] ||
                dimension == k - 1) {
                // because each vertex have sum of coordinates equal 1,
                // 1.0-get_column(vertex,dimension) is proportional to distance
                // to vertex correspond to dimension
                vertex_to_part[get(index_map, vertex)] = dimension;
                break;
            }
    }
    for (auto edge : boost::as_array(edges(graph))) {
        if (vertex_to_part[get(index_map, source(edge, graph))] !=
            vertex_to_part[get(index_map, target(edge, graph))])
            cut_cost += get(weight_map, edge);
    }
    return cut_cost;
}

template <typename Rand = std::default_random_engine,
          typename Distribution = std::uniform_real_distribution<double>,
          typename LP = lp::glp, typename Graph, typename OutputIterator,
          typename VertexIndexMap, typename EdgeWeightMap,
          typename VertexColorMap>
auto multiway_cut_dispatch(const Graph &graph, OutputIterator result,
                           Rand &&random_engine, int iterations,
                           VertexIndexMap index_map, EdgeWeightMap weight_map,
                           VertexColorMap color_map)
    ->typename boost::property_traits<EdgeWeightMap>::value_type {
    using CostType = detail::CostType<Graph>;
    Distribution dis(0, 1);
    int terminals = 0;
    for (auto vertex : boost::as_array(vertices(graph))) {
        assign_max(terminals, get(color_map, vertex));
    }
    detail::multiway_cut_lp<LP> multiway_cut_lp;
    multiway_cut_lp.init(graph, terminals, index_map, weight_map, color_map);
    multiway_cut_lp.m_lp.solve_simplex(lp::DUAL);
    CostType cut_cost = std::numeric_limits<CostType>::max();
    std::vector<int> best_solution;
    std::vector<int> solution;
    for (int i = 0; i < iterations; ++i) {
        solution.clear();
        int res = detail::make_cut(graph, terminals, index_map, weight_map, dis,
                                   random_engine, multiway_cut_lp, solution);
        if (res < cut_cost) {
            swap(solution, best_solution);
            cut_cost = res;
        }
    }
    for (auto v : boost::as_array(vertices(graph))) {
        *result = std::make_pair(v, best_solution[get(index_map, v)]);
        ++result;
    }
    return cut_cost;
}

template <typename Rand = std::default_random_engine,
          typename Distribution = std::uniform_real_distribution<double>,
          typename LP = lp::glp, typename Graph, typename OutputIterator,
          typename VertexIndexMap, typename EdgeWeightMap,
          typename VertexColorMap>
auto multiway_cut_dispatch(const Graph &graph, OutputIterator result,
                           Rand &&random_engine, boost::param_not_found,
                           VertexIndexMap index_map, EdgeWeightMap weight_map,
                           VertexColorMap color_map)
    ->typename boost::property_traits<EdgeWeightMap>::value_type {
    int vertices = num_vertices(graph);
    const static int MIN_NUMBER_OF_REPEATS = 100;
    auto number_of_repeats =
        vertices * vertices +
        MIN_NUMBER_OF_REPEATS; // This variable is not supported by any proof
    return multiway_cut_dispatch(graph, result, random_engine,
                                 number_of_repeats, index_map, weight_map,
                                 color_map);
}

} //!detail

/**
 * @brief this is solve multiway_cut problem
 * and return cut_cost
 * example:
 *  \snippet multiway_cut_example.cpp  Multiway Cut Example
 * @param Graph graph
 * @param OutputIterator result pairs of vertex descriptor and number form (1,2,
 * ... ,k) id of part
 * @param random_engine
 * @param params
 * @tparam Graph
 * @tparam OutputIterator
 * @tparam Rand random engine
 * @tparam Distribution used to chose random radius
 * @tparam LP
 * @tparam P
 * @tparam T
 * @tparam R
 */
template <typename Rand = std::default_random_engine,
          typename Distribution = std::uniform_real_distribution<double>,
          typename LP = lp::glp, typename Graph, typename OutputIterator,
          typename P, typename T, typename R>
auto multiway_cut(const Graph &g, OutputIterator out,
                  const boost::bgl_named_params<P, T, R> &params,
                  Rand &&random_engine = std::default_random_engine(5426u))
    ->typename boost::property_traits<puretype(
          boost::choose_const_pmap(get_param(params, boost::edge_weight), g,
                                   boost::edge_weight))>::value_type {
    return detail::multiway_cut_dispatch(
        g, out, random_engine, get_param(params, boost::iterations_t()),
        boost::choose_const_pmap(get_param(params, boost::vertex_index), g,
                                 boost::vertex_index),
        boost::choose_const_pmap(get_param(params, boost::edge_weight), g,
                                 boost::edge_weight),
        boost::choose_const_pmap(get_param(params, boost::vertex_color), g,
                                 boost::vertex_color));
}

/**
 * @brief this is solve multiway_cut problem
 * and return cut_cost
 * example:
 *  \snippet multiway_cut_example.cpp  Multiway Cut Example
 *
 * example file is  multiway_cut_example.cpp
 * @param Graph graph
 * @param int repeats number of sets of radius
 * @param OutputIterator result pairs of vertex descriptor and number form (1,2,
* ... ,k) id of part
 * @tparam Rand random engine
 * @tparam Distribution used to chose random radius
 * @tparam LP
 * @tparam Graph
 * @tparam OutputIterator
 */
template <typename Rand = std::default_random_engine,
          typename Distribution = std::uniform_real_distribution<double>,
          typename LP = lp::glp, typename Graph, class OutputIterator>
auto multiway_cut(const Graph &graph, OutputIterator result,
                  Rand random_engine = std::default_random_engine(5426u))
    ->detail::CostType<Graph> {
    return multiway_cut(graph, result, boost::no_named_parameters(),
                        random_engine);
}

}      //!paal
#endif // PAAL_MULTIWAY_CUT_HPP