//!
//! @file       TestHEmatrix.cpp, cpp file
//! @brief      defining functions for testing homomorphic matrix computation
//!
//! @author     Miran Kim
//! @date       Dec. 1, 2017
//! @copyright  GNU Pub License
//!

#include <cmath>
#include <iostream>
#include <stdio.h>
#include <vector>
#include <sys/time.h>
#include <chrono>

#include <NTL/RR.h>
#include <NTL/xdouble.h>
#include <NTL/ZZ.h>
#include <NTL/ZZX.h>
#include "NTL/RR.h"
#include "NTL/vec_RR.h"
#include "NTL/mat_RR.h"
#include <NTL/BasicThreadPool.h>
#include <NTL/mat_ZZ.h>

#include "../src/Context.h"
#include "../src/Scheme.h"
#include "../src/SecretKey.h"
#include "../src/TimeUtils.h"

#include "matrix.h"
#include "HEmatrix.h"
#include "TestHEmatrix.h"

using namespace std;
using namespace chrono;

void TestHEmatrix::testHEAAN(long logN, long logQ) {

    long pBits = 25;   // scaling factor for message
    long cBits = 15;   // scaling factor for constant

    long h = 64;

    long ntrial = 10;
    /*---------------------------------------*/
    //  Key Generation
    /*---------------------------------------*/

    TimeUtils timeutils;
    timeutils.start("Scheme generating...");
    Context context(logN, logQ);
    SecretKey secretKey(logN, h);
    Scheme scheme(secretKey, context);
    scheme.addLeftRotKeys(secretKey);
    scheme.addRightRotKeys(secretKey);
    timeutils.stop("Scheme generation");

    /*---------------------------------------*/
    //  Encryption
    /*---------------------------------------*/

    auto start= chrono::steady_clock::now();

    long nslots = (1<<(logN-1));

    double* msg1 = new double[nslots];
    double* msg2 = new double[nslots];

    complex<double>* cmsg1 = new complex<double>[nslots];
    complex<double>* cmsg2 = new complex<double>[nslots];

    double* dmsg1 = new double[nslots];
    double* dmsg2 = new double[nslots];

    for(long i = 0; i < nslots; ++i){
        msg1[i] = 1.0/(i+1);
        msg2[i] = 1.0/(2*i+1);

        cmsg1[i].real(msg1[i]);
        cmsg2[i].real(msg2[i]);
    }

    Ciphertext ct1 = scheme.encrypt(cmsg1, nslots, pBits, logQ);
    Ciphertext ct2 = scheme.encrypt(cmsg2, nslots, pBits, logQ);

    auto end = std::chrono::steady_clock::now();
    auto diff = end - start;
    double timeElapsed = chrono::duration <double, milli> (diff).count()/1000.0;
    cout << "Enc time= " << timeElapsed/2 << " s" << endl;
    cout << "------------------" << endl;

    /*---------------------------------------*/
    //  Rotation
    /*---------------------------------------*/
    long nshift = 1;
    Ciphertext ctrot;

    start= chrono::steady_clock::now();
    for(long i = 0; i < ntrial; ++i){
        ctrot = ct1;
        scheme.leftRotateAndEqual(ctrot, nshift);
    }

    end = std::chrono::steady_clock::now();
    diff = end - start;
    timeElapsed = chrono::duration <double, milli> (diff).count()/1000.0;
    cout << "Rot time= " << timeElapsed/ntrial << " s" << endl;
    cout << "------------------" << endl;

    /*---------------------------------------*/
    //  Decryption
    /*---------------------------------------*/

    cmsg1 = scheme.decrypt(secretKey, ctrot);
    for(long i = 0; i < nslots; ++i){
        dmsg1[i] = cmsg1[i].real();
    }

    double* msgrot = new double[nslots];
    long k = nslots - nshift;
    for(long j = 0; j < k; ++j){
        msgrot[j] = msg1[j + nshift];
    }
    for(long j = k; j < nslots; ++j){
        msgrot[j] = msg1[j - k];
    }

    for(long i = 0 ; i < 20; ++i){
        cout << i << ": " << msgrot[i] << ", " << dmsg1[i] << endl;
    }
     cout << "------------------" << endl;

    /*---------------------------------------*/
    //  Multiplication
    /*---------------------------------------*/

    Ciphertext ctmult;

    start= chrono::steady_clock::now();
    for(long i = 0; i < ntrial; ++i){
        ctmult = ct1;
        scheme.multAndEqual(ctmult, ct2);
    }

    end = std::chrono::steady_clock::now();
    diff = end - start;
    timeElapsed = chrono::duration <double, milli> (diff).count()/1000.0;
    cout << "Mult time= " << timeElapsed/ntrial << " s" << endl;
    cout << "------------------" << endl;

    scheme.reScaleByAndEqual(ctmult, pBits);

    /*---------------------------------------*/
    //  Decryption
    /*---------------------------------------*/

    cmsg1 = scheme.decrypt(secretKey, ctmult);
    for(long i = 0; i < nslots; ++i){
        dmsg1[i] = cmsg1[i].real();
    }

    double* msgmult = new double[nslots];
    for(long i = 0; i < nslots; ++i){
        msgmult[i] = msg1[i] * msg2[i];
    }

    for(long i = 0 ; i < 20; ++i){
        cout << i << ": " << msgmult[i] << ", " << dmsg1[i] << endl;
    }
     cout << "------------------" << endl;
}

void TestHEmatrix::testEnc(long nrows) {

    long logN = 13;

    long ncols = nrows;
    long pBits = 25;   // scaling factor for message
    long cBits = 15;   // scaling factor for constant
    long lBits = pBits + 5;
    long logQ = lBits;

    long h = 64;
    long keydist = 1;

    struct HEMatpar HEmatpar;
    readHEMatpar(HEmatpar, nrows, ncols, pBits, cBits, logQ);

    cout << "------------------" << endl;
    cout << "(dim,po2dim, nslots) = (" << nrows << "," << HEmatpar.dim  << "," << HEmatpar.sqrdim << "," << HEmatpar.nslots << ")" << endl;
    cout << "HEAAN PARAMETER logQ: " << logQ << endl;
    cout << "HEAAN PARAMETER logN: " << logN << endl;

    /*---------------------------------------*/
    //  Key Generation
    /*---------------------------------------*/

    TimeUtils timeutils;
    timeutils.start("Scheme generating...");
    Context context(logN, logQ);
    SecretKey secretKey(logN, h);
    Scheme scheme(secretKey, context);

    scheme.addLeftRotKeys(secretKey);
    scheme.addRightRotKeys(secretKey);
    timeutils.stop("Scheme generation");

    HEmatrix HEmatrix(scheme, secretKey, HEmatpar);

    /*---------------------------------------*/
    //  Generate a random matrix
    /*---------------------------------------*/
    Mat<RR> Amat;
    Amat.SetDims(nrows, ncols);

    for(long i = 0; i < nrows ; i++){
        for(long j = 0; j < ncols; j++){
            Amat[i][j]= to_RR((((i* 2 + ncols * j)%3)/10.0));
        }
    }

    /*---------------------------------------*/
    //  Encryption
    /*---------------------------------------*/

    auto start= chrono::steady_clock::now();
    Ciphertext Actxt;
    HEmatrix.encryptRmat(Actxt, Amat, HEmatpar.pBits);

    auto end = std::chrono::steady_clock::now();
    auto diff = end - start;
    double timeElapsed = chrono::duration <double, milli> (diff).count()/1000.0;
    cout << "Enc time= " << timeElapsed << " s" << endl;
    cout << "------------------" << endl;

    /*---------------------------------------*/
    //  Decryption
    /*---------------------------------------*/
    Mat<RR> resmat;
    HEmatrix.decryptRmat(resmat, Actxt);

    /*---------------------------------------*/
    //  Error
    /*---------------------------------------*/

    RR error = getError(Amat, resmat, nrows, ncols);
    cout << "Error (enc): " << error << endl;

    cout << "------------------" << endl;
    cout << "Plaintext" << endl;
    printRmatrix(Amat, 4);
    cout << "------------------" << endl;
    cout << "Encryption" << endl;
    printRmatrix(resmat, 4);
    cout << "------------------" << endl;
}

void TestHEmatrix::testAdd(long nrows) {

    long logN = 13;

    long ncols = nrows;
    long pBits = 25;   // scaling factor for message
    long cBits = 15;   // scaling factor for constant
    long lBits = pBits + 5;
    long logQ = lBits;

    long h = 64;
    long keydist = 1;

    struct HEMatpar HEmatpar;
    readHEMatpar(HEmatpar, nrows, ncols, pBits, cBits, logQ);

    cout << "------------------" << endl;
    cout << "(dim,po2dim, nslots) = (" << nrows << "," << HEmatpar.dim  << "," << HEmatpar.sqrdim << "," << HEmatpar.nslots << ")" << endl;
    cout << "HEAAN PARAMETER logQ: " << logQ << endl;
    cout << "HEAAN PARAMETER logN: " << logN << endl;

    /*---------------------------------------*/
    //  Key Generation
    /*---------------------------------------*/

    TimeUtils timeutils;
    timeutils.start("Scheme generating...");
    Context context(logN, logQ);
    SecretKey secretKey(logN, h);
    Scheme scheme(secretKey, context);

    scheme.addLeftRotKeys(secretKey);
    scheme.addRightRotKeys(secretKey);
    timeutils.stop("Scheme generation");

    HEmatrix HEmatrix(scheme, secretKey, HEmatpar);

    /*---------------------------------------*/
    //  Generate a random matrix
    /*---------------------------------------*/
    Mat<RR> Amat;
    Mat<RR> Bmat;

    Amat.SetDims(nrows, ncols);
    Bmat.SetDims(nrows, ncols);

    for(long i = 0; i < nrows ; i++){
        for(long j = 0; j < ncols; j++){
            Amat[i][j]= to_RR((((i* 2 + ncols * j) % 3) / 1.0));
            Bmat[i][j]= to_RR((((i*ncols + j ) % 3) / 1.0));
        }
    }

    /*---------------------------------------*/
    //  Encryption
    /*---------------------------------------*/

    auto start = chrono::steady_clock::now();
    Ciphertext Actxt;
    HEmatrix.encryptRmat(Actxt, Amat, HEmatpar.pBits);

    Ciphertext Bctxt;
    HEmatrix.encryptRmat(Bctxt, Bmat, HEmatpar.pBits);

    auto end = std::chrono::steady_clock::now();
    auto diff = end - start;
    double timeElapsed = chrono::duration <double, milli> (diff).count()/1000.0;
    cout << "Enc time= " << timeElapsed << " s" << endl;
    cout << "------------------" << endl;

    /*---------------------------------------*/
    //  Addition
    /*---------------------------------------*/

    start = chrono::steady_clock::now();
    scheme.addAndEqual(Actxt, Bctxt);

    end = std::chrono::steady_clock::now();
    diff = end - start;
    timeElapsed = chrono::duration <double, milli> (diff).count()/1000.0;
    cout << "Add time= " << timeElapsed << " s" << endl;

    /*---------------------------------------*/
    //  Decryption
    /*---------------------------------------*/
    Mat<RR> HEresmat;
    HEmatrix.decryptRmat(HEresmat, Actxt);

    /*---------------------------------------*/
    //  Error
    /*---------------------------------------*/
    Mat<RR> resmat;
    add(resmat, Amat, Bmat);

    RR error = getError(resmat, HEresmat, nrows, ncols);
    cout << "Error (add): " << error << endl;

    cout << "------------------" << endl;
    cout << "Plaintext" << endl;
    printRmatrix(resmat, 4);
    cout << "------------------" << endl;
    cout << "Encryption" << endl;
    printRmatrix(HEresmat, 4);
    cout << "------------------" << endl;
}

void TestHEmatrix::testTrans(long nrows) {

    long logN = 13;

    long ncols = nrows;
    long pBits = 35;   // scaling factor for message
    long cBits = 15;   // scaling factor for constant
    long lBits = pBits + 5;
    long logQ = lBits + cBits;

    long h = 64;
    long keydist = 1;

    struct HEMatpar HEmatpar;
    readHEMatpar(HEmatpar, nrows, ncols, pBits, cBits, logQ);

    cout << "------------------" << endl;
    cout << "(dim,po2dim, nslots) = (" << nrows << "," << HEmatpar.dim  << "," << HEmatpar.sqrdim << "," << HEmatpar.nslots << ")" << endl;
    cout << "HEAAN PARAMETER logQ: " << logQ << endl;
    cout << "HEAAN PARAMETER logN: " << logN << endl;

    /*---------------------------------------*/
    //  Key Generation
    /*---------------------------------------*/

    TimeUtils timeutils;
    timeutils.start("Scheme generating...");
    Context context(logN, logQ);
    SecretKey secretKey(logN, h);
    Scheme scheme(secretKey, context);

    scheme.addLeftRotKeys(secretKey);
    scheme.addRightRotKeys(secretKey);
    timeutils.stop("Scheme generation");

    HEmatrix HEmatrix(scheme, secretKey, HEmatpar);

    /*---------------------------------------*/
    //  Generate a random matrix
    /*---------------------------------------*/
    Mat<RR> Amat;
    Amat.SetDims(nrows, ncols);

    for(long i = 0; i < nrows ; i++){
        for(long j = 0; j < ncols; j++){
            Amat[i][j]= to_RR((((i* 2 + ncols * j)%3)/10.0));
        }
        Amat[i][i] += to_RR("2.0");
    }

    /*---------------------------------------*/
    //  Encryption
    /*---------------------------------------*/

    auto start= chrono::steady_clock::now();
    Ciphertext Actxt;
    HEmatrix.encryptRmat(Actxt, Amat, HEmatpar.pBits);

    auto end = std::chrono::steady_clock::now();
    auto diff = end - start;
    double timeElapsed = chrono::duration <double, milli> (diff).count()/1000.0;
    cout << "Enc time= " << timeElapsed << " s" << endl;
    cout << "------------------" << endl;

    /*---------------------------------------*/
    //  Transposition
    /*---------------------------------------*/

    ZZX* transpoly;
    HEmatrix.genTransPoly(transpoly);

    Ciphertext Tctxt;

    start= chrono::steady_clock::now();

    HEmatrix.transpose(Tctxt, Actxt, transpoly);

    end = std::chrono::steady_clock::now();
    diff = end - start;
    timeElapsed = chrono::duration <double, milli> (diff).count()/1000.0;
    cout << "Trans time= " << timeElapsed << " s" << endl;

    /*---------------------------------------*/
    //  Decryption
    /*---------------------------------------*/
    Mat<RR> HEresmat;
    HEmatrix.decryptRmat(HEresmat, Tctxt);

    /*---------------------------------------*/
    //  Error
    /*---------------------------------------*/
    Mat<RR> resmat;
    transpose(resmat, Amat);

    RR error = getError(resmat, HEresmat, nrows, ncols);
    cout << "Error (trans): " << error << endl;

    cout << "------------------" << endl;
    cout << "Plaintext" << endl;
    printRmatrix(resmat, 4);
    cout << "------------------" << endl;
    cout << "Encryption" << endl;
    printRmatrix(HEresmat, 4);
    cout << "------------------" << endl;
}

void TestHEmatrix::testShift(long nrows, long k) {

    long logN = 13;

    long ncols = nrows;
    long pBits = 25;   // scaling factor for message
    long cBits = 15;   // scaling factor for constant
    long lBits = pBits + 5;
    long logQ = lBits + cBits;

    long h = 64;
    long keydist = 1;

    struct HEMatpar HEmatpar;
    readHEMatpar(HEmatpar, nrows, ncols, pBits, cBits, logQ);

    cout << "------------------" << endl;
    cout << "(dim,po2dim, nslots) = (" << nrows << "," << HEmatpar.dim  << "," << HEmatpar.sqrdim << "," << HEmatpar.nslots << ")" << endl;
    cout << "HEAAN PARAMETER logQ: " << logQ << endl;
    cout << "HEAAN PARAMETER logN: " << logN << endl;

    /*---------------------------------------*/
    //  Key Generation
    /*---------------------------------------*/

    TimeUtils timeutils;
    timeutils.start("Scheme generating...");
    Context context(logN, logQ);
    SecretKey secretKey(logN, h);
    Scheme scheme(secretKey, context);

    scheme.addLeftRotKeys(secretKey);
    scheme.addRightRotKeys(secretKey);
    timeutils.stop("Scheme generation");

    HEmatrix HEmatrix(scheme, secretKey, HEmatpar);

    /*---------------------------------------*/
    //  Generate a random matrix
    /*---------------------------------------*/
    Mat<RR> Amat;
    Amat.SetDims(nrows, ncols);

    for(long i = 0; i < nrows ; i++){
        for(long j = 0; j < ncols; j++){
            Amat[i][j]= to_RR((((i* 2 + ncols * j)%3)/10.0));
        }
    }

    /*---------------------------------------*/
    //  Encryption
    /*---------------------------------------*/

    auto start= chrono::steady_clock::now();
    Ciphertext Actxt;
    HEmatrix.encryptRmat(Actxt, Amat, HEmatpar.pBits);

    auto end = std::chrono::steady_clock::now();
    auto diff = end - start;
    double timeElapsed = chrono::duration <double, milli> (diff).count()/1000.0;
    cout << "Enc time= " << timeElapsed << " s" << endl;
    cout << "------------------" << endl;

    /*---------------------------------------*/
    //  Transposition
    /*---------------------------------------*/

    ZZX* shiftpoly;
    HEmatrix.genShiftPoly(shiftpoly);

    Ciphertext Sctxt;

    start= chrono::steady_clock::now();

    HEmatrix.shiftBycols(Sctxt, Actxt, k, shiftpoly);

    end = std::chrono::steady_clock::now();
    diff = end - start;
    timeElapsed = chrono::duration <double, milli> (diff).count()/1000.0;
    cout << "Shift by" << k << " time= " << timeElapsed << " s" << endl;

    /*---------------------------------------*/
    //  Decryption
    /*---------------------------------------*/
    Mat<RR> HEresmat;
    HEmatrix.decryptRmat(HEresmat, Sctxt);

    /*---------------------------------------*/
    //  Error
    /*---------------------------------------*/
    cout << "------------------" << endl;
    cout << "Plaintext (before shifting)" << endl;
    printRmatrix(Amat, nrows);
    cout << "------------------" << endl;

    cout << "------------------" << endl;
    cout << "Encryption" << endl;
    printRmatrix(HEresmat, nrows);
    cout << "------------------" << endl;
}

void TestHEmatrix::testMult(long nrows) {

    long logN = 13;

    long ncols = nrows;
    long pBits = 25;   // scaling factor for message
    long cBits = 14;   // scaling factor for constant
    long lBits = pBits + 8;
    long logQ = (2*cBits) + pBits + lBits;

    long h = 64;
    long keydist = 1;

    struct HEMatpar HEmatpar;
    readHEMatpar(HEmatpar, nrows, ncols, pBits, cBits, logQ);

    cout << "------------------" << endl;
    cout << "(dim,po2dim, nslots) = (" << nrows << "," << HEmatpar.dim  << "," << HEmatpar.sqrdim << "," << HEmatpar.nslots << ")" << endl;
    cout << "HEAAN PARAMETER logQ: " << logQ << endl;
    cout << "HEAAN PARAMETER logN: " << logN << endl;

    /*---------------------------------------*/
    //  Key Generation
    /*---------------------------------------*/

    TimeUtils timeutils;
    timeutils.start("Scheme generating...");
    Context context(logN, logQ);
    SecretKey secretKey(logN, h);
    Scheme scheme(secretKey, context);

    scheme.addLeftRotKeys(secretKey);
    scheme.addRightRotKeys(secretKey);
    timeutils.stop("Scheme generation");

    HEmatrix HEmatrix(scheme, secretKey, HEmatpar);

    /*---------------------------------------*/
    //  Generate a random matrix
    /*---------------------------------------*/
    Mat<RR> Amat;
    Mat<RR> Bmat;

    Amat.SetDims(nrows, ncols);
    Bmat.SetDims(nrows, ncols);

    for(long i = 0; i < nrows ; i++){
        for(long j = 0; j < ncols; j++){
            Amat[i][j]= to_RR((((i* 2 + ncols * j)%3) / 1.0));
            Bmat[i][j]= to_RR((((i*ncols + j )%3) / 1.0));
        }
    }
    cout << "Amat: " << Amat << endl;
    cout << "Bmat: " << Bmat << endl;
    /*---------------------------------------*/
    //  Encryption
    /*---------------------------------------*/

    auto start= chrono::steady_clock::now();
    Ciphertext Actxt;
    HEmatrix.encryptRmat(Actxt, Amat, HEmatpar.pBits);

    Ciphertext Bctxt;
    HEmatrix.encryptRmat(Bctxt, Bmat, HEmatpar.pBits);

    auto end = std::chrono::steady_clock::now();
    auto diff = end - start;
    double timeElapsed = chrono::duration <double, milli> (diff).count()/1000.0;
    cout << "Enc time= " << timeElapsed << " s" << endl;
    cout << "------------------" << endl;

    /*---------------------------------------*/
    //  GenPoly
    /*---------------------------------------*/

    start= chrono::steady_clock::now();

    ZZX** Initpoly;
    HEmatrix.genMultPoly(Initpoly);

    ZZX* shiftpoly;
    HEmatrix.genShiftPoly(shiftpoly);

    end = std::chrono::steady_clock::now();
    diff = end - start;
    timeElapsed = chrono::duration <double, milli> (diff).count()/1000.0;
    cout << "GenPoly time= " << timeElapsed << " s" << endl;
    cout << "---------------------------" << endl;

    /*---------------------------------------*/
    //  Mult
    /*---------------------------------------*/

    start= chrono::steady_clock::now();
    Ciphertext Cctxt;
    HEmatrix.HEmatmul(Cctxt, Actxt, Bctxt, Initpoly, shiftpoly);

    end = std::chrono::steady_clock::now();
    diff = end - start;
    timeElapsed = chrono::duration <double, milli> (diff).count()/1000.0;
    cout << "Mult time= " << timeElapsed << " s" << endl;

    /*---------------------------------------*/
    //  Decryption
    /*---------------------------------------*/
    Mat<RR> HEresmat;
    HEmatrix.decryptRmat(HEresmat, Cctxt);

    /*---------------------------------------*/
    //  Error
    /*---------------------------------------*/
    Mat<RR> resmat;
    mul(resmat, Amat, Bmat);

    RR error = getError(resmat, HEresmat, nrows, ncols);
    cout << "Error (mul): " << error << endl;

    cout << "------------------" << endl;
    cout << "Plaintext" << endl;
    printRmatrix(resmat, 4);
    cout << "------------------" << endl;
    cout << "Encryption" << endl;
    printRmatrix(HEresmat, 4);
    cout << "------------------" << endl;
}

void TestHEmatrix::testRMult(long nrows, long subdim) {

    long logN = 13;

    long ncols = nrows;
    long pBits = 25;   // scaling factor for message
    long cBits = 14;   // scaling factor for constant
    long lBits = pBits + 8;
    long logQ = (2*cBits) + pBits + lBits;

    long h = 64;
    long keydist = 1;

    struct HEMatpar HEmatpar;
    readHEMatpar(HEmatpar, nrows, ncols, pBits, cBits, logQ, subdim);

    cout << "------------------" << endl;
    cout << "(dim,po2dim, nslots) = (" << nrows << "," << HEmatpar.dim  << "," << HEmatpar.sqrdim << "," << HEmatpar.nslots << ")" << endl;
    cout << "HEAAN PARAMETER logQ: " << logQ << endl;
    cout << "HEAAN PARAMETER logN: " << logN << endl;

    /*---------------------------------------*/
    //  Key Generation
    /*---------------------------------------*/

    TimeUtils timeutils;
    timeutils.start("Scheme generating...");
    Context context(logN, logQ);
    SecretKey secretKey(logN, h);
    Scheme scheme(secretKey, context);

    scheme.addLeftRotKeys(secretKey);
    scheme.addRightRotKeys(secretKey);
    timeutils.stop("Scheme generation");

    HEmatrix HEmatrix(scheme, secretKey, HEmatpar);

    /*---------------------------------------*/
    //  Generate a random matrix
    /*---------------------------------------*/
    Mat<RR> rAmat;
    rAmat.SetDims(subdim, ncols);

    Mat<RR> Bmat;
    Bmat.SetDims(nrows, ncols);

    for(long i = 0; i < subdim; i++){
        for(long j = 0; j < ncols; j++){
            rAmat[i][j]= to_RR((((i + ncols * j)%3)/10.0));
        }
    }

    for(long i = 0; i < nrows ; i++){
        for(long j = 0; j < ncols; j++){
            Bmat[i][j]= to_RR((((i*ncols + j )%3)/10.0));
        }
    }

    // replicate
    Mat<RR> rAmat1;
    rAmat1.SetDims(nrows, ncols);
    for(long i = 0; i < nrows/subdim; ++i){
        for(long j = 0; j < subdim; ++j){
            for(long k = 0; k< ncols; k++){
                rAmat1[i*subdim + j][k] = rAmat[j][k];
            }
        }
    }

    /*---------------------------------------*/
    //  Encryption
    /*---------------------------------------*/

    auto start= chrono::steady_clock::now();
    Ciphertext Actxt;
    HEmatrix.encryptRmat(Actxt, rAmat1, HEmatpar.pBits);

    Ciphertext Bctxt;
    HEmatrix.encryptRmat(Bctxt, Bmat, HEmatpar.pBits);

    auto end = std::chrono::steady_clock::now();
    auto diff = end - start;
    double timeElapsed = chrono::duration <double, milli> (diff).count()/1000.0;
    cout << "Enc time= " << timeElapsed << " s" << endl;
    cout << "------------------" << endl;

    /*---------------------------------------*/
    //  GenPoly
    /*---------------------------------------*/

    start= chrono::steady_clock::now();

    ZZX** Initpoly;
    HEmatrix.genMultPoly(Initpoly);

    ZZX* shiftpoly;
    HEmatrix.genShiftPoly(shiftpoly);

    end = std::chrono::steady_clock::now();
    diff = end - start;
    timeElapsed = chrono::duration <double, milli> (diff).count()/1000.0;
    cout << "GenPoly time= " << timeElapsed << " s" << endl;
    cout << "---------------------------" << endl;

    /*---------------------------------------*/
    //  Mult
    /*---------------------------------------*/

    start= chrono::steady_clock::now();
    Ciphertext Cctxt;
    HEmatrix.HErmatmul(Cctxt, Actxt, Bctxt, Initpoly, shiftpoly);

    end = std::chrono::steady_clock::now();
    diff = end - start;
    timeElapsed = chrono::duration <double, milli> (diff).count()/1000.0;
    cout << "R-Mult time= " << timeElapsed << " s" << endl;

    /*---------------------------------------*/
    //  Decryption
    /*---------------------------------------*/
    Mat<RR> HEresmat;
    HEmatrix.decryptRmat(HEresmat, Cctxt);

    /*---------------------------------------*/
    //  Error
    /*---------------------------------------*/
    Mat<RR> resmat;
    mul(resmat, rAmat, Bmat);

    RR error = getError(resmat, HEresmat, subdim, ncols);
    cout << "Error (mul): " << error << endl;

    cout << "------------------" << endl;
    cout << "Plaintext" << endl;
    printRmatrix(resmat, 4);
    cout << "------------------" << endl;
    cout << "Encryption" << endl;
    printRmatrix(HEresmat, 4);
    cout << "------------------" << endl;
}

// If Amat is given in plaintext

void TestHEmatrix::testMult_preprocessing(long nrows) {

    long logN = 13;

    long ncols = nrows;
    long pBits = 25;   // scaling factor for message
    long cBits = 14;   // scaling factor for constant
    long lBits = pBits + 8;
    long logQ = cBits + pBits + lBits;

    long h = 64;
    long keydist = 1;

    struct HEMatpar HEmatpar;
    readHEMatpar(HEmatpar, nrows, ncols, pBits, cBits, logQ);

    cout << "------------------" << endl;
    cout << "(dim,po2dim, nslots) = (" << nrows << "," << HEmatpar.dim  << "," << HEmatpar.sqrdim << "," << HEmatpar.nslots << ")" << endl;
    cout << "HEAAN PARAMETER logQ: " << logQ << endl;
    cout << "HEAAN PARAMETER logN: " << logN << endl;

    /*---------------------------------------*/
    //  Key Generation
    /*---------------------------------------*/

    TimeUtils timeutils;
    timeutils.start("Scheme generating...");
    Context context(logN, logQ);
    SecretKey secretKey(logN, h);
    Scheme scheme(secretKey, context);

    scheme.addLeftRotKeys(secretKey);
    scheme.addRightRotKeys(secretKey);
    timeutils.stop("Scheme generation");

    HEmatrix HEmatrix(scheme, secretKey, HEmatpar);

    /*---------------------------------------*/
    //  Generate a random matrix
    /*---------------------------------------*/
    Mat<RR> Amat;
    Mat<RR> Bmat;

    Amat.SetDims(nrows, ncols);
    Bmat.SetDims(nrows, ncols);

    for(long i = 0; i < nrows ; i++){
        for(long j = 0; j < ncols; j++){
            Amat[i][j]= to_RR((((i* 2 + ncols * j)%3)/10.0));
            Bmat[i][j]= to_RR((((i*ncols + j )%3)/10.0));
        }
    }

    /*---------------------------------------*/
    //  Encryption
    /*---------------------------------------*/

    auto start= chrono::steady_clock::now();
    Ciphertext* Actxts;
    HEmatrix.genInitActxt(Actxts, Amat);

    Ciphertext Bctxt;
    HEmatrix.encryptRmat(Bctxt, Bmat, HEmatpar.pBits);

    auto end = std::chrono::steady_clock::now();
    auto diff = end - start;
    double timeElapsed = chrono::duration <double, milli> (diff).count()/1000.0;
    cout << "Enc time= " << timeElapsed << " s" << endl;
    cout << "------------------" << endl;

    /*---------------------------------------*/
    //  GenPoly
    /*---------------------------------------*/

    start= chrono::steady_clock::now();

    ZZX* Initpoly;
    HEmatrix.genMultBPoly(Initpoly);

    end = std::chrono::steady_clock::now();
    diff = end - start;
    timeElapsed = chrono::duration <double, milli> (diff).count()/1000.0;
    cout << "GenPoly time= " << timeElapsed << " s" << endl;
    cout << "---------------------------" << endl;

    /*---------------------------------------*/
    //  Mult
    /*---------------------------------------*/

    start= chrono::steady_clock::now();
    Ciphertext Cctxt;
    HEmatrix.HEmatmul_preprocessing(Cctxt, Actxts, Bctxt, Initpoly);

    end = std::chrono::steady_clock::now();
    diff = end - start;
    timeElapsed = chrono::duration <double, milli> (diff).count()/1000.0;
    cout << "Mult time (preprocessing)= " << timeElapsed << " s" << endl;

    /*---------------------------------------*/
    //  Decryption
    /*---------------------------------------*/
    Mat<RR> HEresmat;
    HEmatrix.decryptRmat(HEresmat, Cctxt);

    /*---------------------------------------*/
    //  Error
    /*---------------------------------------*/
    Mat<RR> resmat;
    mul(resmat, Amat, Bmat);

    RR error = getError(resmat, HEresmat, nrows, ncols);
    cout << "Error (mul): " << error << endl;

    cout << "------------------" << endl;
    cout << "Plaintext" << endl;
    printRmatrix(resmat, 4);
    cout << "------------------" << endl;
    cout << "Encryption" << endl;
    printRmatrix(HEresmat, 4);
    cout << "------------------" << endl;
}

void TestHEmatrix::testRMult_preprocessing(long nrows, long subdim) {

    long logN = 13;

    long ncols = nrows;
    long pBits = 25;   // scaling factor for message
    long cBits = 14;   // scaling factor for constant
    long lBits = pBits + 8;
    long logQ = (cBits) + pBits + lBits;

    long h = 64;
    long keydist = 1;

    struct HEMatpar HEmatpar;
    readHEMatpar(HEmatpar, nrows, ncols, pBits, cBits, logQ, subdim);

    cout << "------------------" << endl;
    cout << "(dim,po2dim, nslots) = (" << nrows << "," << HEmatpar.dim  << "," << HEmatpar.sqrdim << "," << HEmatpar.nslots << ")" << endl;
    cout << "HEAAN PARAMETER logQ: " << logQ << endl;
    cout << "HEAAN PARAMETER logN: " << logN << endl;

    /*---------------------------------------*/
    //  Key Generation
    /*---------------------------------------*/

    TimeUtils timeutils;
    timeutils.start("Scheme generating...");
    Context context(logN, logQ);
    SecretKey secretKey(logN, h);
    Scheme scheme(secretKey, context);

    scheme.addLeftRotKeys(secretKey);
    scheme.addRightRotKeys(secretKey);
    timeutils.stop("Scheme generation");

    HEmatrix HEmatrix(scheme, secretKey, HEmatpar);

    /*---------------------------------------*/
    //  Generate a random matrix
    /*---------------------------------------*/
    Mat<RR> rAmat;
    rAmat.SetDims(subdim, ncols);

    Mat<RR> Bmat;
    Bmat.SetDims(nrows, ncols);

    for(long i = 0; i < subdim; i++){
        for(long j = 0; j < ncols; j++){
            rAmat[i][j]= to_RR((((i + ncols * j)%3)/10.0));
        }
    }

    for(long i = 0; i < nrows ; i++){
        for(long j = 0; j < ncols; j++){
            Bmat[i][j]= to_RR((((i*ncols + j )%3)/10.0));
        }
    }

    /*---------------------------------------*/
    //  Encryption
    /*---------------------------------------*/

    auto start= chrono::steady_clock::now();
    Ciphertext* Actxts;
    HEmatrix.genInitRecActxt(Actxts, rAmat);

    Ciphertext Bctxt;
    HEmatrix.encryptRmat(Bctxt, Bmat, HEmatpar.pBits);

    auto end = std::chrono::steady_clock::now();
    auto diff = end - start;
    double timeElapsed = chrono::duration <double, milli> (diff).count()/1000.0;
    cout << "Enc time= " << timeElapsed << " s" << endl;
    cout << "------------------" << endl;

    /*---------------------------------------*/
    //  GenPoly
    /*---------------------------------------*/

    start= chrono::steady_clock::now();

    ZZX* Initpoly;
    HEmatrix.genMultBPoly(Initpoly);

    end = std::chrono::steady_clock::now();
    diff = end - start;
    timeElapsed = chrono::duration <double, milli> (diff).count()/1000.0;
    cout << "GenPoly time= " << timeElapsed << " s" << endl;
    cout << "---------------------------" << endl;

    /*---------------------------------------*/
    //  Mult
    /*---------------------------------------*/

    start= chrono::steady_clock::now();
    Ciphertext Cctxt;
    HEmatrix.HErmatmul_preprocessing(Cctxt, Actxts, Bctxt, Initpoly);

    end = std::chrono::steady_clock::now();
    diff = end - start;
    timeElapsed = chrono::duration <double, milli> (diff).count()/1000.0;
    cout << "R-Mult time= " << timeElapsed << " s" << endl;

    /*---------------------------------------*/
    //  Decryption
    /*---------------------------------------*/
    Mat<RR> HEresmat;
    HEmatrix.decryptRmat(HEresmat, Cctxt);

    /*---------------------------------------*/
    //  Error
    /*---------------------------------------*/
    Mat<RR> resmat;
    mul(resmat, rAmat, Bmat);

    RR error = getError(resmat, HEresmat, subdim, ncols);
    cout << "Error (mul): " << error << endl;

    cout << "------------------" << endl;
    cout << "Plaintext" << endl;
    printRmatrix(resmat, 4);
    cout << "------------------" << endl;
    cout << "Encryption" << endl;
    printRmatrix(HEresmat, 4);
    cout << "------------------" << endl;
}

//----------------------------------------------------
// SIMD (parallel computation)

void TestHEmatrix::testSIMDAdd(long nrows, long nbatching, const long niter) {

    long logN = 13;

    long ncols = nrows;

    if(nrows * ncols * nbatching > (1<< (logN-1))){
        cout << "Cannot support the parallism " << endl;
    }

    long pBits = 25;   // scaling factor for message
    long cBits = 15;   // scaling factor for constant
    long lBits = pBits + 5;
    long logQ = lBits;

    long h = 64;
    long keydist = 1;

    struct HEMatpar HEmatpar;
    readHEMatpar(HEmatpar, nrows, ncols, pBits, cBits, logQ, 0, nbatching);

    cout << "------------------" << endl;
    cout << "(dim, nslots, nbatching) = (" << nrows << ","  << HEmatpar.nslots << ","  << HEmatpar.nbatching << ")" << endl;
    cout << "HEAAN PARAMETER logQ: " << logQ << endl;
    cout << "HEAAN PARAMETER logN: " << logN << endl;

    /*---------------------------------------*/
    //  Key Generation
    /*---------------------------------------*/

    TimeUtils timeutils;
    timeutils.start("Scheme generating...");
    Context context(logN, logQ);
    SecretKey secretKey(logN, h);
    Scheme scheme(secretKey, context);

    scheme.addLeftRotKeys(secretKey);
    scheme.addRightRotKeys(secretKey);
    timeutils.stop("Scheme generation");

    HEmatrix HEmatrix(scheme, secretKey, HEmatpar);

    /*---------------------------------------*/
    //  Generate a random matrix
    /*---------------------------------------*/
    Mat<RR>* Amat = new Mat<RR>[nbatching];
    Mat<RR>* Bmat = new Mat<RR>[nbatching];

    for(long k = 0; k < nbatching; ++k){
        Amat[k].SetDims(nrows, ncols);
        Bmat[k].SetDims(nrows, ncols);

        for(long i = 0; i < nrows ; i++){
            for(long j = 0; j < ncols; j++){
                //Amat[k][i][j] = to_RR((((i*2 + ncols * j + 1)%3)/10.0));
                Amat[k][i][j] = to_RR(((i*ncols + j)%3/10.0));
                Bmat[k][i][j] = to_RR(((i*ncols + j)%3/10.0));
            }
        }
    }

    double totalEnctime = 0.0;
    double totalEvaltime = 0.0;
    double amortizedtime = 0.0;
    double totalDectime = 0.0;

    for(long l = 0; l < niter; ++l){
        /*---------------------------------------*/
        //  Encryption
        /*---------------------------------------*/

        auto start= chrono::steady_clock::now();
        Ciphertext Actxt;
        HEmatrix.encryptParallelRmat(Actxt, Amat, HEmatpar.pBits, nbatching);

        Ciphertext Bctxt;
        HEmatrix.encryptParallelRmat(Bctxt, Bmat, HEmatpar.pBits, nbatching);

        auto end = std::chrono::steady_clock::now();
        auto diff = end - start;
        double timeElapsed = chrono::duration <double, milli> (diff).count()/1000.0;
        totalEnctime += timeElapsed;

        cout << "Encryption" << endl;

        /*---------------------------------------*/
        //  Addition
        /*---------------------------------------*/

        start= chrono::steady_clock::now();

        scheme.addAndEqual(Actxt, Bctxt);

        end = std::chrono::steady_clock::now();
        diff = end - start;
        timeElapsed = chrono::duration <double, milli> (diff).count()/1000.0;

        totalEvaltime += timeElapsed;
        amortizedtime += timeElapsed/nbatching;

        /*---------------------------------------*/
        //  Decryption
        /*---------------------------------------*/
        start= chrono::steady_clock::now();

        Mat<RR>* HEresmat;
        HEmatrix.decryptParallelRmat(HEresmat, Actxt);

        end = std::chrono::steady_clock::now();
        diff = end - start;
        timeElapsed = chrono::duration <double, milli> (diff).count()/1000.0;
        totalDectime += timeElapsed;

        /*---------------------------------------*/
        //  Error
        /*---------------------------------------*/
        RR avgerror_add = to_RR("0");

        Mat<RR>* resmat = new Mat<RR>[nbatching];
        for(long k = 0; k < nbatching; ++k){
            add(resmat[k], Amat[k], Bmat[k]);
            avgerror_add += getError(resmat[k], HEresmat[k], nrows, ncols);
        }
        avgerror_add /= nbatching;

        if(l == niter -1) cout << "Error (add): " << avgerror_add << endl;
    }

    cout << "------------------" << endl;
    cout << "Enc time= " << totalEnctime/niter << " s" << endl;
    cout << "Add time= " << totalEvaltime/niter << " s/ " ;
    cout << "Amortized time= " << amortizedtime/niter << " s" << endl;
    cout << "Dec time= " <<  totalDectime/niter << " s" << endl;

}

void TestHEmatrix::testSIMDTrans(long nrows, long nbatching, const long niter) {

    long logN = 13;

    long ncols = nrows;

    if(nrows * ncols * nbatching > (1<< (logN-1))){
        cout << "Cannot support the parallism " << endl;
    }

    long pBits = 25;   // scaling factor for message
    long cBits = 15;   // scaling factor for constant
    long lBits = pBits + 5;
    long logQ = cBits + lBits;

    long h = 64;
    long keydist = 1;

    struct HEMatpar HEmatpar;
    readHEMatpar(HEmatpar, nrows, ncols, pBits, cBits, logQ, 0, nbatching);

    cout << "------------------" << endl;
    cout << "(dim, nslots, nbatching) = (" << nrows << ","  << HEmatpar.nslots << ","  << HEmatpar.nbatching << ")" << endl;
    cout << "HEAAN PARAMETER logQ: " << logQ << endl;
    cout << "HEAAN PARAMETER logN: " << logN << endl;

    /*---------------------------------------*/
    //  Key Generation
    /*---------------------------------------*/

    TimeUtils timeutils;
    timeutils.start("Scheme generating...");
    Context context(logN, logQ);
    SecretKey secretKey(logN, h);
    Scheme scheme(secretKey, context);

    scheme.addLeftRotKeys(secretKey);
    scheme.addRightRotKeys(secretKey);
    timeutils.stop("Scheme generation");

    HEmatrix HEmatrix(scheme, secretKey, HEmatpar);

    /*---------------------------------------*/
    //  Generate a random matrix
    /*---------------------------------------*/
    Mat<RR>* Amat = new Mat<RR>[nbatching];

    for(long k = 0; k < nbatching; ++k){
        Amat[k].SetDims(nrows, ncols);
        for(long i = 0; i < nrows ; i++){
            for(long j = 0; j < ncols; j++){
                Amat[k][i][j] = to_RR(((i*ncols + j)%3/10.0));
            }
        }
    }

    double totalEnctime = 0.0;
    double totalEvaltime = 0.0;
    double amortizedtime = 0.0;
    double totalDectime = 0.0;

    for(long l = 0; l < niter; ++l){
        /*---------------------------------------*/
        //  Encryption
        /*---------------------------------------*/

        auto start= chrono::steady_clock::now();
        Ciphertext Actxt;
        HEmatrix.encryptParallelRmat(Actxt, Amat, HEmatpar.pBits, nbatching);

        auto end = std::chrono::steady_clock::now();
        auto diff = end - start;
        double timeElapsed = chrono::duration <double, milli> (diff).count()/1000.0;
        totalEnctime += timeElapsed;

        /*---------------------------------------*/
        //  Transpose
        /*---------------------------------------*/
        ZZX* transpoly;
        HEmatrix.genTransPoly_Parallel(transpoly);

        start= chrono::steady_clock::now();
        Ciphertext Tctxt;
        HEmatrix.transpose_Parallel(Tctxt, Actxt, transpoly);

        end = std::chrono::steady_clock::now();
        diff = end - start;
        timeElapsed = chrono::duration <double, milli> (diff).count()/1000.0;

        totalEvaltime += timeElapsed;
        amortizedtime += timeElapsed/nbatching;

        /*---------------------------------------*/
        //  Decryption
        /*---------------------------------------*/
        start= chrono::steady_clock::now();

        Mat<RR>* HEresmat;
        HEmatrix.decryptParallelRmat(HEresmat, Tctxt);

        end = std::chrono::steady_clock::now();
        diff = end - start;
        timeElapsed = chrono::duration <double, milli> (diff).count()/1000.0;
        totalDectime += timeElapsed;

        /*---------------------------------------*/
        //  Error
        /*---------------------------------------*/
        RR avgerror = to_RR("0");
        Mat<RR>* resmat = new Mat<RR>[nbatching];
        for(long k = 0; k < nbatching; ++k){
            transpose(resmat[k], Amat[k]);
            avgerror += getError(HEresmat[k], resmat[k], nrows, ncols);
        }
        avgerror /= nbatching;

        if(l == niter -1) cout << "Error (trans): " << avgerror << endl;
    }

    cout << "------------------" << endl;
    cout << "Enc time= " << totalEnctime/niter << " s" << endl;
    cout << "Trans time= " << totalEvaltime/niter << " s/ " ;
    cout << "Amortized time= " << amortizedtime/niter << " s" << endl;
    cout << "Dec time= " <<  totalDectime/niter << " s" << endl;

}

void TestHEmatrix::testSIMDMult(long nrows, long nbatching, const long niter) {

    long logN = 13;

    long ncols = nrows;

    if(nrows * ncols * nbatching > (1<< (logN-1))){
        cout << "Cannot support the parallism " << endl;
    }

    long pBits = 25;   // scaling factor for message
    long cBits = 15;   // scaling factor for constant
    long lBits = pBits + 5;
    long logQ = (2*cBits + pBits) + lBits;

    long h = 64;
    long keydist = 1;

    struct HEMatpar HEmatpar;
    readHEMatpar(HEmatpar, nrows, ncols, pBits, cBits, logQ, 0, nbatching);

    cout << "------------------" << endl;
    cout << "(dim, nslots, nbatching) = (" << nrows << ","  << HEmatpar.nslots << ","  << HEmatpar.nbatching << ")" << endl;
    cout << "HEAAN PARAMETER logQ: " << logQ << endl;
    cout << "HEAAN PARAMETER logN: " << logN << endl;

    /*---------------------------------------*/
    //  Key Generation
    /*---------------------------------------*/

    TimeUtils timeutils;
    timeutils.start("Scheme generating...");
    Context context(logN, logQ);
    SecretKey secretKey(logN, h);
    Scheme scheme(secretKey, context);

    scheme.addLeftRotKeys(secretKey);
    scheme.addRightRotKeys(secretKey);
    timeutils.stop("Scheme generation");

    HEmatrix HEmatrix(scheme, secretKey, HEmatpar);

    /*---------------------------------------*/
    //  Generate a random matrix
    /*---------------------------------------*/
    Mat<RR>* Amat = new Mat<RR>[nbatching];
    Mat<RR>* Bmat = new Mat<RR>[nbatching];

    for(long k = 0; k < nbatching; ++k){
        Amat[k].SetDims(nrows, ncols);
        Bmat[k].SetDims(nrows, ncols);

        for(long i = 0; i < nrows ; i++){
            for(long j = 0; j < ncols; j++){
                //Amat[k][i][j] = to_RR((((i*2 + ncols * j + 1)%3)/10.0));
                Amat[k][i][j] = to_RR(((i*ncols + j)%3/10.0));
                Bmat[k][i][j] = to_RR(((i*ncols + j)%3/10.0));
            }
        }
    }

    double totalEnctime = 0.0;
    double totalEvaltime = 0.0;
    double amortizedtime = 0.0;
    double totalDectime = 0.0;
    double genpolytime = 0.0;

    for(long l = 0; l < niter; ++l){
        /*---------------------------------------*/
        //  Encryption
        /*---------------------------------------*/

        auto start= chrono::steady_clock::now();
        Ciphertext Actxt;
        HEmatrix.encryptParallelRmat(Actxt, Amat, HEmatpar.pBits, nbatching);

        Ciphertext Bctxt;
        HEmatrix.encryptParallelRmat(Bctxt, Bmat, HEmatpar.pBits, nbatching);

        auto end = std::chrono::steady_clock::now();
        auto diff = end - start;
        double timeElapsed = chrono::duration <double, milli> (diff).count()/1000.0;
        totalEnctime += timeElapsed;

        /*---------------------------------------*/
        //  Genpoly
        /*---------------------------------------*/

        start= chrono::steady_clock::now();

        ZZX** Initpoly;
        HEmatrix.genMultPoly_Parallel(Initpoly);

        ZZX* shiftpoly;
        HEmatrix.genShiftPoly_Parallel(shiftpoly);

        end = std::chrono::steady_clock::now();
        diff = end - start;
        timeElapsed = chrono::duration <double, milli> (diff).count()/1000.0;
        genpolytime += timeElapsed;

        /*---------------------------------------*/
        //  Mult
        /*---------------------------------------*/

        start= chrono::steady_clock::now();

        Ciphertext Cctxt;
        HEmatrix.HEmatmul_Parallel(Cctxt, Actxt, Bctxt, Initpoly, shiftpoly);

        end = std::chrono::steady_clock::now();
        diff = end - start;
        timeElapsed = chrono::duration <double, milli> (diff).count()/1000.0;

        totalEvaltime += timeElapsed;
        amortizedtime += timeElapsed/nbatching;

        /*---------------------------------------*/
        //  Decryption
        /*---------------------------------------*/
        start= chrono::steady_clock::now();

        Mat<RR>* HEresmat;
        HEmatrix.decryptParallelRmat(HEresmat, Cctxt);

        end = std::chrono::steady_clock::now();
        diff = end - start;
        timeElapsed = chrono::duration <double, milli> (diff).count()/1000.0;
        totalDectime += timeElapsed;

        /*---------------------------------------*/
        //  Error
        /*---------------------------------------*/
        RR avgerror_mult = to_RR("0");
        Mat<RR>* resmat = new Mat<RR>[nbatching];
        for(long k = 0; k < nbatching; ++k){
            mul(resmat[k], Amat[k], Bmat[k]);
            avgerror_mult += getError(HEresmat[k], resmat[k], nrows, ncols);
        }
        avgerror_mult /= nbatching;

        if(l == niter -1) cout << "Error (mult): " << avgerror_mult << endl;
    }

    cout << "------------------" << endl;
    cout << "Enc time= " << totalEnctime/niter << " s" << endl;
    cout << "Mult time= " << totalEvaltime/niter << " s/ " ;
    cout << "Amortized time= " << amortizedtime/niter << " s" << endl;
    cout << "Dec time= " <<  totalDectime/niter << " s" << endl;

}

void TestHEmatrix::testSIMDMult_Huang(long Arows, long Acols, long Brows, long Bcols, long nbatching)
{

    long logN = 13;

    long ncols = Bcols;
    long nrows = Bcols;

    if(nrows * ncols * nbatching > (1<< (logN-1)))
    {
        cout << "Cannot support the parallism " << endl;
    }

    long pBits = 25;   // scaling factor for message
    long cBits = 15;   // scaling factor for constant
    long lBits = pBits + 5;
    long logQ = (2*cBits + pBits) + lBits;

    long h = 64;
    long keydist = 1;

    struct HEMatpar HEmatpar;
    readHEMatpar(HEmatpar, nrows, ncols, pBits, cBits, logQ, 0, nbatching);

    cout << "------------------" << endl;
    cout << "(dim, nslots, nbatching) = (" << nrows << ","  << HEmatpar.nslots << ","  << HEmatpar.nbatching << ")" << endl;
    cout << "HEAAN PARAMETER logQ: " << logQ << endl;
    cout << "HEAAN PARAMETER logN: " << logN << endl;

    /*---------------------------------------*/
    //  Initialization random seed
    /*---------------------------------------*/
    srand(time(NULL));
    /*---------------------------------------*/
    //  Key Generation
    /*---------------------------------------*/

    TimeUtils timeutils;
    timeutils.start("Scheme generating...");
    Context context(logN, logQ);
    SecretKey secretKey(logN, h);
    Scheme scheme(secretKey, context);

    scheme.addLeftRotKeys(secretKey);
    scheme.addRightRotKeys(secretKey);
    timeutils.stop("Scheme generation");

    HEmatrix HEmatrix(scheme, secretKey, HEmatpar);

    /*---------------------------------------*/
    //  Generate a random matrix
    /*---------------------------------------*/
    Mat<RR>* Mat_Vertical_Stack = new Mat<RR>[2];
    Mat_Vertical_Stack[0].SetDims(Acols, Acols);Mat_Vertical_Stack[1].SetDims(Acols, Acols);
    Mat<RR>* Mat_Horizontal_Stack = new Mat<RR>[2];
    Mat_Horizontal_Stack[0].SetDims(Brows, Brows);Mat_Horizontal_Stack[1].SetDims(Brows, Brows);

    Mat<RR> vertical_Amat, Amat, Random_Amat;
    vertical_Amat.SetDims(2 * Arows, Acols);Amat.SetDims(Acols, Acols);Random_Amat.SetDims(Acols, Acols);
    Mat<RR> horizontal_Bmat_trans,Bmat, Random_Bmat;
    horizontal_Bmat_trans.SetDims(2 * Brows, Brows);Bmat.SetDims(Brows, Brows);Random_Bmat.SetDims(Brows, Brows);

    for(long i = 0; i < Acols; i++)
    {
        for(long j = 0; j < Acols; j++)
        {
            Amat[i][j] = to_RR(rand() % 5);
            Random_Amat[i][j] = to_RR(rand() % 5);
        }
        vertical_Amat[i] = Amat[i];
        vertical_Amat[i + Acols] = Random_Amat[i];
    }

    cout << "Amat : " << endl;
    cout << Amat << endl;
    cout << "Random_Amat : " << endl;
    cout << Random_Amat << endl;
    cout << "vertical_Amat : " << endl;
    cout << vertical_Amat << endl;

    for(long i = 0; i < Bcols; i++)
    {
        for(long j = 0; j < Bcols; j++)
        {
            Bmat[i][j] = to_RR(rand() % 5);
            Random_Bmat[i][j] = to_RR(rand() % 5);
        }
    }
    Mat<RR> Bmat_trans = transpose(Bmat);
    Mat<RR> Random_Bmat_trans = transpose(Random_Bmat);
    for(long j = 0; j < Bcols; j++)
    {
        horizontal_Bmat_trans[j] = Bmat_trans[j];
        horizontal_Bmat_trans[j + Bcols] = Random_Bmat_trans[j];
    }

    cout << "Bmat : " << endl;
    cout << Bmat << endl;
    cout << "Random_Bmat : " << endl;
    cout << Random_Bmat << endl;
    cout << "horizontal_Bmat : " << endl;
    cout << horizontal_Bmat_trans << endl;

    Mat<RR> permutation_Amat, permutation_Bmat;
    permutation_Amat.SetDims(2 * Arows, 2 * Arows);
    permutation_Bmat.SetDims(2 * Bcols, 2 * Bcols);
    generate_random_permutation_matrix(permutation_Amat);   // ???
    generate_random_permutation_matrix(permutation_Bmat);

    cout << "permutation_Amat : " << endl;
    cout << permutation_Amat << endl;

    cout << "permutation_Bmat : " << endl;
    cout << permutation_Bmat << endl;

    Mat<RR> vertical_Amat_mul = permutation_Amat * vertical_Amat;
    Mat<RR> horizontal_Bmat_mul = transpose(horizontal_Bmat_trans) * permutation_Bmat;

    for(long i = 0; i < Arows; i++)
    {
        Mat_Vertical_Stack[0][i] = vertical_Amat_mul[i];
        Mat_Vertical_Stack[1][i] = vertical_Amat_mul[i + Arows];
    }
    for(long i = 0; i < Arows; i++)
    {
        for(long j = 0; j < Arows; j++)
        {
            Mat_Horizontal_Stack[0][i][j] = horizontal_Bmat_mul[i][j];
            Mat_Horizontal_Stack[1][i][j] = horizontal_Bmat_mul[i][j + Arows];
        }
    }

    cout << "Mat_Vertical_Stack[0] : " << endl;
    cout << Mat_Vertical_Stack[0] << endl;
    cout << "Mat_Vertical_Stack[1] : " << endl;
    cout << Mat_Vertical_Stack[1] << endl;

    cout << "transpose Mat_Horizontal_Stack[0] : " << endl;
    cout << Mat_Horizontal_Stack[0] << endl;

    cout << "transpose Mat_Horizontal_Stack[1] : " << endl;
    cout << Mat_Horizontal_Stack[1] << endl;

    /* 0th diagonal multiply vector */
    Ciphertext Actxt;
    HEmatrix.encryptParallelRmat(Actxt, Mat_Vertical_Stack, HEmatpar.pBits, nbatching);

    Ciphertext Bctxt;
    HEmatrix.encryptParallelRmat(Bctxt, Mat_Horizontal_Stack, HEmatpar.pBits, nbatching);

    ZZX** Initpoly;
    HEmatrix.genMultPoly_Parallel_Huang(Initpoly, 0);

    ZZX* shiftpoly;
    HEmatrix.genShiftPoly_Parallel(shiftpoly);

    Ciphertext* res = new Ciphertext[nbatching];
    HEmatrix.HEmatmul_Parallel_Huang(res[0], Actxt, Bctxt, Initpoly, shiftpoly);

    /* 1th diagonal multiply vector */
    Ciphertext Bctxt_rot_one = scheme.leftRotate(Bctxt, 1);

    HEmatrix.genMultPoly_Parallel_Huang(Initpoly, 1);

    HEmatrix.HEmatmul_Parallel_Huang(res[1], Actxt, Bctxt_rot_one, Initpoly, shiftpoly);

    /*---------------------------------------*/
    //  Decryption
    /*---------------------------------------*/
    Mat<RR> DecryptionRes;DecryptionRes.SetDims(2 * Arows, 2 * Arows);
    Mat<RR>** HEresmat = new Mat<RR>*[nbatching];
    for(long i = 0; i < nbatching; i++)
    {
        HEmatrix.decryptParallelRmat(HEresmat[i], res[i]);
        cout << "(" << i << ",0)" << endl;
        cout << HEresmat[i][0] << endl;
        cout << "(" << i << ",1)" << endl;
        cout << HEresmat[i][1] << endl;
        for(long j = 0; j < nbatching; j++)
        {
            for(long k = 0; k < Arows; k++)
            {
                for(long l = 0; l < Arows; l++)
                {
                    DecryptionRes[(k + j * Arows)][(l + (i + j) * Arows) % (2 * Arows)] = HEresmat[i][j][k][l];
                }
            }

        }
    }
    cout << "DecryptionRes : " << endl;
    cout << DecryptionRes << endl;
    // plaintext random matrix multiplication

    Mat<RR> permutation_res = transpose(permutation_Amat) * DecryptionRes * transpose(permutation_Bmat);
    Mat<RR> random_matrix_multiplication = Random_Amat * Random_Bmat;

    cout << "permutation_res : " << endl;
    cout << permutation_res << endl;
    cout << "random_matrix_multiplication : " << endl;
    cout << random_matrix_multiplication << endl;
    bool flag = false;
    for(long i = Arows; i < 2 * Arows; i++)
    {
        for(long j = Arows; j < 2 * Arows; j++)
        {
            if(abs(permutation_res[i][j] - random_matrix_multiplication[i - Arows][j - Arows]) > to_RR(0.01))
            {
                flag = true;
                cout << "The result is not correct !!!" << endl;
                break;
            }
        }
    }
    if(!flag)
    {
        cout << "The result is correct" << endl;
    }
}


//  (C) Copyright Matt Borland 2021.
//  Use, modification and distribution are subject to 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)
//

#ifndef BOOST_MATH_CCMATH_HPP
#define BOOST_MATH_CCMATH_HPP

#include <boost/math/ccmath/sqrt.hpp>
#include <boost/math/ccmath/isinf.hpp>
#include <boost/math/ccmath/isnan.hpp>
#include <boost/math/ccmath/abs.hpp>
#include <boost/math/ccmath/fabs.hpp>
#include <boost/math/ccmath/isfinite.hpp>
#include <boost/math/ccmath/isnormal.hpp>
#include <boost/math/ccmath/fpclassify.hpp>
#include <boost/math/ccmath/frexp.hpp>
#include <boost/math/ccmath/div.hpp>
#include <boost/math/ccmath/logb.hpp>
#include <boost/math/ccmath/ilogb.hpp>
#include <boost/math/ccmath/scalbn.hpp>
#include <boost/math/ccmath/scalbln.hpp>
#include <boost/math/ccmath/floor.hpp>
#include <boost/math/ccmath/ceil.hpp>
#include <boost/math/ccmath/trunc.hpp>
#include <boost/math/ccmath/modf.hpp>
#include <boost/math/ccmath/round.hpp>
#include <boost/math/ccmath/fmod.hpp>
#include <boost/math/ccmath/remainder.hpp>
#include <boost/math/ccmath/copysign.hpp>

#endif // BOOST_MATH_CCMATH_HPP


/**
 * \Description:
 * this node is used to implement the Extended kalman Filter equations, its input are the prediction and the measurement, and its output is final estimation.
 * It subscribes to:
 * - /odom2D: this topic contains the 2D pose (x,y,theta) of the turtlrbot and their covariance.
 * - /beacon distances : this topic contains three distance from the robot to each beacon in order "blue, orange/red, green".
 * It publishes to:
 * - /robot_pose ekf: it publishes the 2D pose (x,y,theta) of the turtlrbot and their corresponding covariance after modifying the predicted pose using the measurement.
 *
 * ALI&SERRANO_ECN_M1_2017
 */

//Cpp
#include <sstream>
#include <stdio.h>
#include <vector>
#include <iostream>
#include <stdlib.h>
#include <math.h>
#include <Eigen/Dense>

using namespace Eigen;
using namespace std;
using Eigen::MatrixXd;

//ROS
#include "ros/ros.h"

// Include here the ".h" files corresponding to the topic type you use.
#include <nav_msgs/Odometry.h>
#include "tf/transform_datatypes.h"

#include "std_msgs/MultiArrayLayout.h"
#include "std_msgs/MultiArrayDimension.h"
#include <std_msgs/Float64MultiArray.h>

#include <geometry_msgs/Point.h>
#include <geometry_msgs/Pose2D.h>
#include <turtle_ekf/Pose2DWithCovariance.h>

//#include <geometry_msgs/PoseWithCovariance.h>
//#include <std_msgs/Float32.h>

// global variables...
turtle_ekf::Pose2DWithCovariance odom2D;
std_msgs::Float64MultiArray dist;
bool is_odom_available= false, is_dist_available=false;

// Callback functions...
void odom_Callback( turtle_ekf::Pose2DWithCovariance odom2D_msg) {
    is_odom_available= true;
    odom2D = odom2D_msg;
}

void msurmnt_Callback ( std_msgs::Float64MultiArray  dist_msg) {
    is_dist_available= true;
    dist = dist_msg;
}

int main (int argc, char** argv){

    //ROS Initialization
    ros::init(argc, argv, "ekf_node");

    // Define your node handles
    ros::NodeHandle nh_loc("~"), nh_glob;

    // Read the node parameters if any

    // Declare your node's subscriptions and service clients
    ros::Subscriber odom2D_sub = nh_glob.subscribe<turtle_ekf::Pose2DWithCovariance>("/odom2D",1, odom_Callback) ;
    ros::Subscriber msurmnt_sub = nh_glob.subscribe<std_msgs::Float64MultiArray>("/beacon_distances",1, msurmnt_Callback) ;

    // Declare you publishers and service servers
    ros::Publisher  pose_pub = nh_glob.advertise<turtle_ekf::Pose2DWithCovariance>("/robot_pose_ekf",1) ;

    // node initilization
    MatrixXd Q_alpha =  MatrixXd::Zero(3,3);
    MatrixXd Q_gamma =  MatrixXd::Zero(3,3);
    MatrixXd P =  MatrixXd::Zero(3,3);
    MatrixXd C =  MatrixXd::Zero(3,3);
    MatrixXd D =  MatrixXd::Zero(3,3);
    MatrixXd K =  MatrixXd::Zero(3,3);
    MatrixXd I =  MatrixXd::Zero(3,3);
    I(0,0)= 1;   I(1,1)= 1;     I(2,2)= 1;

    VectorXd X(3);
    VectorXd Y(3);
    VectorXd Y_hat(3);

    geometry_msgs::Point B1, B2, B3;
    B1.x=0;  B1.y=0; B3.x=4.40; B3.y=3.20; B2.x=0;  B2.y=5.20;

    // Init
    X(0)=1;   X(1)=1;   X(2)= 0; //M_PI/2;
// init matrices
//    for (int i=0; i++; i<3){
//        for(int j=0; j++; j<3){
//            P(i,j)= 0;
//            C(i,j)= 0;
//            D(i,j)= 0;
//            K(i,j)= 0;
//        }
//    }

    ros::Rate rate(10);

    while (ros::ok()){

        ros::spinOnce();
                if( ! is_dist_available || !is_odom_available ){
                    ROS_INFO("Waiting for odom2D and/or dist") ;
                    rate.sleep() ;
                    continue ;
                }

        // odometry part:
        // Prediction step: predicted_state Xk+1/k "extract pose from odom"
        X(0)= odom2D.pose.x;
        X(1)= odom2D.pose.y;
        X(2)= odom2D.pose.theta;

        // state_trans_uncertainty_noise: Q_alpha "extract covariance matrix from odom"
        Q_alpha(0,0)= odom2D.Covariance[0];
        Q_alpha(1,1)= odom2D.Covariance[4];
        Q_alpha(2,2)= odom2D.Covariance[8];

        // jacobian_matrix: A = I

        // pred_covariance_est: Pk+1/k = A. PK/k. A^T + Q_alpha  "P= A*P*A.transpose() + Q_alpha;" but A=I
        P=P+Q_alpha;

        // measurement part:

        // Actual measurement  Y
        Y(0)= dist.data[0]/100;
        Y(1)= dist.data[1]/100;
        Y(2)= dist.data[2]/100;

        // Expected measurement Y_hat
        Y_hat(0)= sqrt( (B1.x-X(0))*(B1.x-X(0)) + (B1.y-X(1))*(B1.y-X(1)) );
        Y_hat(1)= sqrt( (B2.x-X(0))*(B2.x-X(0)) + (B2.y-X(1))*(B2.y-X(1)) );
        Y_hat(2)= sqrt( (B3.x-X(0))*(B3.x-X(0)) + (B3.y-X(1))*(B3.y-X(1)) );

cout<< "Y:  "<< Y(0) << ", " << Y(1) << ", " <<  Y(2) << endl<< endl;
cout<< "Y_hat:  " << Y_hat(0) << ", " << Y_hat(1) << ", " <<  Y_hat(2) <<  endl ;

        // Measurement jacobian C_matrix
        C(0,0)= 2*(X(0) - B1.x) / Y_hat(0);
        C(0,1)= 2*(X(1) - B1.y) / Y_hat(0);
        C(1,0)= 2*(X(0) - B1.x) / Y_hat(1);
        C(1,1)= 2*(X(1) - B1.y) / Y_hat(1);
        C(2,0)= 2*(X(0) - B1.x) / Y_hat(2);
        C(2,1)= 2*(X(1) - B1.y) / Y_hat(2);

        // measurement covariance Q_gamma
        Q_gamma(0,0)= 10;
        Q_gamma(1,1)= 10;
        Q_gamma(2,2)= 10;

        if (Y(0)> 0.3 && Y(0)< 6)
        Q_gamma(0,0)= (.0000842*Y(0))+0.0158;
        if (Y(1)> 0.3 && Y(1)< 6)
        Q_gamma(1,1)= (.0000842*Y(1))+0.0158;
        if (Y(2)> 0.3 && Y(2)< 6)
        Q_gamma(2,2)= (.0000842*Y(2))+0.0158;

if (Y(0)> 0.3 && Y(0)< 6 && Y(1)> 0.3 && Y(1)< 6 && Y(2)> 0.3 && Y(2)< 6 ){

       // Residual covariance D is Q_(Y-Y_hat)
       D= ( C*P*C.transpose() + Q_gamma );

       // Correction propotional gain  K= P*C.transpose() * (C*P*C.transpose() + Q_gamma )^-1
        K= P*C.transpose() * D.reverse();

        // Estimation Step: estimated_state Xk+1/k+1 = Xk+1/k + Kk (Yk - Yk_hat)
        X = X + K*(Y-Y_hat);

        // Estimated_covariance_est: Pk+1/k+1 = ( I - Kk*Ck ) * Pk+1/k
        P= ( I - K*C ) * P;
}
        std::cout<<  "P" << P << endl << endl;
        std::cout<<  "C" << C << endl << endl;
        std::cout<<  "D" << D << endl << endl;
        std::cout<<  "K" << K << endl << endl;

        // publish the final pose ........................................................................................................
        turtle_ekf::Pose2DWithCovariance pose_ekf;
        pose_ekf.pose.x = X(0);
        pose_ekf.pose.y = X(1);
        pose_ekf.pose.theta =  X(2);

        pose_ekf.Covariance[0] =  P(0,0);
        pose_ekf.Covariance[1] =  0;
        pose_ekf.Covariance[2] =  0;
        pose_ekf.Covariance[3] =  0;
        pose_ekf.Covariance[4] =  P(1,1);
        pose_ekf.Covariance[5] =  0;
        pose_ekf.Covariance[6] =  0;
        pose_ekf.Covariance[7] =  0;
        pose_ekf.Covariance[8] =  P(2,2);

        pose_pub.publish(pose_ekf);

        rate.sleep();
    }

}


// Copyright (c) 2018 Evan S Weinberg
// Thick restarted and deflated Lanczos.
// Finds some number of exterior largest or
// smallest values, NOT in magnitude.
// (Doesn't try solve the interior problem,
// though there is a way to trick it with the
// operator (A-sigma)^2 ...)
// Based on arXiv:1512:08135

#include <iostream>
#include <iomanip>
#include <cmath>
#include <string>
#include <sstream>
#include <complex>
#include <random>
#include <vector>
#include <functional>

// Borrow dense matrix eigenvalue routines.
#include <Eigen/Dense>

#include "blas/generic_vector.h"

#include "../square_laplace.h"

// Operator class
#include "../operator.h"
#include "../poly_operator.h"

// Lanczos
#include "../lanczos.h"
#include "../thick_deflate_lanczos.h"

using namespace std;
using namespace Eigen;

typedef Matrix<double, Dynamic, Dynamic, ColMajor> dMatrix;
typedef Matrix<std::complex<double>, Dynamic, Dynamic, ColMajor> cMatrix;

int main(int argc, char** argv)
{
  complex<double> *rhs_cplx;

  // Set output precision to be long.
  cout << setprecision(10);

  // RNG related things.
  std::mt19937 generator (1337u); // RNG, 1337u is the seed.

  ///////////////////////////////////
  // Properties of Linear operator //
  ///////////////////////////////////
  double inv_variance = 6.0; // inverse of variance for gaussian non-compact U(1) links.

  // Basic information about the lattice.
  int length = 24;
  double m_sq = 0.001;

  // Some start-up.
  int volume = length*length;

  // Create a random compact U(1) link.
  complex<double>* gauge_links = allocate_vector<complex<double>>(2*length*length);
  gaussian_real(gauge_links, 2*length*length, generator, 1.0/inv_variance);
  polar(gauge_links, 2*length*length);

  // Vectors.
  rhs_cplx = allocate_vector<complex<double>>(volume);

  // Zero out the vector.
  zero_vector(rhs_cplx, length*length);

  // Structure which gets passed to the function.
  laplace_gauged_struct lapstr_gauged;
  lapstr_gauged.length = length;
  lapstr_gauged.m_sq = m_sq;
  lapstr_gauged.gauge_links = gauge_links;

  // Uncomment this to get the free field.
  //constant_vector(gauge_links, 1.0, 2*volume);
  //std::cout << "Free case.\n\n";

  std::cout << "Interacting case.\n\n";

  //////////////////////////////////////
  // Create a linear operator object. //
  //////////////////////////////////////

  FunctionWrapper<complex<double>> lap_fcn(square_laplacian_gauged, &lapstr_gauged, volume);

  //////////////////////////////////////
  // Properties for Lanczos algorithm //
  //////////////////////////////////////

  // Fill a struct
  TRCLStruct lanc_props;
  lanc_props.n_ev = 10; // get 10 eigenvalues
  lanc_props.m = 30; // subspace size of 20
  lanc_props.tol = 1e-10; // lock at a tolerance of 1e-8
  lanc_props.max_restarts = 500; // maximum of 10 restarts
  lanc_props.preserved_space = 10; // space preserved after restart,
                                  // set to -1 to default to m/4+1
  lanc_props.deflate = false; // don't deflate locked eigenvalues
  lanc_props.generator = &generator; // passed by reference
  lanc_props.verbose = false;

  // Get the smallest eigenvalues
  ThickRestartComplexLanczos<double,std::less<double>> lanczos(&lap_fcn, lanc_props);

  //////////////
  // Compute! //
  //////////////

  lanczos.compute();

  // Get number of converged eigenvalues
  int n_converged = lanczos.num_converged();

  // Get the eigenvalues
  double* eigenvalues = new double[n_converged];
  lanczos.ritzvalues(eigenvalues);

  // Print the Ritz values
  std::cout << "The " << n_converged << " converged eigenvalues are:\n";
  for (int i = 0; i < n_converged; i++) {
    std::cout << eigenvalues[i] << "\n";
  }

  // Get converged eigenvectors
  complex<double>** eigenvectors = new complex<double>*[n_converged];
  for (int i = 0; i < n_converged; i++) {
    eigenvectors[i] = allocate_vector<complex<double>>(volume);
  }

  // Get the Ritz vectors
  lanczos.ritzvectors(eigenvectors);

  ///////////////////////////////////////////
  // Do the polynomial accelerated version //
  ///////////////////////////////////////////

  std::cout << "\nPolynomial Accleration\n";

  // Get approximate bounds from a small Lanczos
  const int m_mini = 6;
  SimpleComplexLanczos<double> simp_lanczos(&lap_fcn, m_mini, generator);
  simp_lanczos.compute();
  double approx_eigs[m_mini];
  simp_lanczos.ritzvalues((double*)approx_eigs);
  // print the approximate eigenvalues
  std::cout << "The " << m_mini << " Ritz values are:\n";
  for (int i = 0; i < m_mini; i++) {
    std::cout << approx_eigs[i] << "\n";
  }
  std::cout << "\n";
  double approx_min = approx_eigs[0]; // get the approximate min
  double approx_max = approx_eigs[m_mini-1]*1.2; // and overshoot the max
  std::cout << "The linear op window is " << approx_min << " to " << approx_max << "\n\n";

  // Make a linear interpolation of the laplace op.
  LinearMapToUnit<complex<double>,double> linear_lap(&lap_fcn, approx_min, approx_max);

  // And make the 20th order poly accelerated form
  OneOverOnePlusX<complex<double> > poly_accel(&linear_lap, 50);

  // Fill a struct
  TRCLStruct poly_lanc_props;
  poly_lanc_props.n_ev = 10; // get 10 eigenvalues
  poly_lanc_props.m = 30; // subspace size of 20
  poly_lanc_props.tol = 1e-10; // lock at a tolerance of 1e-8
  poly_lanc_props.max_restarts = 20; // maximum of 10 restarts
  poly_lanc_props.preserved_space = 10; // space preserved after restart,
                                  // set to -1 to default to m/4+1
  poly_lanc_props.deflate = true; // deflate locked eigenvalues
  poly_lanc_props.generator = &generator; // passed by reference
  poly_lanc_props.verbose = false;

  // Get the largest eigenvalues, since we're poly acceling
  ThickRestartComplexLanczos<double,std::greater<double>> poly_lanczos(&poly_accel, poly_lanc_props);
  poly_lanczos.compute();
  int n_poly_converged = poly_lanczos.num_converged();
  std::cout << n_poly_converged << "\n";

  // Get the poly eigenvalues
  double* poly_eigenvalues = new double[n_poly_converged];
  poly_lanczos.ritzvalues(poly_eigenvalues);

  // Print the Ritz values
  std::cout << "The " << n_poly_converged << " converged eigenvalues are:\n";
  for (int i = 0; i < n_poly_converged; i++) {
    std::cout << poly_eigenvalues[i] << "\n";
  }

  // Get the poly eigenvectors
  complex<double>** poly_eigenvectors = new complex<double>*[n_poly_converged];
  for (int i = 0; i < n_poly_converged; i++) {
    poly_eigenvectors[i] = allocate_vector<complex<double>>(volume);
  }

  // Get the Ritz vectors
  poly_lanczos.ritzvectors(poly_eigenvectors);

  // Get eigenvalues of original system
  complex<double>* intermediate = allocate_vector<complex<double>>(volume);
  std::cout << "\nThe " << n_poly_converged << " reconstructed eigenvalues are:\n";
  for (int i = 0; i < n_poly_converged; i++) {
    zero_vector(intermediate, volume);
    lap_fcn(intermediate, poly_eigenvectors[i]);
    double rec_eval = re_dot(poly_eigenvectors[i], intermediate, volume);
    std::cout << rec_eval << "\n";
  }

  //////////////////////////
  // Get all eigenvalues! //
  //////////////////////////

  if (volume <= 256) {

    // Let's get the eigenvalues of the full operator!

    // Allocate a sufficiently gigantic matrix.
    cMatrix mat_cplx = cMatrix::Zero(volume, volume);

    // Form matrix elements. This is where it's important that
    // dMatrix and cMatrix are column major.
    // I should probably make this safer by using a "Map".
    for (int i = 0; i < volume; i++)
    {
      // Set a point on the rhs for a matrix element.
      zero_vector(rhs_cplx, volume);
      rhs_cplx[i] = 1.0;

      // Where we put the result of the matrix element.
      complex<double>* mptr = &(mat_cplx(i*volume));

      square_laplacian_gauged(mptr, rhs_cplx, &lapstr_gauged);
    }

    // Get the eigenvalues.
    SelfAdjointEigenSolver<cMatrix> eigsolve_cplx(volume);
    eigsolve_cplx.compute(mat_cplx);

    std::cout << "The eigenvalues are:\n" << eigsolve_cplx.eigenvalues() << "\n";

    ////////////////////////////////
    // COMPARE LOWEST EIGENVECTOR //
    ////////////////////////////////

    std::cout << "Compare results:\n\n";
    std::cout << "Lanczos Exact Ratio\n";
    for (int i = 0; i < volume; i++) {
      std::cout << eigenvectors[0][i] << " " << eigsolve_cplx.eigenvectors().col(0)(i)
                << " " << eigenvectors[0][i]/eigsolve_cplx.eigenvectors().col(0)(i) << "\n";
    }
  }

  //////////////
  // CLEAN UP //
  //////////////

  delete[] eigenvalues;
  delete[] poly_eigenvalues;
  for (int i = 0; i < n_converged; i++) {
    deallocate_vector(&eigenvectors[i]);
  }
  delete[] eigenvectors;
  for (int i = 0; i < n_poly_converged; i++) {
    deallocate_vector(&poly_eigenvectors[i]);
  }
  delete[] poly_eigenvectors;
  deallocate_vector(&intermediate);

  deallocate_vector(&rhs_cplx);
  deallocate_vector(&gauge_links);
  return 0;
}


#include "AnalysisGraph.hpp"
#include <range/v3/all.hpp>
#include <unsupported/Eigen/MatrixFunctions>
#include <boost/range/adaptors.hpp>
#include <tqdm.hpp>

using namespace std;
using Eigen::VectorXd, Eigen::MatrixXd;
using fmt::print, fmt::format;
namespace rs = ranges;
using boost::adaptors::transformed;

using fmt::print;
/*
 ============================================================================
 Private: Prediction
 ============================================================================
*/

void AnalysisGraph::generate_latent_state_sequences(
    double initial_prediction_step) {

  // Allocate memory for prediction_latent_state_sequences
  this->predicted_latent_state_sequences.clear();
  this->predicted_latent_state_sequences = vector<vector<VectorXd>>(
      this->res,
      vector<VectorXd>(this->pred_timesteps, VectorXd(this->num_vertices() * 2)));

  cout << "\nPredicting for " << this->pred_timesteps << " time steps..." << endl;
  for (int samp : tq::trange(this->res)) {
      // The sampled transition matrices would be either of matrix exponential
      // (continuous) version or discretized version depending on whether the
      // matrix exponential (continuous) version or the discretized transition
      // matrix version had been used to train the model. This allows us to
      // make the prediction routine common for both the versions, except for
      // the exponentiation of the matrix exponential (continuous) transition
      // matrices.
      MatrixXd A;

      this->generate_head_node_latent_sequences(
          samp, initial_prediction_step + this->pred_timesteps);

      if (this->continuous) {
//          // Here A = Ac = this->transition_matrix_collection[samp] (continuous)
//
//          // Evolving the system till the initial_prediction_step
//          A = (this->transition_matrix_collection[samp] *
//                   initial_prediction_step).exp();
//
//          this->predicted_latent_state_sequences[samp][0] =
//                               A * this->initial_latent_state_collection[samp];
//
//          // After jumping to time step ips - 1, we take one step of length Δt
//          // at a time.
//          // So compute the transition matrix for a single step.
//          // Computing the matrix exponential for a Δt time step.
//          // By default we are using Δt = 1
//          // A = e^{Ac * Δt)
//          A = (this->transition_matrix_collection[samp] * this->delta_t).exp();

          /////////////////
          A = this->transition_matrix_collection[samp].exp();

          // Evolving the system till the initial_prediction_step
          this->current_latent_state = this->initial_latent_state_collection[samp];

          for (int ts = 0; ts < initial_prediction_step; ts++) {
            this->update_latent_state_with_generated_derivatives(ts, ts + 1);

            // Set derivatives for frozen nodes
            for (const auto & [ v, deriv_func ] : this->external_concepts) {
              const Indicator& ind = this->graph[v].indicators[0];
              this->current_latent_state[2 * v + 1] = deriv_func(ts, ind.mean);
            }

            this->current_latent_state = A * this->current_latent_state;
          }
          //this->update_latent_state_with_generated_derivatives(0, initial_prediction_step);
          //this->current_latent_state = A * this->current_latent_state;

          this->update_latent_state_with_generated_derivatives(
              initial_prediction_step, initial_prediction_step + 1);

          // Set derivatives for frozen nodes
          for (const auto & [ v, deriv_func ] : this->external_concepts) {
            const Indicator& ind = this->graph[v].indicators[0];
            this->current_latent_state[2 * v + 1] =
                deriv_func(initial_prediction_step, ind.mean);
          }

          this->predicted_latent_state_sequences[samp][0] = this->current_latent_state;
          /////////////////
      } else {
          // Here A = Ad = this->transition_matrix_collection[samp] (discrete)
          // This is the discrete transition matrix to take a single step of
          // length Δt
          A = this->transition_matrix_collection[samp];

          // Evolving the system till the initial_prediction_step
          this->current_latent_state = this->initial_latent_state_collection[samp];

          for (int ts = 0; ts < initial_prediction_step; ts++) {
            this->update_latent_state_with_generated_derivatives(ts, ts + 1);

            // Set derivatives for frozen nodes
            for (const auto & [ v, deriv_func ] : this->external_concepts) {
              const Indicator& ind = this->graph[v].indicators[0];
              this->current_latent_state[2 * v + 1] = deriv_func(ts, ind.mean);
            }

            this->current_latent_state = A * this->current_latent_state;
          }

          this->update_latent_state_with_generated_derivatives(
              initial_prediction_step, initial_prediction_step + 1);

          // Set derivatives for frozen nodes
          for (const auto & [ v, deriv_func ] : this->external_concepts) {
            const Indicator& ind = this->graph[v].indicators[0];
            this->current_latent_state[2 * v + 1] =
                                  deriv_func(initial_prediction_step, ind.mean);
          }

          this->predicted_latent_state_sequences[samp][0] = this->current_latent_state;
      }

      // Clear out perpetual constraints residual from previous sample
      this->perpetual_constraints.clear();

      if (this->clamp_at_derivative) {
          // To clamp a latent state value to x_c at prediction step 1 via
          // clamping the derivative, we have to perturb the derivative at
          // prediction step 0, before evolving it to prediction time step 1.
          // So we have to look one time step ahead whether we have to clamp
          // at 1.
          //
          if (delphi::utils::in(this->one_off_constraints, 1)) {
              this->perturb_predicted_latent_state_at(1, samp);
          }
      }

      // Since we used ts = 0 for the initial_prediction_step - 1, this loop is
      // one ahead of the prediction time steps. That is, index 1 in the
      // prediction data structures is the 0th index for the requested
      // prediction sequence. Hence, when we are at time step ts, we should
      // check whether there are any constconstraints
        // time step index is 1
      for (int ts = 1; ts < this->pred_timesteps; ts++) {
          // When continuous: The standard matrix exponential equation is,
          //                        s_{t+Δt} = e^{Ac * Δt } * s_t
          //                  Since vector indices are integral values, and in
          //                  the implementation s is the vector, to index into
          //                  the vector we uses consecutive integers. Thus in
          //                  the implementation, the matrix exponential
          //                  equation becomes,
          //                      s_{t+1} = e^{Ac * Δt } * s_t
          //                  What this equation says is that although vector
          //                  indices advance by 1, the duration between two
          //                  predictions stored in two adjacent vector cells
          //                  need not be 1. The time duration is actually Δt.
          //                  The actual line of code represents,
          //                        s_t = e^{Ac * Δt } * s_{t-1}
          // When discrete  : s_t = Ad * s_{t-1}
          this->current_latent_state = A * this->predicted_latent_state_sequences[samp][ts - 1];
          this->update_latent_state_with_generated_derivatives(
              initial_prediction_step + ts, initial_prediction_step + ts + 1);
          this->predicted_latent_state_sequences[samp][ts] = this->current_latent_state;

          // Set derivatives for frozen nodes
          for (const auto & [ v, deriv_func ] : this->external_concepts) {
            const Indicator& ind = this->graph[v].indicators[0];
            this->predicted_latent_state_sequences[samp][ts][2 * v + 1] =
                            deriv_func(initial_prediction_step + ts, ind.mean);
          }

          if (this->clamp_at_derivative ) {
              if (ts == this->rest_derivative_clamp_ts) {
                  // We have perturbed the derivative at ts - 1. If we do not
                  // take any action, that clamped derivative will be in effect
                  // until another clamping or end of prediction. This is where
                  // we take the necessary actions.

                  if (is_one_off_constraints) {
                      // We should revert the derivative to its original value
                      // at this time step so that clamping does not affect
                      // time step ts + 1.
                      for (auto constraint : this->one_off_constraints.at(ts)) {
                          int node_id = constraint.first;

                          this->predicted_latent_state_sequences[samp][ts]
                              (2 * node_id + 1) =
                              this->initial_latent_state_collection[samp]
                              (2 * node_id + 1);
                  }
                  } else {
                      for (auto [node_id, value]: this->one_off_constraints.at(ts)) {
                          this->predicted_latent_state_sequences[samp][ts]
                                                        (2 * node_id + 1) = 0;
                      }
                  }
              }

              // To clamp a latent state value to x_c at prediction step ts + 1
              // via clamping the derivative, we have to perturb the derivative
              // at prediction step ts, before evolving it to prediction time
              // step ts + 1. So we have to look one time step ahead whether we
              // have to clamp at ts + 1 and clamp the derivative now.
              //
              if (delphi::utils::in(this->one_off_constraints, ts + 1) ||
                      !this->perpetual_constraints.empty()) {
                  this->perturb_predicted_latent_state_at(ts + 1, samp);
              }
          } else {
              // Apply constraints to latent state if any
              // Logic of this condition:
              //    Initially perpetual_constraints = ∅
              //
              //    one_off_constraints.at(ts) = ∅ => Unconstrained ∀ ts
              //
              //    one_off_constraints.at(ts) ‡ ∅
              //        => ∃ some constraints (But we do not know what kind)
              //        => Perturb latent state
              //           Call perturb_predicted_latent_state_at()
              //           one_off_constraints == true
              //                => We are applying One-off constraints
              //           one_off_constraints == false
              //                => We are applying perpetual constraints
              //                => We add constraints to perpetual_constraints
              //                => perpetual_constraints ‡ ∅
              //                => The if condition is true ∀ subsequent time steps
              //                   after ts (the first time step s.t.
              //                   one_off_constraints.at(ts) ‡ ∅
              //                => Constraints are perpetual
              //
              if (delphi::utils::in(this->one_off_constraints, ts) ||
                      !this->perpetual_constraints.empty()) {
                  this->perturb_predicted_latent_state_at(ts, samp);
              }
          }
      }
  }
}

/*
 * Applying constraints (interventions) to latent states
 * Check the data structure definition to get more descriptions
 *
 * Clamping at time step ts for value v
 * Two methods work in tandem to achieve this. Let us label them as follows:
 *      glss  : generate_latent_state_sequences()
 *      ppls@ : perturb_predicted_latent_state_at()
 *                                 ┍━━━━━━━━━━━━━━━━━━━━━━━━┑
 *                                 │           How          ┃
 *                                 ┝━━━━━━━━━━┳━━━━━━━━━━━━━┫
 *                                 │ One-off  ┃  Perpetual  ┃
 * ━━━━━━┳━━━━━━━━━━━━┯━━━━━━━━━━━━┿━━━━━━━━━━┻━━━━━━━━━━━━━┫
 *       ┃            │ Clamp at   │           v - x_{ts-1} ┃
 *       ┃            │ (by ppls@) │   ts-1 to ──────────── ┃
 *       ┃            │            │               Δt       ┃
 *       ┃ Derivative ├┈┈┈┈┈┈┈┈┈┈┈┈┼┈┈┈┈┈┈┈┈┈┈┰┈┈┈┈┈┈┈┈┈┈┈┈┈┨
 * Where ┃            │ Reset at   │ ts to ẋ₀ ┃ ts to 0     ┃
 *       ┃            │ (by glss)  │ from S₀  ┃             ┃
 *       ┣━━━━━━━━━━━━┿━━━━━━━━━━━━┿━━━━━━━━━━╋━━━━━━━━━━━━━┫
 *       ┃ Value      │ Clamp at   │ ts to v  ┃ ∀ t≥ts to v ┃
 *       ┃            │ (by ppls@) │          ┃             ┃
 * ━━━━━━┻━━━━━━━━━━━━┷━━━━━━━━━━━━┷━━━━━━━━━━┻━━━━━━━━━━━━━┛
 */
void AnalysisGraph::perturb_predicted_latent_state_at(int timestep, int sample_number) {
    // Let vertices of the CAG be v = 0, 1, 2, 3, ...
    // Then,
    //    indices 2*v keeps track of the state of each variable v
    //    indices 2*v+1 keeps track of the state of ∂v/∂t

    if (this->clamp_at_derivative) {
        for (auto [node_id, value]: this->one_off_constraints.at(timestep)) {
            // To clamp the latent state value to x_c at time step t via
            // clamping the derivative, we have to clamp the derivative
            // appropriately at time step t-1.
            // Example:
            //      Say we want to clamp the latent state at t=6 to value
            //      x_c (i.e. we want x₆ = x_c. So we have to set the
            //      derivative at t=6-1=5, ẋ₅, as follows:
            //                  x_c - x₅
            //             ẋ₅ = --------- ........... (1)
            //                     Δt
            //             x₆ = x₅ + (ẋ₅ × Δt)
            //                = x_c
            //      Thus clamping ẋ₅ (at t = 5) as described in (1) gives
            //      us the desired clamping at t = 5 + 1 = 6
            double clamped_derivative = (value -
                    this->predicted_latent_state_sequences[sample_number]
                    [timestep - 1](2 * node_id)) / this->delta_t;

            this->predicted_latent_state_sequences[sample_number]
                [timestep - 1](2 * node_id + 1) = clamped_derivative;
        }

        // Clamping the derivative at t-1 changes the value at t.
        // According to our model, derivatives never chance. So if we
        // do not revert it, clamped derivative stays till another
        // clamping or the end of prediction. Since this is a one-off
        // clamping, we have to return the derivative back to its
        // original value at time step t, before we use it to evolve
        // time step t + 1. Thus we remember the time step at which we
        // have to perform this.
        this->rest_derivative_clamp_ts = timestep;

        return;
    }

    if (this->is_one_off_constraints) {

        for (auto [node_id, value]: this->one_off_constraints.at(timestep)) {
            this->predicted_latent_state_sequences[sample_number][timestep](2 * node_id) = value;
        }
    } else { // Perpetual constraints
        if (delphi::utils::in(this->one_off_constraints, timestep)) {
            // Update any previous perpetual constraints
            for (auto [node_id, value]: this->one_off_constraints.at(timestep)) {

                this->perpetual_constraints[node_id] = value;
            }
        }

        // Apply perpetual constraints
        for (auto [node_id, value]: this->perpetual_constraints) {

            this->predicted_latent_state_sequences[sample_number][timestep](2 * node_id) = value;
        }
    }
}

void AnalysisGraph::generate_observed_state_sequences() {
  using rs::to, rs::views::transform;

  // Allocate memory for observed_state_sequences
  this->predicted_observed_state_sequences.clear();
  this->predicted_observed_state_sequences =
      vector<PredictedObservedStateSequence>(
          this->res,
          PredictedObservedStateSequence(this->pred_timesteps,
                                         vector<vector<double>>()));

  for (int samp = 0; samp < this->res; samp++) {
    vector<VectorXd>& sample = this->predicted_latent_state_sequences[samp];

    this->predicted_observed_state_sequences[samp] =
        sample | transform([this](VectorXd latent_state) {
          return this->generate_observed_state(latent_state);
        }) |
        to<vector>();
  }
}

vector<vector<double>>
AnalysisGraph::generate_observed_state(VectorXd latent_state) {
  using rs::to, rs::views::transform;

  int num_verts = this->num_vertices();

  vector<vector<double>> observed_state(num_verts);

  for (int v = 0; v < num_verts; v++) {
    vector<Indicator>& indicators = (*this)[v].indicators;

    observed_state[v] = vector<double>(indicators.size());

    observed_state[v] = indicators | transform([&](Indicator ind) {
                          return ind.mean * latent_state[2 * v];
                        }) |
                        to<vector>();
  }

  return observed_state;
}

FormattedPredictionResult AnalysisGraph::format_prediction_result() {

  // NOTE: To facilitate clamping derivatives, we start prediction one time
  //       step before the requested prediction start time. We are omitting
  //       that additional time step from the results returned to the user.
  this->pred_timesteps--;

  // Access
  // [ sample ][ time_step ][ vertex_name ][ indicator_name ]
  auto result = FormattedPredictionResult(
      this->res,
      vector<unordered_map<string, unordered_map<string, double>>>(
          this->pred_timesteps));

  for (int samp = 0; samp < this->res; samp++) {
    // NOTE: To facilitate clamping derivatives, we start prediction one time
    //       step before the requested prediction start time. We are omitting
    //       that additional time step from the results returned to the user.
    for (int ts = 1; ts <= this->pred_timesteps; ts++) {
      for (auto [vert_name, vert_id] : this->name_to_vertex) {
        for (auto [ind_name, ind_id] : (*this)[vert_id].nameToIndexMap) {
          result[samp][ts - 1][vert_name][ind_name] =
              this->predicted_observed_state_sequences[samp][ts][vert_id]
                                                      [ind_id];
        }
      }
    }
  }

  return result;
}

void AnalysisGraph::run_model(int start_year,
                              int start_month,
                              int end_year,
                              int end_month) {
  if (!this->trained) {
    print("Passed untrained Causal Analysis Graph (CAG) Model. \n",
          "Try calling <CAG>.train_model(...) first!");
    throw "Model not yet trained";
  }

  // Check for sensible prediction time step ranges.
  // NOTE: To facilitate clamping derivatives, we start prediction one time
  //       step before the requested prediction start time. Therefor, the
  //       possible valid earliest prediction time step is one time step after
  //       the training start time.
  if (start_year < this->training_range.first.first ||
      (start_year == this->training_range.first.first &&
       start_month <= this->training_range.first.second)) {
    print("The initial prediction date can't be before the "
         "initial training date. Defaulting initial prediction date "
         "to initial training date + 1 month.");
    start_month = this->training_range.first.second + 1;
    if (start_month == 13) {
        start_year = this->training_range.first.first + 1;
        start_month = 1;
    } else {
        start_year = this->training_range.first.first;
    }
  }

  /*
   *              total_timesteps
   *   ____________________________________________
   *  |                                            |
   *  v                                            v
   * start training                          end prediction
   *  |--------------------------------------------|
   *  :           |--------------------------------|
   *  :         start prediction                   :
   *  ^           ^                                ^
   *  |___________|________________________________|
   *      diff              pred_timesteps
   */
  int total_timesteps =
      this->calculate_num_timesteps(this->training_range.first.first,
                                    this->training_range.first.second,
                                    end_year,
                                    end_month);

  this->pred_timesteps = this->calculate_num_timesteps(
      start_year, start_month, end_year, end_month);

  int pred_init_timestep = total_timesteps - pred_timesteps;

  int year = start_year;
  int month = start_month;

  this->pred_range.clear();
  this->pred_range = vector<string>(this->pred_timesteps);

  for (int t = 0; t < this->pred_timesteps; t++) {
    this->pred_range[t] = to_string(year) + "-" + to_string(month);

    if (month == 12) {
      year++;
      month = 1;
    }
    else {
      month++;
    }
  }

  // NOTE: To facilitate clamping derivatives, we start prediction one time
  //       step before the requested prediction start time. When we are
  //       returning results, we have to remove the predictions at the 0th
  //       index of each predicted observed state sequence.
  //       Adding that additional time step.
  // t     = Requested prediction start time step
  //       = pred_init_timestep
  // t - 1 = Prediction start time step
  //       = pred_init_timestep - 1
  pred_init_timestep--;
  this->pred_timesteps++;

  this->generate_latent_state_sequences(pred_init_timestep);
  this->generate_observed_state_sequences();
}

void AnalysisGraph::add_constraint(int step, string concept_name, string indicator_name,
                                                double indicator_clamp_value) {
    // When constraining latent state derivatives, to reach the
    // prescribed value for an observation at projection step t, we have
    // to  clamp the derivative at projection step t-1 appropriately.
    // Therefor, to constrain at projection step 0, we have to clamp the
    // derivative at projection step -1.
    // To facilitate this, we start projection one time step before the
    // requested projection start time.
    // To correct for that and make internal vector indexing easier,
    // nicer and less error prone, we shift all the constraints by one
    // step up.
    step++;

    // Check whether this concept is in the CAG
    if (!delphi::utils::in(this->name_to_vertex, concept_name)) {
        print("Concept \"{0}\" not in CAG!\n", concept_name);
        return;
    }

    int concept_id = this->name_to_vertex.at(concept_name);
    Node& n = (*this)[concept_id];

    if (n.indicators.size() <= 0) {
        // This concept does not have any indicators attached.
        print("Concept \"{0}\" does not have any indicators attached!\n", concept_name);
        print("\tCannot set constraint\n");
        return;
    }

    // If the indicator name is empty we clamp the first indicator attached to
    // this concept
    int ind_id = 0;

    if (!indicator_name.empty()) {
        if (!delphi::utils::in(this->indicators_in_CAG, indicator_name)) {
            print("Indicator \"{0}\" is not in CAG!\n", indicator_name);
            return;
        }

        if (!delphi::utils::in(n.nameToIndexMap, indicator_name)) {
            print("Indicator \"{0}\" is not attached to {1} in CAG!\n",
                    indicator_name,
                    concept_name);
            return;
        }

        ind_id = n.nameToIndexMap.at(indicator_name);
    }

    Indicator& ind = n.indicators[ind_id];

    // We have to clamp the latent state value corresponding to this
    // indicator such that the probability where the emission Gaussian
    // emitting the requested indicator value is the highest. For a
    // Gaussian emission function, the highest probable value is its
    // mean. So we have to set:
    //      μ = ind_value
    //      latent_clamp_value * scaling_factor = ind_value
    //      latent_clamp_value = ind_value / scaling_factor
    // NOTE: In our code we incorrectly call scaling_factor as
    //       indicator mean. To avoid confusion here, I am using the
    //       correct terminology.
    //       (Gosh, when we have incorrect terminology and we know it
    //       and we have not fixed it, I have to type a lot of
    //       comments)
    double latent_clamp_value = indicator_clamp_value / ind.get_mean();

    if (this->head_nodes.find(concept_id) == this->head_nodes.end()) {
      if (!delphi::utils::in(this->one_off_constraints, step)) {
        this->one_off_constraints[step] = vector<pair<int, double>>();
      }

      this->one_off_constraints[step].push_back(
          make_pair(concept_id, latent_clamp_value));
    } else {
      step += this->pred_start_timestep;
      if (!delphi::utils::in(this->head_node_one_off_constraints, step)) {
        this->head_node_one_off_constraints[step] = vector<pair<int, double>>();
      }

      this->head_node_one_off_constraints[step].push_back(
          make_pair(concept_id, latent_clamp_value));
    }
}

/*
 ============================================================================
 Public: Prediction
 ============================================================================
*/

Prediction AnalysisGraph::generate_prediction(int start_year,
                                              int start_month,
                                              int end_year,
                                              int end_month,
                                              ConstraintSchedule constraints,
                                              bool one_off,
                                              bool clamp_deri) {
  this->is_one_off_constraints = one_off;
  this->clamp_at_derivative = clamp_deri;
  this->one_off_constraints.clear();
  this->head_node_one_off_constraints.clear();

  for (auto [step, const_vec] : constraints) {
      for (auto constraint : const_vec) {
          string concept_name = get<0>(constraint);
          string indicator_name = get<1>(constraint);
          double value = get<2>(constraint);

          this->add_constraint(step, concept_name, indicator_name, value);
      }
  }

  this->run_model(start_year, start_month, end_year, end_month);

  return make_tuple(
      this->training_range, this->pred_range, this->format_prediction_result());
}

void AnalysisGraph::generate_prediction(int pred_start_timestep,
                                        int pred_timesteps,
                                        ConstraintSchedule constraints,
                                        bool one_off,
                                        bool clamp_deri) {
  this->is_one_off_constraints = one_off;
  this->clamp_at_derivative = clamp_deri;
  this->one_off_constraints.clear();
  this->head_node_one_off_constraints.clear();

  this->pred_start_timestep = pred_start_timestep - 1;
  this->pred_timesteps = pred_timesteps + 1;

  for (auto [step, const_vec] : constraints) {
    for (auto constraint : const_vec) {
      string concept_name = get<0>(constraint);
      string indicator_name = get<1>(constraint);
      double value = get<2>(constraint);

      this->add_constraint(step, concept_name, indicator_name, value);
    }
  }

  this->generate_latent_state_sequences(this->pred_start_timestep);
  this->generate_observed_state_sequences();
}

vector<vector<double>> AnalysisGraph::prediction_to_array(string indicator) {
  int vert_id = -1;
  int ind_id = -1;

  auto result =
      vector<vector<double>>(this->res, vector<double>(this->pred_timesteps));

  // Find the vertex id the indicator is attached to and
  // the indicator id of it.
  // TODO: We can make this more efficient by making indicators_in_CAG
  // a map from indicator names to vertices they are attached to.
  // This is just a quick and dirty implementation
  for (auto [v_name, v_id] : this->name_to_vertex) {
    for (auto [i_name, i_id] : (*this)[v_id].nameToIndexMap) {
      if (indicator.compare(i_name) == 0) {
        vert_id = v_id;
        ind_id = i_id;
        goto indicator_found;
      }
    }
  }
  // Program will reach here only if the indicator is not found
  throw IndicatorNotFoundException(format(
      "AnalysisGraph::prediction_to_array - indicator \"{}\" not found!\n",
      indicator));

indicator_found:

  for (int samp = 0; samp < this->res; samp++) {
    for (int ts = 0; ts < this->pred_timesteps; ts++) {
      result[samp][ts] =
          this->predicted_observed_state_sequences[samp][ts][vert_id][ind_id];
    }
  }

  return result;
}