//!
//! @file       HEmatrix.cpp
//! @brief      implementing functions for matrix operations
//!
//! @author     Miran Kim
//! @date       Dec. 1, 2017
//! @copyright  GNU Pub License
//!

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

#include <cmath>
#include <chrono>
#include <map>
#include <math.h>  // pow
#include <sys/time.h>
#include <iostream>
#include <stdio.h>
#include <stdlib.h>
#include <cstdlib>
#include <fstream>
#include <sstream>
#include <string>
#include <vector>
#include "math.h"
#include <cassert>
#include <random>
#include <string>
#include <iomanip>

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

#include "../src/Scheme.h"
#include "../src/SecretKey.h"
#include "../src/SchemeAlgo.h"

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

#define timing 0 // use for printing out the timing

//!@ Function: read the parameters
void readHEMatpar(HEMatpar& HEmatpar, long nrows, long ncols, long pBits, long cBits, long logQ, long subdim, long nbatching){
    HEmatpar.nrows = nrows;
    HEmatpar.ncols = ncols;

    long dim = nrows;
    if(dim < ncols) {
        dim = ncols;
    }

    HEmatpar.dim = (1<< (long)ceil(log2(dim)));  //! power of two
    HEmatpar.dim1 = HEmatpar.dim - 1;
    HEmatpar.logdim = (long) log2(HEmatpar.dim);
    HEmatpar.sqrdim = (long) ceil(sqrt(HEmatpar.dim));

    HEmatpar.pBits = pBits;
    HEmatpar.cBits = cBits;
    HEmatpar.logQ = logQ;

    HEmatpar.subdim = subdim;  //! used for non-squared matrix multiplication
    HEmatpar.nbatching = (1<< (long)ceil(log2(nbatching)));

    if(nbatching == 1){
        HEmatpar.nslots = HEmatpar.dim * HEmatpar.dim;  //! just encode a single matrix
    }
    else{
        HEmatpar.nslots = (HEmatpar.dim * HEmatpar.dim) * HEmatpar.nbatching;
    }
}

void HEmatrix::encryptRmat(Ciphertext& ctxt, mat_RR& mat, long logp){
    complex<double>* cmsg = new complex<double>[HEmatpar.nslots];

    NTL_EXEC_RANGE(HEmatpar.nrows, first, last);
    for(int i = first; i < last; ++i){
        for(long j = 0; j < HEmatpar.ncols; ++j){
            double dtemp;
            conv(dtemp, mat[i][j]);
            cmsg[i*HEmatpar.dim + j].real(dtemp);
        }
    }
    NTL_EXEC_RANGE_END;

    ctxt = scheme.encrypt(cmsg, HEmatpar.nslots, logp, HEmatpar.logQ);
    delete[] cmsg;
}

void HEmatrix::decryptRmat(mat_RR& mat, Ciphertext& ctxt){
    complex<double>* cmsg = scheme.decrypt(secretKey, ctxt);
    mat.SetDims(HEmatpar.dim, HEmatpar.dim);

    long k = 0;
    for(long i = 0; i < HEmatpar.dim; ++i){
        for(long j = 0; j < HEmatpar.dim; ++j){
            mat[i][j] = to_RR(cmsg[k].real());
            k++;
        }
    }
    delete[] cmsg;
}

void HEmatrix::encryptParallelRmat(Ciphertext& ctxt, mat_RR*& mat, long logp, long nbatching){
    complex<double>* cmsg = new complex<double>[HEmatpar.nslots];
    double* dtemp = new double[nbatching];

    NTL_EXEC_RANGE(nbatching, first, last);
    for(int k = first; k < last; ++k){
        // encode the kth matrix
        for(long i = 0; i < HEmatpar.nrows; ++i){
            for(long j = 0; j < HEmatpar.ncols; ++j){
                conv(dtemp[k], mat[k][i][j]);
                cmsg[(i*HEmatpar.dim + j)* HEmatpar.nbatching + k].real(dtemp[k]);
            }
        }
    }
    NTL_EXEC_RANGE_END;

    ctxt = scheme.encrypt(cmsg, HEmatpar.nslots, logp, HEmatpar.logQ);

    delete[] cmsg;
    delete[] dtemp;
}

void HEmatrix::decryptParallelRmat(mat_RR*& mat, Ciphertext& ctxt){
    complex<double>* cmsg = scheme.decrypt(secretKey, ctxt);
    mat = new mat_RR[HEmatpar.nbatching];

    for(long k = 0; k < HEmatpar.nbatching; ++k){
        mat[k].SetDims(HEmatpar.dim, HEmatpar.dim);
        long l = 0;
        for(long i = 0; i < HEmatpar.dim; ++i){
            for(long j = 0; j < HEmatpar.dim; ++j){
                mat[k][i][j] = to_RR(cmsg[l * HEmatpar.nbatching + k].real());
                l++;
            }
        }
    }
    delete[] cmsg;
}

void HEmatrix::msgleftRotate(complex<double>*& res, complex<double>* vals, long dim, long nrot){
    long nshift = (nrot)%HEmatpar.nslots;

    long k = dim - nshift;
    for(long j = 0; j < k; ++j){
        res[j] = vals[j + nshift];
    }
    for(long j = k; j < dim; ++j){
        res[j] = vals[j - k];
    }
}

void HEmatrix::msgrightRotate(complex<double>*& res, complex<double>* vals, long dim, long nrot){
    long nshift = (nrot)%HEmatpar.nslots;
    long k = dim - nshift;

    //! vals[0],vals[1],....,vals[d-nrot-1])
    for(long j = 0; j < k; ++j){
        res[nshift + j] = vals[j];
    }

    //! (vals[d-nrot],vals[d-nrot+1],...,vals[d-1])
    for(long j = k; j < dim; ++j){
        res[j - k] = vals[j];
    }
}

void HEmatrix::msgleftRotateAndEqual(complex<double>*& vals, long dim, long nrot){
    complex<double>* res = new complex<double>[dim];

    long nshift = (nrot)%HEmatpar.nslots;
    long k = dim - nshift;

    for(long j = 0; j < k; ++j){
        res[j] = vals[j + nshift];
    }
    for(long j = k; j < dim; ++j){
        res[j] = vals[j - k];
    }
    for(long j = 0; j < dim; ++j){
        vals[j] = res[j];
    }
}

void HEmatrix::msgrightRotateAndEqual(complex<double>*& vals, long dim, long nrot){
    complex<double>* res = new complex<double>[dim];

    long nshift = (nrot)%HEmatpar.nslots;
    long k = dim - nrot;

    for(long j = 0; j < k; ++j){
        res[nshift + j] = vals[j];
    }
    for(long j = k; j < dim; ++j){
        res[j - k] = vals[j];
    }
    for(long j = 0; j < dim; ++j){
        vals[j] = res[j];
    }
}

//------------------------------------------------
//! Transposition
//------------------------------------------------

//! Output: transpoly ("2*dim"), generate the polynomial for transpose
//! Originally we need to generate
//! rho(p[k]; -(d-1)*sqrt(r)i) where k = sqrt(r) * i + j, r = nrows, 0 <= i, j < sqrt(r)
// (0, 1,   ..., d-1),  (-, d+1, ..., 2d-1)

void HEmatrix::genTransPoly(ZZX*& transpoly){
    complex<double>** lvals   = new complex<double>*[HEmatpar.dim];
    complex<double>** rvals  = new complex<double>*[HEmatpar.dim];
    long dsquare = (HEmatpar.nslots) - 1;
    transpoly = new ZZX[2 * HEmatpar.dim];
    long dimsqrdim = HEmatpar.sqrdim * HEmatpar.dim1;   //! (d-1) * sqrt(d)

    // k = i * sqrt(d) + j < d
    bool btmp;
    if((HEmatpar.dim % HEmatpar.sqrdim) == 0){
        btmp = false;   //! all the terms have the same numbers
    }
    else{
        btmp = true;
    }
    long ibound = (long) ceil((double)HEmatpar.dim/HEmatpar.sqrdim); //! number of "i"

    NTL_EXEC_RANGE(ibound, first, last);
    for(int i = first; i < last; ++i){
        long jbound = HEmatpar.sqrdim;
        if ((btmp)&&(i == ibound - 1)){   //! last term
            jbound = (HEmatpar.dim % HEmatpar.sqrdim);
        }

        for(long j = 0; j < jbound; ++j){
            long k = i * HEmatpar.sqrdim + j;

            lvals[k] = new complex<double>[HEmatpar.nslots];
            rvals[k] = new complex<double>[HEmatpar.nslots];

            for(long l = 0; l < HEmatpar.dim - k; ++l){
                long dtemp = l * (HEmatpar.dim + 1) + k;
                lvals[k][dtemp].real(1.0);
                rvals[k][dsquare - dtemp].real(1.0);
            }

            msgrightRotateAndEqual(lvals[k], HEmatpar.nslots, i*dimsqrdim);  //! Lrho(P[k],-dsqrt(d)*i)
            msgleftRotateAndEqual(rvals[k], HEmatpar.nslots, i*dimsqrdim);   //! Rrho(P[k],-dsqrt(d)*i)

            transpoly[k] = scheme.context.encode(lvals[k], HEmatpar.nslots, HEmatpar.cBits);
            transpoly[k + HEmatpar.dim] = scheme.context.encode(rvals[k], HEmatpar.nslots, HEmatpar.cBits);
        }
    }
    NTL_EXEC_RANGE_END;
    delete[] lvals;
    delete[] rvals;
}

void HEmatrix::transpose(Ciphertext& res, Ciphertext& ctxt, ZZX*& transpoly){

    long dimsqrdim = HEmatpar.sqrdim * HEmatpar.dim1;   // (d-1) * sqrt(d)

    bool btmp;
    if((HEmatpar.dim % HEmatpar.sqrdim) == 0){
        btmp = false;   //! all the terms have the same numbers
    }
    else{
        btmp = true;
    }
    long ibound = (long) ceil((double)HEmatpar.dim/HEmatpar.sqrdim); //! number of "i"

    Ciphertext* Babyctxt1 = new Ciphertext[HEmatpar.sqrdim];
    Ciphertext* Babyctxt2 = new Ciphertext[HEmatpar.sqrdim];

    Ciphertext** ltemp = new Ciphertext*[HEmatpar.sqrdim];
    Ciphertext** rtemp = new Ciphertext*[HEmatpar.sqrdim];

    for(long i = 0; i < HEmatpar.sqrdim; ++i){
        ltemp[i] = new Ciphertext[HEmatpar.sqrdim];
        rtemp[i] = new Ciphertext[HEmatpar.sqrdim];
    }

    //! diagonal
    res = scheme.multByPoly(ctxt, transpoly[0], HEmatpar.cBits);

    Babyctxt1[0] = ctxt;
    Babyctxt2[0] = ctxt;

    //! Babyctxt1[j] = rho(v; j * (d-1))
    //! res[j] = rho(-, (d-1)*srt(d) * i) * rho(v; j * (d-1))
    //HEmatpar.sqrdim

    NTL_EXEC_RANGE(HEmatpar.sqrdim - 1, first, last);
    for(long j = first; j < last; ++j){
        long j1 = (j + 1);  // 1 <= j1 <= sqrdim - 1

        Babyctxt1[j1] = scheme.leftRotate(ctxt, j1 * (HEmatpar.dim1));   // prepare baby ctxt
        Babyctxt2[j1] = scheme.rightRotate(ctxt, j1 * (HEmatpar.dim1));

        ltemp[0][j1] = scheme.multByPoly(Babyctxt1[j1], transpoly[j1], HEmatpar.cBits);
        rtemp[0][j1] = scheme.multByPoly(Babyctxt2[j1], transpoly[j1 + HEmatpar.dim], HEmatpar.cBits);

        scheme.addAndEqual(ltemp[0][j1], rtemp[0][j1]);
    }
    NTL_EXEC_RANGE_END;

    for(long j = 1; j < HEmatpar.sqrdim; ++j){
        scheme.addAndEqual(res, ltemp[0][j]);
    }

    //! 1 <= i < d
    //! [k]: rho(p[k];-dsqrt(d)*i1) * rho(v, j*(d-1)) -> rot by dsqrt(d)*i
    //! res[j] = rho( rho(-, (d-1)*srt(d) * i) * rho(v; j * (d-1)) ; i * (d-1)\sqrt(d) )

    NTL_EXEC_RANGE(ibound - 1, first, last);
    for(long i = first; i < last; ++i){
        long jbound = HEmatpar.sqrdim;
        if (btmp &&(i == (ibound - 2))){
            jbound = (HEmatpar.dim % HEmatpar.sqrdim);
        }

        long i1 = i + 1;
        long k = (i1) * HEmatpar.sqrdim;

        //! 0 <= j < sqrdim
        NTL_EXEC_RANGE(jbound, first, last);
        for(long j = first; j < last; ++j){
            ltemp[i1][j] = scheme.multByPoly(Babyctxt1[j], transpoly[k + j], HEmatpar.cBits);
            rtemp[i1][j] = scheme.multByPoly(Babyctxt2[j], transpoly[k + j + HEmatpar.dim], HEmatpar.cBits);
        }
        NTL_EXEC_RANGE_END;

        for(long j = 1; j < jbound; ++j){
            scheme.addAndEqual(ltemp[i1][0], ltemp[i1][j]);
            scheme.addAndEqual(rtemp[i1][0], rtemp[i1][j]);
        }

        scheme.leftRotateAndEqual(ltemp[i1][0], i1 * dimsqrdim);
        scheme.rightRotateAndEqual(rtemp[i1][0], i1 * dimsqrdim);

        scheme.addAndEqual(ltemp[i1][0], rtemp[i1][0]);
    }
    NTL_EXEC_RANGE_END;

    for(long i = 1; i < ibound; ++i){
        scheme.addAndEqual(res, ltemp[i][0]);
    }

    scheme.reScaleByAndEqual(res, HEmatpar.cBits);   // Rescaling by scalarBits

    delete[] ltemp;
    delete[] rtemp;
}

void HEmatrix::genTransPoly_Parallel(ZZX*& transpoly){

    bool btmp;
    if((HEmatpar.dim % HEmatpar.sqrdim) == 0){ btmp = false; }    //! all the terms have the same numbers
    else{ btmp = true; }
    long ibound = (long) ceil((double)HEmatpar.dim/HEmatpar.sqrdim); //! number of "i"

    complex<double>** lvals   = new complex<double>*[HEmatpar.dim];
    complex<double>** rvals  = new complex<double>*[HEmatpar.dim];

    transpoly = new ZZX[2 * HEmatpar.dim];

    long shiftunit = HEmatpar.sqrdim * HEmatpar.dim1 * HEmatpar.nbatching ;   //! (d-1) * sqrt(d) * l
    long dsquare = ((HEmatpar.dim * HEmatpar.dim) - 1) *HEmatpar.nbatching ;  //! (d^2-1) * l

    //! k = (i * sqrt(d) + j)  < d
    //! poly[k] = poly[i * sqrt(d) + j]  -> rho(poly[k]; i * (d-1)*sqrt(d) * l)
    for(int i = 0; i < ibound; ++i){
        long jbound = HEmatpar.sqrdim;
        if ((btmp)&&(i == ibound - 1)){
            jbound = (HEmatpar.dim % HEmatpar.sqrdim);
        }
        for(long j = 0; j < jbound; ++j){
            long k = (i * HEmatpar.sqrdim + j);

            lvals[k] = new complex<double>[HEmatpar.nslots];
            rvals[k] = new complex<double>[HEmatpar.nslots];

            for(long l = 0; l < HEmatpar.dim - k; ++l){
                long dtemp= (l*(HEmatpar.dim+1) + k) * HEmatpar.nbatching;  //! starting index

                for(long n = 0; n < HEmatpar.nbatching; ++n){
                    lvals[k][dtemp + n].real(1.0);
                    rvals[k][dsquare - dtemp + n].real(1.0);
                }
            }

            //! Lrho(P[k], - i * (d-1) * sqrt(d) * (nbathcing))
            //! Rrho(P[k],- i * (d-1) * sqrt(d) * (nbathcing))
            msgrightRotateAndEqual(lvals[k], HEmatpar.nslots, i * shiftunit);
            msgleftRotateAndEqual(rvals[k], HEmatpar.nslots, i * shiftunit);

            transpoly[k] = scheme.context.encode(lvals[k], HEmatpar.nslots, HEmatpar.cBits);
            transpoly[k + HEmatpar.dim] = scheme.context.encode(rvals[k], HEmatpar.nslots, HEmatpar.cBits);

        }
    }
    delete[] lvals;
    delete[] rvals;
}

void HEmatrix::transpose_Parallel(Ciphertext& res, Ciphertext& ctxt, ZZX*& transpoly){
    long shiftunit = HEmatpar.sqrdim * HEmatpar.dim1 * HEmatpar.nbatching;  // (d-1) * sqrt(d) * l
    long unit  = (HEmatpar.dim1) * HEmatpar.nbatching;

    bool btmp;
    if((HEmatpar.dim % HEmatpar.sqrdim) == 0){ btmp = false; }    //! all the terms have the same numbers
    else{ btmp = true; }
    long ibound = (long) ceil((double)HEmatpar.dim/HEmatpar.sqrdim); //! number of "i"

    Ciphertext* Babyctxt1 = new Ciphertext[HEmatpar.sqrdim];
    Ciphertext* Babyctxt2 = new Ciphertext[HEmatpar.sqrdim];

    Ciphertext** ltemp = new Ciphertext*[HEmatpar.sqrdim];
    Ciphertext** rtemp = new Ciphertext*[HEmatpar.sqrdim];

    for(long i = 0; i < HEmatpar.sqrdim; ++i){
        ltemp[i] = new Ciphertext[HEmatpar.sqrdim];
        rtemp[i] = new Ciphertext[HEmatpar.sqrdim];
    }

    //-------------------------
    // i = 0, sqrt(d) polynomials
    //-------------------------
    res = scheme.multByPoly(ctxt, transpoly[0], HEmatpar.cBits);

    Babyctxt1[0] = ctxt;
    Babyctxt2[0] = ctxt;

    NTL_EXEC_RANGE(HEmatpar.sqrdim - 1, first, last);
    for(long j = first; j < last; ++j){
        long j1 = (j + 1);  //! 1 <= j1 <= sqrdim - 1

        // rho(v; j * (d-1) * l)
        Babyctxt1[j1] = scheme.leftRotate(ctxt, j1 * unit);
        Babyctxt2[j1] = scheme.rightRotate(ctxt, j1 * unit);

        ltemp[0][j1] = scheme.multByPoly(Babyctxt1[j1], transpoly[j1], HEmatpar.cBits);
        rtemp[0][j1] = scheme.multByPoly(Babyctxt2[j1], transpoly[j1 + HEmatpar.dim], HEmatpar.cBits);

        scheme.addAndEqual(ltemp[0][j1], rtemp[0][j1]);
    }
    NTL_EXEC_RANGE_END;

    for(long j = 1; j < HEmatpar.sqrdim; ++j){
        scheme.addAndEqual(res, ltemp[0][j]);
    }

    //! 1 <= i < d
    //! [k]: rho(p[k];-dsqrt(d)*i1) * rho(v, j*(d-1)) -> rot by dsqrt(d)*i

    NTL_EXEC_RANGE(ibound - 1, first, last);
    for(long i = first; i < last; ++i){
        long jbound = HEmatpar.sqrdim;
        if ((btmp)&&(i == ibound - 2)){
            jbound = (HEmatpar.dim % HEmatpar.sqrdim);
        }

        long i1 = i + 1;
        long k = (i1) * HEmatpar.sqrdim;

        NTL_EXEC_RANGE(jbound, first, last);
        for(long j = first; j < last; ++j){
            ltemp[i1][j] = scheme.multByPoly(Babyctxt1[j], transpoly[k + j], HEmatpar.cBits);
            rtemp[i1][j] = scheme.multByPoly(Babyctxt2[j], transpoly[k + j + HEmatpar.dim], HEmatpar.cBits);
        }
        NTL_EXEC_RANGE_END;

        for(long j = 1; j < jbound; ++j){
            scheme.addAndEqual(ltemp[i1][0], ltemp[i1][j]);
            scheme.addAndEqual(rtemp[i1][0], rtemp[i1][j]);
        }

        scheme.leftRotateAndEqual(ltemp[i1][0], i1 * shiftunit);
        scheme.rightRotateAndEqual(rtemp[i1][0], i1 * shiftunit);

        scheme.addAndEqual(ltemp[i1][0], rtemp[i1][0]);
    }
    NTL_EXEC_RANGE_END;

    for(long i = 1; i < ibound; ++i){
        scheme.addAndEqual(res, ltemp[i][0]);
    }

    scheme.reScaleByAndEqual(res, HEmatpar.cBits);   //! Rescaling by scalarBits

    delete[] ltemp;
    delete[] rtemp;
}

//------------------------------------------------
//! Shift
//------------------------------------------------

// num = 0: dimension d-1
// subdim - 1: dimension (subdim-1)
void HEmatrix::genShiftPoly(ZZX*& shiftpoly, long num){
    long length = HEmatpar.dim1;
    if(num != 0){
        length = num;
    }

    complex<double>** vals = new complex<double>*[length];
    shiftpoly = new ZZX[length];

    // i: shifted by (i+1)
    NTL_EXEC_RANGE(length, first, last);
    for(int i = first; i < last; ++i){
        vals[i] = new complex<double>[HEmatpar.nslots];

        for(long j = 0; j < HEmatpar.dim; ++j){
            for(long k = 0; k < i + 1; ++k){
                vals[i][j * HEmatpar.dim + k].real(1.0);
            }
        }
        shiftpoly[i] = scheme.context.encode(vals[i], HEmatpar.nslots, HEmatpar.cBits);
    }
    NTL_EXEC_RANGE_END;

    delete[] vals;
}

void HEmatrix::shiftBycols(Ciphertext& res, Ciphertext& ctxt,  long k , ZZX*& shiftpoly){
    //! ctemp = Enc(m[1],.., m[k],0,,,0 | ....)
    Ciphertext ctemp = scheme.multByPoly(ctxt, shiftpoly[k-1], HEmatpar.cBits);
    scheme.reScaleByAndEqual(ctemp, HEmatpar.cBits);

    //! ctemp = Enc(0,,,0, m[k+1],...,m[d] | ....)
    res = scheme.modDownTo(ctxt, ctemp.logq);
    scheme.subAndEqual(res, ctemp);

    scheme.rightRotateAndEqual(ctemp, HEmatpar.dim - k);
    scheme.leftRotateAndEqual(res, k);
    scheme.addAndEqual(res, ctemp);
}

void HEmatrix::genShiftPoly_Parallel(ZZX*& shiftpoly){
    complex<double>** vals   = new complex<double>*[HEmatpar.dim1];
    shiftpoly = new ZZX[HEmatpar.dim1];

    NTL_EXEC_RANGE(HEmatpar.dim1, first, last);
    for(int i = first; i < last; ++i){
        vals[i] = new complex<double>[HEmatpar.nslots];

        for(long j = 0; j < HEmatpar.dim; ++j){
            for(long k = 0; k < i + 1; ++k){
                long dtemp = (j * HEmatpar.dim + k) * HEmatpar.nbatching;

                for(long n = 0; n < HEmatpar.nbatching; ++n){
                    vals[i][dtemp + n].real(1.0);
                }
            }
        }
        shiftpoly[i] = scheme.context.encode(vals[i], HEmatpar.nslots, HEmatpar.cBits);
    }
    NTL_EXEC_RANGE_END;

    delete[] vals;
}

void HEmatrix::shiftBycols_Parallel(Ciphertext& res, Ciphertext& ctxt, long k, ZZX*& shiftpoly){
    //! ctemp = Enc(m[1],.., m[k],0,,,0 | ....)
    Ciphertext ctemp = scheme.multByPoly(ctxt, shiftpoly[k-1], HEmatpar.cBits);
    scheme.reScaleByAndEqual(ctemp, HEmatpar.cBits);

    //! ctemp = Enc(0,,,0, m[k+1],...,m[d] | ....)
    res = scheme.modDownTo(ctxt, ctemp.logq);
    scheme.subAndEqual(res, ctemp);

    scheme.rightRotateAndEqual(ctemp, (HEmatpar.dim - k) * HEmatpar.nbatching);
    scheme.leftRotateAndEqual(res, k * HEmatpar.nbatching);
    scheme.addAndEqual(res, ctemp);
}

//------------------------------------------------
//! Matrix multiplication
//------------------------------------------------

//!@ Output: Initpoly[0] (constant left-polynomials for Amat): 0,1,....,d-1
//!@         Initpoly[1] (constant right-polynomials for Amat):  -1,...,-(d-1)
//!!         Initpoly[2] (constant polynomials for Bmat)

void HEmatrix::genMultPoly(ZZX**& Initpoly){
    bool btmp;
    if((HEmatpar.dim % HEmatpar.sqrdim) == 0){
        btmp = false;
    }
    else{
        btmp = true;
    }
    long ibound = (long) ceil((double) HEmatpar.dim/HEmatpar.sqrdim); //! number of "i"

    Initpoly = new ZZX*[3];

    Initpoly[0] =  new ZZX[HEmatpar.dim];
    Initpoly[1] =  new ZZX[HEmatpar.dim];
    Initpoly[2] =  new ZZX[HEmatpar.dim];

    complex<double>** fvals1   = new complex<double>*[HEmatpar.dim];
    complex<double>** fvals2   = new complex<double>*[HEmatpar.dim];
    complex<double>** bvals    = new complex<double>*[HEmatpar.dim];

    NTL_EXEC_RANGE(ibound, first, last);
    for(int i = first; i < last; ++i){
        long jbound = HEmatpar.sqrdim;
        if ((btmp)&&(i == ibound - 1)){
            jbound = (HEmatpar.dim % HEmatpar.sqrdim);
        }

        for(long j = 0; j < jbound; ++j){
            long k = i * HEmatpar.sqrdim + j;
            fvals1[k] = new complex<double>[HEmatpar.nslots];
            fvals2[k] = new complex<double>[HEmatpar.nslots];
            bvals[k] = new complex<double>[HEmatpar.nslots];

            for(long l = 0; l < HEmatpar.dim - k; ++l){
                fvals1[k][k * HEmatpar.dim + l].real(1.0);
            }
            msgleftRotate(fvals2[k], fvals1[k], HEmatpar.nslots, k*(2*HEmatpar.dim - 1));

            msgrightRotateAndEqual(fvals1[k], HEmatpar.nslots, i*HEmatpar.sqrdim);
            Initpoly[0][k] = scheme.context.encode(fvals1[k], HEmatpar.nslots, HEmatpar.cBits);

            msgleftRotateAndEqual(fvals2[k], HEmatpar.nslots, i*HEmatpar.sqrdim);
            Initpoly[1][k] = scheme.context.encode(fvals2[k], HEmatpar.nslots, HEmatpar.cBits);

            for(long l = 0; l < HEmatpar.dim; ++l){
                bvals[k][l * HEmatpar.dim + k].real(1.0);
            }
            msgrightRotateAndEqual(bvals[k], HEmatpar.nslots, i*HEmatpar.sqrdim*HEmatpar.dim);
            Initpoly[2][k] = scheme.context.encode(bvals[k], HEmatpar.nslots, HEmatpar.cBits);
        }
    }
    NTL_EXEC_RANGE_END;

    delete[] fvals1;
    delete[] fvals2;
    delete[] bvals;
}

void HEmatrix::genInitCtxt(Ciphertext& resA, Ciphertext& resB, Ciphertext& Actxt, Ciphertext& Bctxt, ZZX**& poly){
    bool btmp;
    if((HEmatpar.dim % HEmatpar.sqrdim) == 0){ btmp = false; }    //! all the terms have the same numbers
    else{ btmp = true; }
    long ibound = (long) ceil((double)HEmatpar.dim/HEmatpar.sqrdim); //! number of "i"

    Ciphertext** Actemp1 = new Ciphertext*[HEmatpar.sqrdim]; //! update right polynomial
    Ciphertext** Actemp2 = new Ciphertext*[HEmatpar.sqrdim]; //! update right polynomial
    Ciphertext** Bctemp = new Ciphertext*[HEmatpar.sqrdim];

    for(long i = 0; i < HEmatpar.sqrdim; ++i){
        Actemp1[i] = new Ciphertext[HEmatpar.sqrdim];
        Actemp2[i] = new Ciphertext[HEmatpar.sqrdim];
        Bctemp[i]  = new Ciphertext[HEmatpar.sqrdim];
    }

    //! 0. Store some ciphertexts (0,1,...,d-1), (,d+1,...2d-1)
    //! v, lrho(v;1), lrho(v;2), ..., lrho(v;d-1)
    //! -, rrho(v;1), rrho(v;2), ..., rrho(v;d-1)

    Ciphertext* BaByctxt1  = new Ciphertext[HEmatpar.dim];
    Ciphertext* BaByctxt2  = new Ciphertext[HEmatpar.dim];
    Ciphertext* BaByctxtB  = new Ciphertext[HEmatpar.dim];

    BaByctxt1[0] = Actxt;
    BaByctxt2[0] = Actxt;
    BaByctxtB[0] = Bctxt;

    //! i = 0:   Actxts[0] = v[0] + p1 * v[1] + ... + p[sqr(d)-1] *  v[sqr(d)-1]
    resA = scheme.multByPoly(BaByctxt1[0], poly[0][0], HEmatpar.cBits);
    resB = scheme.multByPoly(BaByctxtB[0], poly[2][0], HEmatpar.cBits);

    NTL_EXEC_RANGE(HEmatpar.sqrdim - 1, first, last);
    for(long j = first; j < last; ++j){
        long j1 = (j + 1);
        BaByctxt1[j1] = scheme.leftRotate(Actxt, j1);
        BaByctxt2[j1] = scheme.rightRotate(Actxt, j1);
        BaByctxtB[j1] = scheme.leftRotate(Bctxt, (j1)*HEmatpar.dim);

        Actemp1[0][j1] = scheme.multByPoly(BaByctxt1[j1], poly[0][j1], HEmatpar.cBits);
        Actemp2[0][j1] = scheme.multByPoly(BaByctxt2[j1], poly[1][j1], HEmatpar.cBits);
        Bctemp[0][j1]  = scheme.multByPoly(BaByctxtB[j1], poly[2][j1], HEmatpar.cBits);

        scheme.addAndEqual(Actemp1[0][j1], Actemp2[0][j1]);
    }
    NTL_EXEC_RANGE_END;

    for(long j = 1; j < HEmatpar.sqrdim; ++j){
        scheme.addAndEqual(resA, Actemp1[0][j]);
        scheme.addAndEqual(resB, Bctemp[0][j]);
    }

    Ciphertext* Actxts1 = new Ciphertext[HEmatpar.dim];
    Ciphertext* Actxts2 = new Ciphertext[HEmatpar.dim];
    Ciphertext* Bctxts  = new Ciphertext[HEmatpar.dim];

    NTL_EXEC_RANGE(HEmatpar.dim - HEmatpar.sqrdim, first, last);
    for(long k = first; k < last; ++k){
        long k1 = (k + HEmatpar.sqrdim);
        long i = (long)(k1 / HEmatpar.sqrdim);
        long j = (long)(k1 % HEmatpar.sqrdim);

        Actxts1[k1] = scheme.multByPoly(BaByctxt1[j], poly[0][k1], HEmatpar.cBits);
        Actxts2[k1] = scheme.multByPoly(BaByctxt2[j], poly[1][k1], HEmatpar.cBits);
        Bctxts[k1]  = scheme.multByPoly(BaByctxtB[j], poly[2][k1], HEmatpar.cBits);

    }
    NTL_EXEC_RANGE_END;

    NTL_EXEC_RANGE(ibound - 1, first, last);
    for(long k = first; k < last; ++k){
        long k1 = (k+1) * HEmatpar.sqrdim;
        long jbound = HEmatpar.sqrdim;
        if((btmp) && (k == ibound - 2)){
            jbound = (HEmatpar.dim % HEmatpar.sqrdim);
        }
        for(long j = 1; j < jbound; ++j){
            scheme.addAndEqual(Actxts1[k1], Actxts1[k1+j]);
            scheme.addAndEqual(Actxts2[k1], Actxts2[k1+j]);
            scheme.addAndEqual(Bctxts[k1], Bctxts[k1+j]);
        }

        long k2 = (k1 - (k1 % HEmatpar.sqrdim));
        scheme.leftRotateAndEqual(Actxts1[k1], k2);
        scheme.rightRotateAndEqual(Actxts2[k1], k2);
        scheme.leftRotateAndEqual(Bctxts[k1], k2 * HEmatpar.dim);

        scheme.addAndEqual(Actxts1[k1], Actxts2[k1]);
    }
    NTL_EXEC_RANGE_END;

    for(long k = 1; k < ibound; ++k){
        long k1 = k * HEmatpar.sqrdim;
        scheme.addAndEqual(resA, Actxts1[k1]);
        scheme.addAndEqual(resB, Bctxts[k1]);
    }

    scheme.reScaleByAndEqual(resA, HEmatpar.cBits);
    scheme.reScaleByAndEqual(resB, HEmatpar.cBits);

    delete[] Actxts1;
    delete[] Actxts2;
    delete[] Bctxts;

    delete[] Actemp1;
    delete[] Actemp2;
    delete[] Bctemp;

    delete[] BaByctxt1;
    delete[] BaByctxt2;
    delete[] BaByctxtB;
}

void HEmatrix::HEmatmul_Hadamard(Ciphertext& res, Ciphertext* Actxts, Ciphertext* Bctxts, long num){
    //! logq: Actxt > Actxts[0] = Actxts[i + 1] + 2*cBits
    //! logq: Bctxt > Bctxts[0] = Bctxt[i + 1]

    long num1 = num - 1;
    //! 1) Bctxts[0] < Actxts[1] < Actxts[0]
    if(Actxts[1].logq > Bctxts[0].logq){
        res = Actxts[0];
        scheme.modDownToAndEqual(res, Bctxts[0].logq);
        scheme.multAndEqual(res, Bctxts[0]);   //! log(res) = log Bctxt[0]

        NTL_EXEC_RANGE(num1, first, last);
        for(int i = first; i < last; ++i){
            long i1 = (i + 1);
            scheme.modDownToAndEqual(Actxts[i1], Bctxts[0].logq);
            scheme.multAndEqual(Actxts[i1], Bctxts[i1]);   //! log(Actxt[i+1]) = log Bctxt[0]
        }
        NTL_EXEC_RANGE_END;
    }

    //! 2)  Actxts[1] < Actxts[0] < Bctxts[0]
    else if(Actxts[0].logq < Bctxts[0].logq){
        res = Bctxts[0];
        scheme.modDownToAndEqual(res, Actxts[0].logq);
        scheme.multAndEqual(res, Actxts[0]);   //! log(res) = log Actxt[0]

        NTL_EXEC_RANGE(num1, first, last);
        for(int i = first; i < last; ++i){
            long i1 = (i + 1);
            scheme.modDownToAndEqual(Bctxts[i1], Actxts[i1].logq);
            scheme.multAndEqual(Actxts[i1], Bctxts[i1]);
        }
        NTL_EXEC_RANGE_END;

        scheme.modDownToAndEqual(res, Actxts[1].logq);
    }

    //! 3)  Actxts[1] < Bctxts[i] < Actxts[0]: the common case
    else{
        res = Actxts[0];
        scheme.modDownToAndEqual(res, Bctxts[0].logq);
        scheme.multAndEqual(res, Bctxts[0]);   //! log(res) = log Bctxt[0]

        NTL_EXEC_RANGE(num1, first, last);
        for(int i = first; i < last; ++i){
            long i1 = (i + 1);
            scheme.modDownToAndEqual(Bctxts[i1], Actxts[i1].logq);
            scheme.multAndEqual(Actxts[i1], Bctxts[i1]);
        }
        NTL_EXEC_RANGE_END;

        scheme.modDownToAndEqual(res, Actxts[1].logq);
    }

    //! aggregate the results and rescaling
    for(int i = 1; i < num; ++i){
        scheme.addAndEqual(res, Actxts[i]);
    }
    scheme.reScaleByAndEqual(res, res.logp);

}

void HEmatrix::HEmatmul(Ciphertext& res, Ciphertext& Actxt, Ciphertext& Bctxt, ZZX**& Initpoly, ZZX*& shiftpoly){
    Ciphertext* Actxts = new Ciphertext[HEmatpar.dim];
    Ciphertext* Bctxts = new Ciphertext[HEmatpar.dim];

    //! 1. Generate the initial ciphertexts
    genInitCtxt(Actxts[0], Bctxts[0], Actxt, Bctxt, Initpoly);

    //! 2. Column shifting of Actxt[0], Row shifting of Bctxt[0]
    NTL_EXEC_RANGE(HEmatpar.dim1, first, last);
    for(int i = first; i < last; ++i){
        long i1 = (i + 1);
        shiftBycols(Actxts[i1], Actxts[0], i1, shiftpoly);
        Bctxts[i1] = scheme.leftRotate(Bctxts[0], HEmatpar.dim * (i1));
    }
    NTL_EXEC_RANGE_END;

    //! 3. Hadamard multiplication
    HEmatmul_Hadamard(res, Actxts, Bctxts, HEmatpar.dim);

    delete[] Actxts;
    delete[] Bctxts;
}

void HEmatrix::genMultPoly_Parallel(ZZX**& Initpoly){
    long btmp;
    if((HEmatpar.dim % HEmatpar.sqrdim) == 0){ btmp = false; }    //! all the terms have the same numbers
    else{ btmp = true; }
    long ibound = (long) ceil((double)HEmatpar.dim/HEmatpar.sqrdim); //! number of "i"

    Initpoly = new ZZX*[3];

    Initpoly[0] =  new ZZX[HEmatpar.dim];
    Initpoly[1] =  new ZZX[HEmatpar.dim];
    Initpoly[2] =  new ZZX[HEmatpar.dim];

    complex<double>** fvals1   = new complex<double>*[HEmatpar.dim];
    complex<double>** fvals2   = new complex<double>*[HEmatpar.dim];
    complex<double>** bvals    = new complex<double>*[HEmatpar.dim];

    NTL_EXEC_RANGE(ibound, first, last);
    for(int i = first; i < last; ++i){
        long jbound = HEmatpar.sqrdim;
        if ((btmp) && (i == (ibound - 1))){
            jbound = (HEmatpar.dim % HEmatpar.sqrdim);
        }

        for(long j = 0; j < jbound; ++j){
            long k = i * HEmatpar.sqrdim + j;
            fvals1[k] = new complex<double>[HEmatpar.nslots];
            fvals2[k] = new complex<double>[HEmatpar.nslots];
            bvals[k] = new complex<double>[HEmatpar.nslots];

            //! original: k * d <= index <= k * d + d - k -1
            //! parallel: (k * d) * l + 0 <= ... <= (k * d) * l + (l-1)
            //!           (k * d + d - k -1) * l + 0 <= ... <= (k * d + d - k -1) * l + (l-1) <  (k * d + d - k -1) * l + l
            //!           from k * d * l <= index <  (k * d + d - k) * l
            long start = k * HEmatpar.dim * HEmatpar.nbatching;
            long end = ( k * HEmatpar.dim + HEmatpar.dim - k) * HEmatpar.nbatching;
            for(long l = start; l < end; ++l){
                fvals1[k][l].real(1.0);
            }

            msgleftRotate(fvals2[k], fvals1[k], HEmatpar.nslots, k*(2*HEmatpar.dim - 1) * HEmatpar.nbatching);

            msgrightRotateAndEqual(fvals1[k], HEmatpar.nslots, i*HEmatpar.sqrdim * HEmatpar.nbatching);
            Initpoly[0][k] = scheme.context.encode(fvals1[k], HEmatpar.nslots, HEmatpar.cBits);

            msgleftRotateAndEqual(fvals2[k], HEmatpar.nslots, i*HEmatpar.sqrdim * HEmatpar.nbatching);
            Initpoly[1][k] = scheme.context.encode(fvals2[k], HEmatpar.nslots, HEmatpar.cBits);

            for(long l = 0; l < HEmatpar.dim; ++l){
                long dtemp = (l * HEmatpar.dim + k) * HEmatpar.nbatching;

                for(long n = 0; n < HEmatpar.nbatching; ++n){
                    bvals[k][dtemp + n].real(1.0);
                }
            }
            msgrightRotateAndEqual(bvals[k], HEmatpar.nslots, i*HEmatpar.sqrdim*HEmatpar.dim * HEmatpar.nbatching);
            Initpoly[2][k] = scheme.context.encode(bvals[k], HEmatpar.nslots, HEmatpar.cBits);
        }
    }
    NTL_EXEC_RANGE_END;

    delete[] fvals1;
    delete[] fvals2;
    delete[] bvals;
}

void HEmatrix::genInitCtxt_Parallel(Ciphertext& resA, Ciphertext& resB, Ciphertext& Actxt, Ciphertext& Bctxt, ZZX**& poly){
    bool btmp;
    if((HEmatpar.dim % HEmatpar.sqrdim) == 0){ btmp = false; }    //! all the terms have the same numbers
    else{ btmp = true; }
    long ibound = (long) ceil((double)HEmatpar.dim/HEmatpar.sqrdim); //! number of "i"

    Ciphertext** Actemp1 = new Ciphertext*[HEmatpar.sqrdim]; //! update right polynomial
    Ciphertext** Actemp2 = new Ciphertext*[HEmatpar.sqrdim]; //! update right polynomial
    Ciphertext** Bctemp = new Ciphertext*[HEmatpar.sqrdim];

    for(long i = 0; i < HEmatpar.sqrdim; ++i){
        Actemp1[i] = new Ciphertext[HEmatpar.sqrdim];
        Actemp2[i] = new Ciphertext[HEmatpar.sqrdim];
        Bctemp[i]  = new Ciphertext[HEmatpar.sqrdim];
    }

    //---------------------------------------------------------------
    // 0. Store some ciphertexts (0,1,...,d-1), ( ,d+1,...2d-1)
    // v, lrho(v;1), lrho(v;2), ..., lrho(v;d-1)
    // -, rrho(v;1), rrho(v;2), ..., rrho(v;d-1)

    Ciphertext* BaByctxt1  = new Ciphertext[HEmatpar.dim];
    Ciphertext* BaByctxt2  = new Ciphertext[HEmatpar.dim];
    Ciphertext* BaByctxtB  = new Ciphertext[HEmatpar.dim];

    BaByctxt1[0] = Actxt;
    BaByctxt2[0] = Actxt;
    BaByctxtB[0] = Bctxt;

    resA = scheme.multByPoly(BaByctxt1[0], poly[0][0], HEmatpar.cBits);
    resB = scheme.multByPoly(BaByctxtB[0], poly[2][0], HEmatpar.cBits);

    NTL_EXEC_RANGE(HEmatpar.sqrdim - 1, first, last);
    for(long j = first; j < last; ++j){
        long j1 = (j + 1);
        BaByctxt1[j1] = scheme.leftRotate(Actxt, j1 * HEmatpar.nbatching);
        BaByctxt2[j1] = scheme.rightRotate(Actxt, j1 * HEmatpar.nbatching);
        BaByctxtB[j1] = scheme.leftRotate(Bctxt, (j1)*HEmatpar.dim * HEmatpar.nbatching);

        Actemp1[0][j1] = scheme.multByPoly(BaByctxt1[j1], poly[0][j1], HEmatpar.cBits);
        Actemp2[0][j1] = scheme.multByPoly(BaByctxt2[j1], poly[1][j1], HEmatpar.cBits);
        Bctemp[0][j1]  = scheme.multByPoly(BaByctxtB[j1], poly[2][j1], HEmatpar.cBits);

        scheme.addAndEqual(Actemp1[0][j1], Actemp2[0][j1]);
    }
    NTL_EXEC_RANGE_END;

    for(long j = 1; j < HEmatpar.sqrdim; ++j){
        scheme.addAndEqual(resA, Actemp1[0][j]);
        scheme.addAndEqual(resB, Bctemp[0][j]);
    }

    //---------------------------
    Ciphertext* Actxts1 = new Ciphertext[HEmatpar.dim];
    Ciphertext* Actxts2 = new Ciphertext[HEmatpar.dim];
    Ciphertext* Bctxts  = new Ciphertext[HEmatpar.dim];

    NTL_EXEC_RANGE(HEmatpar.dim - ibound, first, last);
    for(long k = first; k < last; ++k){
        long k1 = (k + HEmatpar.sqrdim);
        long i = (long)(k1 / HEmatpar.sqrdim);
        long j = (long)(k1 % HEmatpar.sqrdim);

        Actxts1[k1] = scheme.multByPoly(BaByctxt1[j], poly[0][k1], HEmatpar.cBits);
        Actxts2[k1] = scheme.multByPoly(BaByctxt2[j], poly[1][k1], HEmatpar.cBits);
        Bctxts[k1]  = scheme.multByPoly(BaByctxtB[j], poly[2][k1], HEmatpar.cBits);

    }
    NTL_EXEC_RANGE_END;

    NTL_EXEC_RANGE(ibound - 1, first, last);
    for(long k = first; k < last; ++k){
        long k1 = (k+1) * HEmatpar.sqrdim;

        long jbound = HEmatpar.sqrdim;
        if((btmp) && (k == ibound - 2)){
            jbound = (HEmatpar.dim % HEmatpar.sqrdim);
        }
        for(long j = 1; j < jbound; ++j){
            scheme.addAndEqual(Actxts1[k1], Actxts1[k1+j]);
            scheme.addAndEqual(Actxts2[k1], Actxts2[k1+j]);
            scheme.addAndEqual(Bctxts[k1], Bctxts[k1+j]);
        }

        long k2 = (k1 - (k1 % HEmatpar.sqrdim)) * HEmatpar.nbatching;
        scheme.leftRotateAndEqual(Actxts1[k1], k2);
        scheme.rightRotateAndEqual(Actxts2[k1], k2);
        scheme.leftRotateAndEqual(Bctxts[k1], k2 * HEmatpar.dim);

        scheme.addAndEqual(Actxts1[k1], Actxts2[k1]);
    }
    NTL_EXEC_RANGE_END;

    for(long k = 1; k < ibound; ++k){
        long k1 = k * HEmatpar.sqrdim;
        scheme.addAndEqual(resA, Actxts1[k1]);
        scheme.addAndEqual(resB, Bctxts[k1]);
    }

    scheme.reScaleByAndEqual(resA, HEmatpar.cBits);
    scheme.reScaleByAndEqual(resB, HEmatpar.cBits);

    delete[] Actxts1;
    delete[] Actxts2;
    delete[] Bctxts;

    delete[] Actemp1;
    delete[] Actemp2;
    delete[] Bctemp;

    delete[] BaByctxt1;
    delete[] BaByctxt2;
    delete[] BaByctxtB;
}

void HEmatrix::HEmatmul_Parallel(Ciphertext& res, Ciphertext& Actxt, Ciphertext& Bctxt, ZZX**& Initpoly, ZZX*& shiftpoly){
    Ciphertext* Actxts = new Ciphertext[HEmatpar.dim];
    Ciphertext* Bctxts = new Ciphertext[HEmatpar.dim];

    //! 1. Generate the initial ciphertexts
    genInitCtxt_Parallel(Actxts[0], Bctxts[0], Actxt, Bctxt, Initpoly);

    //! 2. Column shifting of Actxt[0], Row shifting of Bctxt[0]
    long unit = HEmatpar.dim  * HEmatpar.nbatching;
    NTL_EXEC_RANGE(HEmatpar.dim1, first, last);
    for(int i = first; i < last; ++i){
        long i1 = (i + 1);
        shiftBycols_Parallel(Actxts[i1], Actxts[0], i1, shiftpoly);
        Bctxts[i1] = scheme.leftRotate(Bctxts[0], unit * (i1));
    }
    NTL_EXEC_RANGE_END;

    //! 3. Hadamard mult : Actxts[0] * Bctxts[0] + ... + Actxts[d-1] * Bctxts[d-1]
    HEmatmul_Hadamard(res, Actxts, Bctxts, HEmatpar.dim);

    delete[] Actxts;
    delete[] Bctxts;
}

void HEmatrix::HErmatmul(Ciphertext& res, Ciphertext& Actxt, Ciphertext& Bctxt, ZZX**& Initpoly, ZZX*& shiftpoly){
    Ciphertext* Actxts = new Ciphertext[HEmatpar.subdim];
    Ciphertext* Bctxts = new Ciphertext[HEmatpar.subdim];

    //! 1. Generate the initial ciphertexts
    genInitCtxt(Actxts[0], Bctxts[0], Actxt, Bctxt, Initpoly);

    //! 2. Column shifting of Actxt[0], Row shifting of Bctxt[0]
    NTL_EXEC_RANGE(HEmatpar.subdim - 1, first, last);
    for(int i = first; i < last; ++i){
        long i1 = (i + 1);
        shiftBycols(Actxts[i1], Actxts[0], i1, shiftpoly);
        Bctxts[i1] = scheme.leftRotate(Bctxts[0], HEmatpar.dim * (i1));
    }
    NTL_EXEC_RANGE_END;

    //! 3. Hadamard mult : Actxts[0] * Bctxts[0] + ... + Actxts[d-1] * Bctxts[d-1]
    HEmatmul_Hadamard(res, Actxts, Bctxts, HEmatpar.subdim);

    //! 4. shift and aggregate the results
    long index = (long) log2(HEmatpar.dim/HEmatpar.subdim);

    for(long i = 0; i < index; ++i){
        Ciphertext ctemp = scheme.leftRotate(res, HEmatpar.dim*HEmatpar.subdim * (1<<i));
        scheme.addAndEqual(res, ctemp);
    }

    delete[] Actxts;
    delete[] Bctxts;
}

void HEmatrix::genMultBPoly(ZZX*& Initpoly){
    bool btmp;
    if((HEmatpar.dim % HEmatpar.sqrdim) == 0){ btmp = false; }    //! all the terms have the same numbers
    else{ btmp = true; }
    long ibound = (long) ceil((double)HEmatpar.dim/HEmatpar.sqrdim); //! number of "i"

    Initpoly =  new ZZX[HEmatpar.dim];
    complex<double>** bvals  = new complex<double>*[HEmatpar.dim];

    NTL_EXEC_RANGE(ibound, first, last);
    for(int i = first; i < last; ++i){
        long jbound = HEmatpar.sqrdim;
        if ((btmp)&&(i == ibound - 1)){
            jbound = (HEmatpar.dim % HEmatpar.sqrdim);
        }
        for(long j = 0; j < jbound; ++j){
            long k = i * HEmatpar.sqrdim + j;
            bvals[k] = new complex<double>[HEmatpar.nslots];

            for(long l = 0; l < HEmatpar.dim; ++l){
                bvals[k][l * HEmatpar.dim + k].real(1.0);
            }
            msgrightRotateAndEqual(bvals[k], HEmatpar.nslots, i*HEmatpar.sqrdim*HEmatpar.dim);
            Initpoly[k] = scheme.context.encode(bvals[k], HEmatpar.nslots, HEmatpar.cBits);
        }
    }
    NTL_EXEC_RANGE_END;

    delete[] bvals;
}

void HEmatrix::genInitActxt(Ciphertext*& Actxts, Mat<RR>& mat){
    Mat<RR>* Amat = new Mat<RR>[HEmatpar.dim];
    Actxts = new Ciphertext[HEmatpar.dim];
    complex<double>** cmsg = new complex<double>*[HEmatpar.dim];

    NTL_EXEC_RANGE(HEmatpar.dim, first, last);
    for(long k = first; k < last; ++k){
        Amat[k].SetDims(HEmatpar.dim, HEmatpar.dim);
        long dimk = HEmatpar.dim - k;

        //! 0 <= i < d - k: shift by (k+i)-positions from mat[i]
        for(long i = 0; i < dimk; ++i){
            long nshift = k + i;
            long nshift2 = HEmatpar.dim - nshift;

            for(long j = 0; j < nshift2; ++j){
                Amat[k][i][j] = mat[i][j + nshift];
            }

            for(long j = nshift2; j < HEmatpar.dim; ++j){
                Amat[k][i][j] = mat[i][j - nshift2];
            }
        }

        //! i = d - k : Amat[k][i] <- mat[i]
        if(k!=0){
            for(long j = 0; j < HEmatpar.dim; ++j){
                Amat[k][dimk][j] = mat[dimk][j];
            }
        }

        //! d - k + 1 <= i < d: shift by (k+i-d)-positions from mat[i]
        for(long i = dimk + 1; i < HEmatpar.dim; ++i){
            long nshift =  i - dimk;
            long nshift2 = HEmatpar.dim - nshift;

            for(long j = 0; j < nshift2; ++j){
                Amat[k][i][j] = mat[i][j + nshift];
            }
            for(long j = nshift2; j < HEmatpar.dim; ++j){
                Amat[k][i][j] = mat[i][j - nshift2];
            }
        }
    }
    NTL_EXEC_RANGE_END;

    //! encryption of d Rmat
    NTL_EXEC_RANGE(HEmatpar.dim, first, last);
    for(long k = first; k < last; ++k){
        cmsg[k] = new complex<double>[HEmatpar.nslots];

        for(long i = 0; i < HEmatpar.nrows; ++i){
            for(long j = 0; j < HEmatpar.ncols; ++j){
                double dtemp;
                conv(dtemp, Amat[k][i][j]);
                cmsg[k][i*HEmatpar.dim + j].real(dtemp);
            }
        }
        Actxts[k] = scheme.encrypt(cmsg[k], HEmatpar.nslots, HEmatpar.pBits, HEmatpar.logQ);
    }
    NTL_EXEC_RANGE_END;

    delete[] cmsg;
}

void HEmatrix::genInitBctxt(Ciphertext& resB, Ciphertext& Bctxt, ZZX*& poly){
    bool btmp;
    if((HEmatpar.dim % HEmatpar.sqrdim) == 0){
        btmp = false;
    }
    else{
        btmp = true;
    }
    long ibound = (long) ceil((double)HEmatpar.dim/HEmatpar.sqrdim);

    Ciphertext** Bctemp = new Ciphertext*[HEmatpar.sqrdim];
    for(long i = 0; i < HEmatpar.sqrdim; ++i){
        Bctemp[i]  = new Ciphertext[HEmatpar.sqrdim];
    }

    //! 0. Store some ciphertexts (0,1,...,d-1), ( ,d+1,...2d-1)
    Ciphertext* BaByctxtB  = new Ciphertext[HEmatpar.dim];

    BaByctxtB[0] = Bctxt;

    //! i = 0:   Actxts[0] = v[0] + p1 * v[1] + ... + p[sqr(d)-1] *  v[sqr(d)-1]
    resB = scheme.multByPoly(BaByctxtB[0], poly[0], HEmatpar.cBits);

    NTL_EXEC_RANGE(HEmatpar.sqrdim - 1, first, last);
    for(long j = first; j < last; ++j){
        long j1 = (j + 1);
        BaByctxtB[j1] = scheme.leftRotate(Bctxt, (j1)*HEmatpar.dim);
        Bctemp[0][j1]  = scheme.multByPoly(BaByctxtB[j1], poly[j1], HEmatpar.cBits);
    }
    NTL_EXEC_RANGE_END;

    for(long j = 1; j < HEmatpar.sqrdim; ++j){
        scheme.addAndEqual(resB, Bctemp[0][j]);
    }

    Ciphertext* Bctxts  = new Ciphertext[HEmatpar.dim];

    NTL_EXEC_RANGE(HEmatpar.dim - HEmatpar.sqrdim, first, last);
    for(long k = first; k < last; ++k){
        long k1 = (k + HEmatpar.sqrdim);
        long i = (long)(k1 / HEmatpar.sqrdim);
        long j = (long)(k1 % HEmatpar.sqrdim);

        Bctxts[k1]  = scheme.multByPoly(BaByctxtB[j], poly[k1], HEmatpar.cBits);
    }
    NTL_EXEC_RANGE_END;

    NTL_EXEC_RANGE(ibound - 1, first, last);
    for(long k = first; k < last; ++k){
        long k1 = (k+1) * HEmatpar.sqrdim;
        long jbound = HEmatpar.sqrdim;
        if((btmp) && (k == ibound - 2)){
            jbound = (HEmatpar.dim % HEmatpar.sqrdim);
        }

        for(long j = 1; j < jbound; ++j){
            scheme.addAndEqual(Bctxts[k1], Bctxts[k1+j]);
        }

        long k2 = (k1 - (k1 % HEmatpar.sqrdim));
        scheme.leftRotateAndEqual(Bctxts[k1], k2 * HEmatpar.dim);
    }
    NTL_EXEC_RANGE_END;

    for(long k = 1; k < ibound; ++k){
        long k1 = k * HEmatpar.sqrdim;
        scheme.addAndEqual(resB, Bctxts[k1]);
    }

    scheme.reScaleByAndEqual(resB, HEmatpar.cBits);

    delete[] Bctxts;
    delete[] Bctemp;
    delete[] BaByctxtB;
}

void HEmatrix::HEmatmul_preprocessing(Ciphertext& res, Ciphertext*& Actxts, Ciphertext& Bctxt, ZZX*& Initpoly){

    //! 1. Generate the initial ciphertexts
    Ciphertext* Bctxts = new Ciphertext[HEmatpar.dim];
    genInitBctxt(Bctxts[0], Bctxt, Initpoly);

    //! 2. Row shifting of Bctxt[0]
    NTL_EXEC_RANGE(HEmatpar.dim1, first, last);
    for(int i = first; i < last; ++i){
        long i1 = (i + 1);
        Bctxts[i1] = scheme.leftRotate(Bctxts[0], HEmatpar.dim * (i1));
    }
    NTL_EXEC_RANGE_END;

    //! 3. Hadamard mult : Actxts[0] * Bctxts[0] + ... + Actxts[d-1] * Bctxts[d-1]
    NTL_EXEC_RANGE(HEmatpar.dim, first, last);
    for(int i = first; i < last; ++i){
        scheme.modDownToAndEqual(Actxts[i], Bctxts[0].logq);
        scheme.multAndEqual(Actxts[i], Bctxts[i]);   // log(Actxt[i+1]) = log Bctxt[0]
    }
    NTL_EXEC_RANGE_END;

    res = Actxts[0];
    for(int i = 1; i < HEmatpar.dim; ++i){
        scheme.addAndEqual(res, Actxts[i]);
    }

    scheme.reScaleByAndEqual(res, res.logp);

    delete[] Bctxts;
}

void HEmatrix::genInitRecActxt(Ciphertext*& Actxts, Mat<RR>& mat){
    // rep_mat: (mat; mat; ... ;mat) square mat
    Mat<RR> replicate_mat;
    replicate_mat.SetDims(HEmatpar.dim, HEmatpar.dim);

    long index_rows = HEmatpar.dim/HEmatpar.subdim;

    NTL_EXEC_RANGE(index_rows, first, last);
    for(long k = first; k < last; ++k){
        for(long i = 0; i < HEmatpar.subdim; ++i){
            for(long j = 0; j < HEmatpar.dim; ++j){
                replicate_mat[k*HEmatpar.subdim + i][j] = mat[i][j];
            }
        }
    }
    NTL_EXEC_RANGE_END;

    //! generate the (linear transformed) matrices
    Mat<RR>* Amat = new Mat<RR>[HEmatpar.subdim];
    Actxts = new Ciphertext[HEmatpar.subdim];

    NTL_EXEC_RANGE(HEmatpar.subdim, first, last);
    for(long k = first; k < last; ++k){
        Amat[k].SetDims(HEmatpar.dim, HEmatpar.dim);
        long dimk= HEmatpar.dim - k;

        //! 0 <= i < d - k: shift by (k+i)-positions from mat[i]
        for(long i = 0; i < dimk; ++i){
            long nshift = k + i;
            long nshift2 = HEmatpar.dim - nshift;

            for(long j = 0; j < nshift2; ++j){
                Amat[k][i][j] = replicate_mat[i][j + nshift];
            }

            for(long j = nshift2; j < HEmatpar.dim; ++j){
                Amat[k][i][j] = replicate_mat[i][j - nshift2];
            }
        }

        //! i = d - k : Amat[k][i] <- mat[i]
        if(k!=0){
            for(long j = 0; j < HEmatpar.dim; ++j){
                Amat[k][dimk][j] = replicate_mat[dimk][j];
            }
        }

        //! d - k + 1 <= i < d: shift by (k+i-d)-positions from mat[i]
        for(long i = dimk + 1; i < HEmatpar.dim; ++i){
            long nshift =  i - dimk;
            long nshift2 = HEmatpar.dim - nshift;

            for(long j = 0; j < nshift2; ++j){
                Amat[k][i][j] = replicate_mat[i][j + nshift];
            }
            for(long j = nshift2; j < HEmatpar.dim; ++j){
                Amat[k][i][j] = replicate_mat[i][j - nshift2];
            }
        }
    }
    NTL_EXEC_RANGE_END;

    //! encryption of d Rmat
    complex<double>** cmsg = new complex<double>*[HEmatpar.subdim];
    NTL_EXEC_RANGE(HEmatpar.subdim, first, last);
    for(long k = first; k < last; ++k){
        cmsg[k] = new complex<double>[HEmatpar.nslots];
        for(int i = 0; i < HEmatpar.nrows; ++i){
            for(long j = 0; j < HEmatpar.ncols; ++j){
                double dtemp;
                conv(dtemp, Amat[k][i][j]);
                cmsg[k][i*HEmatpar.dim + j].real(dtemp);
            }
        }
        Actxts[k] = scheme.encrypt(cmsg[k], HEmatpar.nslots, HEmatpar.pBits, HEmatpar.logQ);
    }
    NTL_EXEC_RANGE_END;

    delete[] cmsg;
}

void HEmatrix::HErmatmul_preprocessing(Ciphertext& res, Ciphertext*& Actxts, Ciphertext& Bctxt, ZZX*& Initpoly){

    //! 1. Generate the initial ciphertexts
    Ciphertext* Bctxts = new Ciphertext[HEmatpar.subdim];
    genInitBctxt(Bctxts[0], Bctxt, Initpoly);

    //! 2. Row shifting of Bctxt[0]
    NTL_EXEC_RANGE(HEmatpar.subdim - 1, first, last);
    for(int i = first; i < last; ++i){
        long i1 = (i + 1);
        Bctxts[i1] = scheme.leftRotate(Bctxts[0], HEmatpar.dim * (i1));
    }
    NTL_EXEC_RANGE_END;

    //! 3. Hadamard mult : Actxts[0] * Bctxts[0] + ... + Actxts[d-1] * Bctxts[d-1]
    NTL_EXEC_RANGE(HEmatpar.subdim, first, last);
    for(int i = first; i < last; ++i){
        scheme.modDownToAndEqual(Actxts[i], Bctxts[0].logq);
        scheme.multAndEqual(Actxts[i], Bctxts[i]);   //! log(Actxt[i+1]) = log Bctxt[0]
    }
    NTL_EXEC_RANGE_END;

    //! aggregate the results
    res = Actxts[0];
    for(int i = 1; i < HEmatpar.subdim; ++i){
        scheme.addAndEqual(res, Actxts[i]);
    }
    scheme.reScaleByAndEqual(res, res.logp);

    //! 4. shift and aggregate the results
    long index = (long) log2(HEmatpar.dim/HEmatpar.subdim);
    for(long i = 0; i < index; ++i){
        Ciphertext ctemp = scheme.leftRotate(res, HEmatpar.dim*HEmatpar.subdim * (1<<i));
        scheme.addAndEqual(res, ctemp);
    }

    delete[] Bctxts;
}


// ----------------------------------------------------------------------------
// -                        Open3D: www.open3d.org                            -
// ----------------------------------------------------------------------------
// The MIT License (MIT)
//
// Copyright (c) 2018 www.open3d.org
//
// Permission is hereby granted, free of charge, to any person obtaining a copy
// of this software and associated documentation files (the "Software"), to deal
// in the Software without restriction, including without limitation the rights
// to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
// copies of the Software, and to permit persons to whom the Software is
// furnished to do so, subject to the following conditions:
//
// The above copyright notice and this permission notice shall be included in
// all copies or substantial portions of the Software.
//
// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
// IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
// FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
// AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
// LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING
// FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS
// IN THE SOFTWARE.
// ----------------------------------------------------------------------------

#include "Feature.h"

#include <Eigen/Dense>
#include <Open3D/Utility/Console.h>
#include <Open3D/Geometry/PointCloud.h>
#include <Open3D/Geometry/KDTreeFlann.h>

namespace open3d {

namespace {

Eigen::Vector4d ComputePairFeatures(const Eigen::Vector3d &p1,
                                    const Eigen::Vector3d &n1,
                                    const Eigen::Vector3d &p2,
                                    const Eigen::Vector3d &n2) {
    Eigen::Vector4d result;
    Eigen::Vector3d dp2p1 = p2 - p1;
    result(3) = dp2p1.norm();
    if (result(3) == 0.0) {
        return Eigen::Vector4d::Zero();
    }
    auto n1_copy = n1;
    auto n2_copy = n2;
    double angle1 = n1_copy.dot(dp2p1) / result(3);
    double angle2 = n2_copy.dot(dp2p1) / result(3);
    if (acos(fabs(angle1)) > acos(fabs(angle2))) {
        n1_copy = n2;
        n2_copy = n1;
        dp2p1 *= -1.0;
        result(2) = -angle2;
    } else {
        result(2) = angle1;
    }
    auto v = dp2p1.cross(n1_copy);
    double v_norm = v.norm();
    if (v_norm == 0.0) {
        return Eigen::Vector4d::Zero();
    }
    v /= v_norm;
    auto w = n1_copy.cross(v);
    result(1) = v.dot(n2_copy);
    result(0) = atan2(w.dot(n2_copy), n1_copy.dot(n2_copy));
    return result;
}

std::shared_ptr<Feature> ComputeSPFHFeature(
        const PointCloud &input,
        const KDTreeFlann &kdtree,
        const KDTreeSearchParam &search_param) {
    auto feature = std::make_shared<Feature>();
    feature->Resize(33, (int)input.points_.size());
#ifdef _OPENMP
#pragma omp parallel for schedule(static)
#endif
    for (int i = 0; i < (int)input.points_.size(); i++) {
        const auto &point = input.points_[i];
        const auto &normal = input.normals_[i];
        std::vector<int> indices;
        std::vector<double> distance2;
        if (kdtree.Search(point, search_param, indices, distance2) > 1) {
            // only compute SPFH feature when a point has neighbors
            double hist_incr = 100.0 / (double)(indices.size() - 1);
            for (size_t k = 1; k < indices.size(); k++) {
                // skip the point itself, compute histogram
                auto pf = ComputePairFeatures(point, normal,
                                              input.points_[indices[k]],
                                              input.normals_[indices[k]]);
                int h_index = (int)(floor(11 * (pf(0) + M_PI) / (2.0 * M_PI)));
                if (h_index < 0) h_index = 0;
                if (h_index >= 11) h_index = 10;
                feature->data_(h_index, i) += hist_incr;
                h_index = (int)(floor(11 * (pf(1) + 1.0) * 0.5));
                if (h_index < 0) h_index = 0;
                if (h_index >= 11) h_index = 10;
                feature->data_(h_index + 11, i) += hist_incr;
                h_index = (int)(floor(11 * (pf(2) + 1.0) * 0.5));
                if (h_index < 0) h_index = 0;
                if (h_index >= 11) h_index = 10;
                feature->data_(h_index + 22, i) += hist_incr;
            }
        }
    }
    return feature;
}

}  // unnamed namespace

std::shared_ptr<Feature> ComputeFPFHFeature(
        const PointCloud &input,
        const KDTreeSearchParam &search_param /* = KDTreeSearchParamKNN()*/) {
    auto feature = std::make_shared<Feature>();
    feature->Resize(33, (int)input.points_.size());
    if (input.HasNormals() == false) {
        PrintDebug(
                "[ComputeFPFHFeature] Failed because input point cloud has no "
                "normal.\n");
        return feature;
    }
    KDTreeFlann kdtree(input);
    auto spfh = ComputeSPFHFeature(input, kdtree, search_param);
#ifdef _OPENMP
#pragma omp parallel for schedule(static)
#endif
    for (int i = 0; i < (int)input.points_.size(); i++) {
        const auto &point = input.points_[i];
        std::vector<int> indices;
        std::vector<double> distance2;
        if (kdtree.Search(point, search_param, indices, distance2) > 1) {
            double sum[3] = {0.0, 0.0, 0.0};
            for (size_t k = 1; k < indices.size(); k++) {
                // skip the point itself
                double dist = distance2[k];
                if (dist == 0.0) continue;
                for (int j = 0; j < 33; j++) {
                    double val = spfh->data_(j, indices[k]) / dist;
                    sum[j / 11] += val;
                    feature->data_(j, i) += val;
                }
            }
            for (int j = 0; j < 3; j++)
                if (sum[j] != 0.0) sum[j] = 100.0 / sum[j];
            for (int j = 0; j < 33; j++) {
                feature->data_(j, i) *= sum[j / 11];
                // The commented line is the fpfh function in the paper.
                // But according to PCL implementation, it is skipped.
                // Our initial test shows that the full fpfh function in the
                // paper seems to be better than PCL implementation. Further
                // test required.
                feature->data_(j, i) += spfh->data_(j, i);
            }
        }
    }
    return feature;
}

}  // namespace open3d


#ifndef GEOMETRY_HPP
#define GEOMETRY_HPP

#include "point.hpp"

#include <boost/assert.hpp>

#include <vector>
#include <algorithm>

namespace geometry
{
    enum class monoticity : unsigned char
    {
        INCREASING_X = 1,
        INCREASING_Y = 2,
        DECREASING_X = 4,
        DECREASING_Y = 8,
        CONSTANT_X   = 5,
        CONSTANT_Y   = 10,
        CONSTANT     = 15,
        INVALID      = 0

    };

    enum class point_position : char
    {
        LEFT_OF_LINE,
        RIGHT_OF_LINE,
        ON_LINE
    };

    inline monoticity operator&(monoticity lhs, monoticity rhs)
    {
        return static_cast<geometry::monoticity>(static_cast<char>(lhs) & static_cast<char>(rhs));
    }

    /// Transforms the path from the goven monoticity to be x-monotone-increasing
    inline void make_x_monotone_increasing(monoticity mono, std::vector<coordinate>& path)
    {
        BOOST_ASSERT(mono != monoticity::INVALID);

        if ((mono & monoticity::INCREASING_X) != monoticity::INVALID)
        {
            return;
        }

        if ((mono & monoticity::DECREASING_X) == monoticity::INVALID)
        {
            // make x-monotone: swap x and y (mirror on (0,0)->(1,1))
            if ((mono & monoticity::INCREASING_Y) != monoticity::INVALID ||
                (mono & monoticity::DECREASING_Y) != monoticity::INVALID)
            {
                for (unsigned i = 0; i < path.size(); i++)
                {
                    std::swap(path[i].y, path[i].x);
                }
            }

            // we are now x-monotone increasing
            if ((mono & monoticity::INCREASING_Y) != monoticity::INVALID)
            {
                return;
            }
        }

        // at this point we are always x-monotone descreasing: Mirror on y-Axis
        std::transform(path.begin(), path.end(), path.begin(),
                       [](coordinate c)
                       {
                            c.x *= -1;
                            return c;
                       });
    }

    /// Returns a vector that is orthogonal to the line implied by
    /// first_line_point and second_line_point.
    ///
    /// This vector is not normalized!
    inline coordinate line_normal(const coordinate& first_line_point,
                                  const coordinate& second_line_point)
    {
        return coordinate {-second_line_point.y + first_line_point.y,
                            second_line_point.x - first_line_point.x};
    }

    inline coordinate::value_type cross(const coordinate& a, const coordinate& b)
    {
        return a.x * b.y - a.y * b.x;
    }

    struct intersection_params
    {
        double first_param;
        double second_param;
        bool colinear;
    };

    inline intersection_params segment_intersection(
            const coordinate& first_segment_a, const coordinate& first_segment_b,
            const coordinate& second_segment_a, const coordinate& second_segment_b)
    {
        intersection_params params {0, 0, false};
        const coordinate first_delta = first_segment_b - first_segment_a;
        const coordinate second_delta = second_segment_b - second_segment_a;
        auto direction_cross = cross(first_delta, second_delta);
        // colinear
        if (direction_cross == 0)
        {
            params.colinear = true;
        }
        else
        {
            params.first_param  = cross((second_segment_a - first_segment_a), second_delta) / direction_cross;
            params.second_param = cross((second_segment_a - first_segment_a), first_delta) / direction_cross;
        }

        return params;
    }

    inline bool segments_intersect(const coordinate& first_segment_a, const coordinate& first_segment_b,
                                   const coordinate& second_segment_a, const coordinate& second_segment_b)
    {
        auto params = segment_intersection(first_segment_a, first_segment_b, second_segment_a, second_segment_b);
        auto u = params.first_param;
        auto t = params.second_param;
        if (params.colinear)
        {
            return false;
        }

        return (u >= 0 && u <= 1.0) && (t >= 0 && t <= 1.0);
    }

    inline point_position position_to_line(const coordinate& first_line_point,
                                    const coordinate& second_line_point,
                                    const coordinate& point)
    {
        auto delta = point - first_line_point;
        auto p = glm::dot(line_normal(first_line_point, second_line_point), delta);

        if (p > 0)
            return point_position::LEFT_OF_LINE;
        else if (p < 0)
            return point_position::RIGHT_OF_LINE;

        return point_position::ON_LINE;
    }

    /// returns an angle beteen 0 and 2*M_PI
    inline float normalize_angle(float angle)
    {
        float normalized = angle;
        while (normalized < 0) normalized      += 2*M_PI;
        while (normalized > 2*M_PI) normalized -= 2*M_PI;

        return normalized;
    }

    /// returns the shortest angle between two vectors with the given angles
    inline float angle_diff(float first_angle, float second_angle)
    {
        float diff = std::abs(normalize_angle(first_angle) - normalize_angle(second_angle));
        if (diff > M_PI)
        {
            diff = diff - M_PI;
        }
        return diff;
    }

    /// compares the slope of origin -> lhs and origin -> rhs and returns true
    /// if the slope of lhs is bigger than rhs.
    inline bool slope_compare(const coordinate& origin, const coordinate& lhs, const coordinate& rhs)
    {
        BOOST_ASSERT(lhs.x >= origin.x);
        BOOST_ASSERT(rhs.x >= origin.x);

        auto position = position_to_line(origin, lhs, rhs);

        return position == point_position::RIGHT_OF_LINE ||
              (position == point_position::ON_LINE &&
               // origin -> rhs points in opposite direction
               glm::dot(lhs - origin, rhs - origin) < 0);
    }
};
#endif


/*
 * reclayer.cpp
 *
 * Feed-forward:
 *  state:
 *    zz_h(t) = ww_hh * hh(t-1) + ww_xh*x(t) + b_h
 *    hh(t)   = state_activation_function(zz_h(t))
 *  output:
 *    zz_a(t) = ww_ha * hh(t) + b_a
 *    a(t)    = output_activation_function(zz_a(t))
 *
 * Back propagation:
 *   output:
 *    delta_a(t)  = elem_prod(gradient(C, a(t)), activation_derivative(zz_a(t)))
 *    dC/d b_a(t)   = delta_a(t)
 *    dC/d ww_ha(t) = hh(t) * delta_a(t)
 *    total_gradient(C, hh(t)) = gradient(C, hh(t)) + transp(ww_ha) * delta_a(t)
 *  state:
 *    delta_h(t) = elem_prod(total_gradient(C, hh(t)), activation_derivative(zz_h(t)))
 *    dC/d b_h(t) = delta_h(t)
 *    dC/d ww_xh(t) = x(t) * delta_h(t)
 *    dC/d ww_hh(t) = hh(t-1) * delta_h(t)
 *    gradient(C, hh(t-1)) = transp(ww_hh) * delta_h(t)
 *    gradient(C, x(t)) = transp(ww_xh) * delta_h(t)
 */
#include <boost/assert.hpp>

#include "core/utils.h"
#include "core/random.h"
#include "core/functions.h"
#include "reclayer.h"

using namespace std;
using namespace boost;
using namespace yann;

namespace yann {

////////////////////////////////////////////////////////////////////////////////////////////////
//
// RecurrentLayer_Context implementation
//
class RecurrentLayer_Context :
    public Layer::Context
{
  typedef Layer::Context Base;

  friend class RecurrentLayer;

public:
  RecurrentLayer_Context(
      const MatrixSize & state_size,
      const MatrixSize & output_size,
      const MatrixSize & max_batch_size) :
    Base(output_size, max_batch_size),
    _pos(0)
  {
    init(state_size);
  }

  RecurrentLayer_Context(
      const MatrixSize & state_size,
      const RefVectorBatch & output) :
    Base(output),
    _pos(0)
  {
    init(state_size);
  }

  // Layer::Context overwrites
  virtual void reset_state()
  {
    Base::reset_state();

    _pos = 0;
  }

private:
  void init(const MatrixSize & state_size)
  {
    YANN_CHECK_GT(state_size, 0);

    _hh.resize(get_batch_size(), state_size);
    _zz_h.resize(get_batch_size(), state_size);
    _zz_a.resize(get_batch_size(), get_output_size());
  }

protected:
  MatrixSize _pos;
  VectorBatch  _hh;       // state (save it for each step)
  VectorBatch  _zz_h;
  VectorBatch  _zz_a;
}; // class RecurrentLayer_Context

////////////////////////////////////////////////////////////////////////////////////////////////
//
// RecurrentLayer_TrainingContext implementation
//
class RecurrentLayer_TrainingContext :
    public RecurrentLayer_Context
{
  typedef RecurrentLayer_Context Base;
  friend class RecurrentLayer;

public:
  RecurrentLayer_TrainingContext(
      const MatrixSize & input_size,
      const MatrixSize & state_size,
      const MatrixSize & output_size,
      const MatrixSize & batch_size,
      const unique_ptr<Layer::Updater> & updater) :
    Base(state_size, output_size, batch_size),
    _ww_hh_updater(updater->copy()),
    _ww_xh_updater(updater->copy()),
    _bb_h_updater(updater->copy()),
    _ww_ha_updater(updater->copy()),
    _bb_a_updater(updater->copy())
  {
    init(input_size, state_size);
  }

  RecurrentLayer_TrainingContext(
      const MatrixSize & input_size,
      const MatrixSize & state_size,
      const RefVectorBatch & output,
      const unique_ptr<Layer::Updater> & updater) :
    Base(state_size, output),
    _ww_hh_updater(updater->copy()),
    _ww_xh_updater(updater->copy()),
    _bb_h_updater(updater->copy()),
    _ww_ha_updater(updater->copy()),
    _bb_a_updater(updater->copy())
  {
    init(input_size, state_size);
  }

  // Layer::Context overwrites
  virtual void start_epoch()
  {
    YANN_SLOW_CHECK(_ww_hh_updater);
    YANN_SLOW_CHECK(_ww_xh_updater);
    YANN_SLOW_CHECK(_bb_h_updater);
    YANN_SLOW_CHECK(_ww_ha_updater);
    YANN_SLOW_CHECK(_bb_a_updater);

    Base::start_epoch();

    _ww_hh_updater->start_epoch();
    _ww_xh_updater->start_epoch();
    _bb_h_updater->start_epoch();
    _ww_ha_updater->start_epoch();
    _bb_a_updater->start_epoch();
  }

  virtual void reset_state()
  {
    YANN_SLOW_CHECK(_ww_hh_updater);
    YANN_SLOW_CHECK(_ww_xh_updater);
    YANN_SLOW_CHECK(_bb_h_updater);
    YANN_SLOW_CHECK(_ww_ha_updater);
    YANN_SLOW_CHECK(_bb_a_updater);

    Base::reset_state();

    _gradient_h.setZero();

    _delta_ww_hh.setZero();
    _delta_ww_xh.setZero();
    _delta_bb_h.setZero();
    _delta_ww_ha.setZero();
    _delta_bb_a.setZero();

    _ww_hh_updater->reset();
    _ww_xh_updater->reset();
    _bb_h_updater->reset();
    _ww_ha_updater->reset();
    _bb_a_updater->reset();
  }

private:
  void init(const MatrixSize & input_size, const MatrixSize & state_size)
  {
    YANN_CHECK_GT(input_size, 0);
    YANN_CHECK_GT(state_size, 0);

    _delta_ww_hh.resize(state_size, state_size);
    _delta_ww_xh.resize(input_size, state_size);
    _delta_bb_h.resize(state_size);
    _delta_ww_ha.resize(state_size, get_output_size());
    _delta_bb_a.resize(get_output_size());

    _gradient_h.resize(state_size);
    _delta_h.resize(state_size);
    _delta_a.resize(get_output_size());

    _sigma_derivative_zz_h.resize(state_size);
    _sigma_derivative_zz_a.resize(get_output_size());

    _ww_hh_updater->init(_delta_ww_hh.rows(), _delta_ww_hh.cols());
    _ww_xh_updater->init(_delta_ww_xh.rows(), _delta_ww_xh.cols());
    _bb_h_updater->init(1, _delta_bb_h.size()); // RowMajor
    _ww_ha_updater->init(_delta_ww_ha.rows(), _delta_ww_ha.cols());
    _bb_a_updater->init(1, _delta_bb_a.size()); // RowMajor
  }

private:
  Matrix _delta_ww_hh;
  Matrix _delta_ww_xh;
  Vector _delta_bb_h;
  Matrix _delta_ww_ha;
  Vector _delta_bb_a;

  Vector _gradient_h;
  Vector _delta_h;
  Vector _delta_a;

  Vector _sigma_derivative_zz_h;
  Vector _sigma_derivative_zz_a;

  unique_ptr<Layer::Updater> _ww_hh_updater;
  unique_ptr<Layer::Updater> _ww_xh_updater;
  unique_ptr<Layer::Updater> _bb_h_updater;
  unique_ptr<Layer::Updater> _ww_ha_updater;
  unique_ptr<Layer::Updater> _bb_a_updater;
}; // class RecurrentLayer_TrainingContext

}; // namespace yann

////////////////////////////////////////////////////////////////////////////////////////////////
//
// yann::RecurrentLayer implementation
//
yann::RecurrentLayer::RecurrentLayer(
    const MatrixSize & input_size,
    const MatrixSize & state_size,
    const MatrixSize & output_size) :
    _state_activation_function(new SigmoidFunction()),
    _output_activation_function(new SigmoidFunction())
{
  YANN_CHECK_GT(input_size, 0);
  YANN_CHECK_GT(state_size, 0);
  YANN_CHECK_GT(output_size, 0);

  _ww_hh.resize(state_size, state_size);
  _ww_xh.resize(input_size, state_size);
  _bb_h.resize(state_size);

  _ww_ha.resize(state_size, output_size);
  _bb_a.resize(output_size);
}

yann::RecurrentLayer::~RecurrentLayer()
{
}

void yann::RecurrentLayer::set_values(
    const Matrix & ww_hh, const Matrix & ww_xh, const Vector & bb_h,
    const Matrix & ww_ha, const Vector & bb_a)
{
  YANN_CHECK(is_same_size(ww_hh, _ww_hh));
  YANN_CHECK(is_same_size(ww_xh, _ww_xh));
  YANN_CHECK(is_same_size(bb_h, _bb_h));
  YANN_CHECK(is_same_size(ww_ha, _ww_ha));
  YANN_CHECK(is_same_size(bb_a, _bb_a));
  _ww_hh = ww_hh;
  _ww_xh = ww_xh;
  _bb_h  = bb_h;
  _ww_ha = ww_ha;
  _bb_a  = bb_a;
}

void yann::RecurrentLayer::set_activation_functions(
    const std::unique_ptr<ActivationFunction> & state_activation_function,
    const std::unique_ptr<ActivationFunction> & output_activation_function)
{
  YANN_CHECK(state_activation_function);
  YANN_CHECK(output_activation_function);
  _state_activation_function = state_activation_function->copy();
  _output_activation_function = output_activation_function->copy();
}

// Layer overwrites
bool yann::RecurrentLayer::is_valid() const
{
  if(!Base::is_valid()) {
    return false;
  }
  if(!_state_activation_function || !_output_activation_function) {
    return false;
  }
  if(_ww_hh.rows() != _ww_hh.cols()) {
    return false;
  }
  if(_ww_xh.cols() != _ww_hh.cols()) {
    return false;
  }
  if(_ww_hh.cols() != _bb_h.size()) {
    return false;
  }
  if(_ww_hh.cols() != _ww_ha.rows()) {
    return false;
  }
  if(_ww_ha.cols() != _bb_a.size()) {
    return false;
  }

  return true;
}

std::string yann::RecurrentLayer::get_name() const
{
  return "RecurrentLayer";
}

string yann::RecurrentLayer::get_info() const
{
  YANN_CHECK(is_valid());

  ostringstream oss;
  oss << Base::get_info()
      << " state size: " << get_state_size()
      << ", state activation: " << _state_activation_function->get_info()
      << ", output activation: " << _output_activation_function->get_info()
      ;
  return oss.str();
}

bool yann::RecurrentLayer::is_equal(const Layer & other, double tolerance) const
{
  if(!Base::is_equal(other, tolerance)) {
    return false;
  }
  auto the_other = dynamic_cast<const RecurrentLayer*>(&other);
  if(the_other == nullptr) {
    return false;
  }
  // TOOD: add deep compare
  if(_state_activation_function->get_info() != the_other->_state_activation_function->get_info()) {
    return false;
  }
  if(_output_activation_function->get_info() != the_other->_output_activation_function->get_info()) {
    return false;
  }

  if(!_ww_hh.isApprox(the_other->_ww_hh, tolerance)) {
    return false;
  }
  if(!_ww_xh.isApprox(the_other->_ww_xh, tolerance)) {
    return false;
  }
  if(!_bb_h.isApprox(the_other->_bb_h, tolerance)) {
    return false;
  }
  if(!_ww_ha.isApprox(the_other->_ww_ha, tolerance)) {
    return false;
  }
  if(!_bb_a.isApprox(the_other->_bb_a, tolerance)) {
    return false;
  }
  return true;
}

MatrixSize yann::RecurrentLayer::get_input_size() const
{
  YANN_SLOW_CHECK(is_valid());
  return _ww_xh.rows();
}

MatrixSize yann::RecurrentLayer::get_state_size() const
{
  YANN_SLOW_CHECK(is_valid());
  return _ww_xh.cols();
}

MatrixSize yann::RecurrentLayer::get_output_size() const
{
  YANN_SLOW_CHECK(is_valid());
  return _ww_ha.cols();
}

unique_ptr<Layer::Context> yann::RecurrentLayer::create_context(const MatrixSize & batch_size) const
{
  YANN_CHECK(is_valid());
  return make_unique<RecurrentLayer_Context>(get_state_size(), get_output_size(), batch_size);
}
unique_ptr<Layer::Context> yann::RecurrentLayer::create_context(const RefVectorBatch & output) const
{
  YANN_CHECK(is_valid());
  return make_unique<RecurrentLayer_Context>(get_state_size(), output);
}
unique_ptr<Layer::Context> yann::RecurrentLayer::create_training_context(
    const MatrixSize & batch_size, const std::unique_ptr<Layer::Updater> & updater) const
{
  YANN_CHECK(is_valid());
  YANN_CHECK(updater);
  return make_unique<RecurrentLayer_TrainingContext>(
      get_input_size(), get_state_size(), get_output_size(), batch_size, updater);
}
unique_ptr<Layer::Context> yann::RecurrentLayer::create_training_context(
    const RefVectorBatch & output, const std::unique_ptr<Layer::Updater> & updater) const
{
  YANN_CHECK(is_valid());
  YANN_CHECK(updater);
  return make_unique<RecurrentLayer_TrainingContext>(
      get_input_size(), get_state_size(), output, updater);
}

template<typename InputType>
void yann::RecurrentLayer::feedforward_internal(
    const InputType & input,
    Context * context,
    enum OperationMode mode) const
{
  auto ctx = dynamic_cast<RecurrentLayer_Context *>(context);
  YANN_CHECK(ctx);
  YANN_CHECK(is_valid());
  YANN_CHECK_LE(get_batch_size(input) + ctx->_pos, ctx->get_batch_size());
  YANN_CHECK_LE(get_batch_size(input), get_batch_size(ctx->get_output()));
  YANN_CHECK_EQ(get_batch_item_size(input), get_input_size());
  YANN_SLOW_CHECK(is_same_size(ctx->_zz_h, ctx->_hh));
  YANN_SLOW_CHECK(is_same_size(ctx->_zz_a, ctx->get_output()));

  for(MatrixSize ii = 0; ii < get_batch_size(input); ++ii, ++ctx->_pos) {
    // hh(t) = state_activation_function(ww_hh * hh(t-1) + ww_xh*x(t) + b_h)  // state
    auto in = get_batch(input, ii);
    auto out = get_batch(ctx->get_output(), ctx->_pos);
    auto hh = get_batch(ctx->_hh, ctx->_pos);
    auto zz_h = get_batch(ctx->_zz_h, ctx->_pos);
    auto zz_a = get_batch(ctx->_zz_a, ctx->_pos);

    if(ctx->_pos > 0) {
      auto hh_prev = get_batch(ctx->_hh, ctx->_pos - 1);
      zz_h = MatrixFunctions<InputType>::product(hh_prev, _ww_hh) +
             MatrixFunctions<InputType>::product(in, _ww_xh) +
             _bb_h;
    } else {
      zz_h = MatrixFunctions<InputType>::product(in, _ww_xh) +
            _bb_h;
    }
    _state_activation_function->f(zz_h, hh, Operation_Assign);

    // a(t) = output_activation_function(ww_ha * hh(t) + b_a) // output
    zz_a = MatrixFunctions<InputType>::product(hh, _ww_ha) +
        _bb_a;
    _state_activation_function->f(zz_a, out, mode);
  }
}

void yann::RecurrentLayer::feedforward(
    const RefConstVectorBatch & input,
    Context * context,
    enum OperationMode mode) const
{
  feedforward_internal(input, context, mode);
}

void yann::RecurrentLayer::feedforward(
    const RefConstSparseVectorBatch & input,
    Context * context,
    enum OperationMode mode) const
{
  feedforward_internal(input, context, mode);
}

template<typename InputType>
void yann::RecurrentLayer::backprop_internal(
    const RefConstVectorBatch & gradient_output,
    const InputType & input,
    optional<RefVectorBatch> gradient_input,
    Context * context) const
{
  auto ctx = dynamic_cast<RecurrentLayer_TrainingContext *>(context);
  YANN_CHECK(ctx);
  YANN_CHECK(is_valid());
  YANN_CHECK_GT(get_batch_size(gradient_output), 0);
  YANN_CHECK_LE(get_batch_size(gradient_output), ctx->_pos);
  YANN_CHECK_EQ(get_batch_item_size(gradient_output), get_output_size());
  YANN_CHECK_EQ(get_batch_size(input), get_batch_size(gradient_output));
  YANN_CHECK_EQ(get_batch_item_size(input), get_input_size());
  YANN_CHECK_EQ(get_batch_item_size(input), get_input_size());
  YANN_CHECK(!gradient_input || is_same_size(input, *gradient_input));

  for(MatrixSize ii = get_batch_size(gradient_output) - 1; ii >= 0 && (--ctx->_pos) >= 0; --ii) {
    auto gradient_out = get_batch(gradient_output, ii);
    auto in = get_batch(input, ii);
    auto hh = get_batch(ctx->_hh, ctx->_pos);
    auto zz_h = get_batch(ctx->_zz_h, ctx->_pos);
    auto zz_a = get_batch(ctx->_zz_a, ctx->_pos);

    //
    // output:
    //
    auto & sigma_derivative_zz_a = ctx->_sigma_derivative_zz_a;
    auto & delta_a = ctx->_delta_a;
    auto & delta_ww_ha = ctx->_delta_ww_ha;
    auto & delta_bb_a = ctx->_delta_bb_a;
    auto & gradient_h = ctx->_gradient_h;

    // delta_a(t)    = elem_prod(gradient(C, a(t)), activation_derivative(zz_a(t)))
    YANN_SLOW_CHECK(is_same_size(zz_a, sigma_derivative_zz_a));
    YANN_SLOW_CHECK(is_same_size(zz_a, gradient_out));
    YANN_SLOW_CHECK(is_same_size(zz_a, delta_a));
    _output_activation_function->derivative(zz_a, sigma_derivative_zz_a);
    delta_a.array() = gradient_out.array() * sigma_derivative_zz_a.array();

    // dC/d b_a(t)   = delta_a(t)
    YANN_SLOW_CHECK(is_same_size(delta_bb_a, delta_a));
    delta_bb_a.noalias()  += delta_a;

    // dC/d ww_ha(t) = hh(t) * delta_a(t)
    YANN_SLOW_CHECK(is_same_size(delta_ww_ha, _ww_ha));
    delta_ww_ha.noalias() += MatrixFunctions<InputType>::product(hh.transpose(), delta_a);

    // total_gradient(C, hh(t)) = gradient(C, hh(t)) + transp(ww_ha) * delta_a(t)
    gradient_h.noalias()  += MatrixFunctions<InputType>::product(delta_a, _ww_ha.transpose());

    //
    // state:
    //
    auto & sigma_derivative_zz_h = ctx->_sigma_derivative_zz_h;
    auto & delta_h = ctx->_delta_h;
    auto & delta_ww_xh = ctx->_delta_ww_xh;
    auto & delta_ww_hh = ctx->_delta_ww_hh;
    auto & delta_bb_h = ctx->_delta_bb_h;

    // delta_h(t) = elem_prod(total_gradient(C, hh(t)), activation_derivative(zz_h(t)))
    YANN_CHECK(is_same_size(zz_h, sigma_derivative_zz_h));
    YANN_CHECK(is_same_size(zz_h, gradient_h));
    _state_activation_function->derivative(zz_h, sigma_derivative_zz_h);
    delta_h.array() = gradient_h.array() * sigma_derivative_zz_h.array();

    // dC/d b_h(t) = delta_h(t)
    YANN_CHECK(is_same_size(delta_bb_h, delta_h));
    delta_bb_h.noalias() += delta_h;

    // dC/d ww_xh(t) = x(t) * delta_h(t)
    YANN_SLOW_CHECK(is_same_size(delta_ww_xh, _ww_xh));
    delta_ww_xh += MatrixFunctions<InputType>::product(in.transpose(), delta_h);

    // dC/d ww_hh(t) = hh(t-1) * delta_h(t)
    if(ctx->_pos > 0) {
      auto prev_hh = get_batch(ctx->_hh, ctx->_pos - 1);
      YANN_SLOW_CHECK(is_same_size(delta_ww_hh, _ww_hh));
      delta_ww_hh.noalias() += MatrixFunctions<InputType>::product(prev_hh.transpose(), delta_h);
    }

    // gradient(C, hh(t-1)) = transp(ww_hh) * delta_h(t)
    if(ctx->_pos > 0) {
      gradient_h.noalias() = MatrixFunctions<InputType>::product(delta_h, _ww_hh.transpose());
    }

    // gradient(C, x(t)) = transp(ww_xh) * delta_h(t)
    if(gradient_input) {
      auto gradient_in = get_batch(*gradient_input, ii);
      gradient_in.noalias() = MatrixFunctions<InputType>::product(delta_h, _ww_xh.transpose());
    }
  }
}

void yann::RecurrentLayer::backprop(
    const RefConstVectorBatch & gradient_output,
    const RefConstVectorBatch & input,
    optional<RefVectorBatch> gradient_input,
    Context * context) const
{
  backprop_internal(gradient_output, input, gradient_input, context);
}

void yann::RecurrentLayer::backprop(
    const RefConstVectorBatch & gradient_output,
    const RefConstSparseVectorBatch & input,
    optional<RefVectorBatch> gradient_input,
    Context * context) const
{
  // <RefConstSparseMatrix> is required for MatrixFunctions<>::product
  backprop_internal<RefConstSparseMatrix>(gradient_output, input, gradient_input, context);
}

void yann::RecurrentLayer::init(enum InitMode mode, boost::optional<InitContext> init_context)
{
  switch (mode) {
  case InitMode_Zeros:
    _ww_hh.setZero();
    _ww_xh.setZero();
    _bb_h.setZero();
    _ww_ha.setZero();
    _bb_a.setZero();
    break;
  case InitMode_Random:
    {
      unique_ptr<RandomGenerator> gen01 = RandomGenerator::normal_distribution(0, 1,
          init_context ? optional<Value>(init_context->seed()) : boost::none);
      gen01->generate(_ww_hh);
      gen01->generate(_ww_xh);
      gen01->generate(_bb_h);
      gen01->generate(_ww_ha);
      gen01->generate(_bb_a);
    }
    break;
  }
}

void yann::RecurrentLayer::update(Context * context, const size_t & tests_num)
{
  auto ctx = dynamic_cast<RecurrentLayer_TrainingContext *>(context);
  YANN_CHECK(ctx);
  YANN_SLOW_CHECK(ctx->_ww_hh_updater);
  YANN_SLOW_CHECK(ctx->_ww_xh_updater);
  YANN_SLOW_CHECK(ctx->_bb_h_updater);
  YANN_SLOW_CHECK(ctx->_ww_ha_updater);
  YANN_SLOW_CHECK(ctx->_bb_a_updater);
  YANN_SLOW_CHECK(is_same_size(_ww_hh, ctx->_delta_ww_hh));
  YANN_SLOW_CHECK(is_same_size(_ww_xh, ctx->_delta_ww_xh));
  YANN_SLOW_CHECK(is_same_size(_bb_h, ctx->_delta_bb_h));
  YANN_SLOW_CHECK(is_same_size(_ww_ha, ctx->_delta_ww_ha));
  YANN_SLOW_CHECK(is_same_size(_bb_a, ctx->_delta_bb_a));

  ctx->_ww_hh_updater->update(ctx->_delta_ww_hh, tests_num, _ww_hh);
  ctx->_ww_xh_updater->update(ctx->_delta_ww_xh, tests_num, _ww_xh);
  ctx->_bb_h_updater->update(ctx->_delta_bb_h, tests_num, _bb_h);
  ctx->_ww_ha_updater->update(ctx->_delta_ww_ha, tests_num, _ww_ha);
  ctx->_bb_a_updater->update(ctx->_delta_bb_a, tests_num, _bb_a);
}

// the format is (wh:<_ww_hh>,wx:<_ww_xh>,bh:<_bb_h>,wa:<_ww_ha>,ba:<_bb_a>)
void yann::RecurrentLayer::read(std::istream & is)
{
  Base::read(is);

  read_char(is, '(');
  read_object(is, "wh", _ww_hh);
  read_char(is, ',');
  read_object(is, "wx", _ww_xh);
  read_char(is, ',');
  read_object(is, "bh", _bb_h);
  read_char(is, ',');
  read_object(is, "wa", _ww_ha);
  read_char(is, ',');
  read_object(is, "ba", _bb_a);
  read_char(is, ')');
}

// the format is (wh:<_ww_hh>,wx:<_ww_xh>,bh:<_bb_h>,wa:<_ww_ha>,ba:<_bb_a>)
void yann::RecurrentLayer::write(std::ostream & os) const
{
  Base::write(os);

  os << "(";
  write_object(os, "wh", _ww_hh);
  os << ",";
  write_object(os, "wx", _ww_xh);
  os << ",";
  write_object(os, "bh", _bb_h);
  os << ",";
  write_object(os, "wa", _ww_ha);
  os << ",";
  write_object(os, "ba", _bb_a);
  os << ")";
}


#include "NlMeans.hpp"

#include "thread/ThreadUtils.hpp"
#include "thread/ThreadPool.hpp"

#include "Logging.hpp"

#include <Eigen/Dense>
#include <iostream>
#include <vector>

namespace Tungsten {

Pixmap3f collaborativeRegression(const Pixmap3f &image, const Pixmap3f &guide,
        const std::vector<PixmapF> &features, const Pixmap3f &imageVariance,
        int F, int R, float k)
{
    int w = image.w();
    int h = image.h();
    int d = features.size() + 3;

    // We parallellize by dicing up the image into 32x32 tiles
    const int TileSize = 32;
    int padSize = TileSize + 2*F;

    struct Tile
    {
        Vec2i pos;
        Box2i dstRect;
        Pixmap3f result;
        PixmapF resultWeights;
        Tile(int x, int y) : pos(x, y) {}
    };

    std::vector<Tile> tiles;
    for (int tileY : range(0, h, TileSize))
        for (int tileX : range(0, w, TileSize))
            tiles.emplace_back(tileX, tileY);

    struct PerThreadData { Pixmap3f tmpBufA, tmpBufB; std::vector<PixmapF> weights; };
    std::vector<std::unique_ptr<PerThreadData>> threadData(ThreadUtils::idealThreadCount());

    ThreadUtils::pool->enqueue([&](uint32 i, uint32, uint32 threadId) {
        printProgressBar(i, tiles.size());

        if (!threadData[threadId]) {
            threadData[threadId].reset(new PerThreadData());
            threadData[threadId]->tmpBufA = Pixmap3f(padSize, padSize);
            threadData[threadId]->tmpBufB = Pixmap3f(padSize, padSize);
            for (int i = 0; i < (2*R + 1)*(2*R + 1); ++i)
                threadData[threadId]->weights.emplace_back(TileSize, TileSize);
        }
        auto &data = *threadData[threadId];
        Tile &tile = tiles[i];

        for (auto &w : data.weights)
            w.clear();

        Box2i srcRect(tile.pos, min(tile.pos + TileSize, Vec2i(w, h)));
        tile.dstRect = srcRect;
        tile.dstRect.grow(R);
        tile.dstRect.intersect(Box2i(Vec2i(0), Vec2i(w, h)));
        int dstW = tile.dstRect.diagonal().x(), dstH = tile.dstRect.diagonal().y();

        tile.result = Pixmap3f(dstW, dstH);
        tile.resultWeights = PixmapF(dstW, dstH);

        // Precompute weights for entire tile
        for (int dy = -R, idxW = 0; dy <= R; ++dy)
            for (int dx = -R; dx <= R; ++dx, ++idxW)
                nlMeansWeights(data.weights[idxW], data.tmpBufA, data.tmpBufB, guide, imageVariance,
                        srcRect, F, k, dx, dy, 2.0f);

        for (int y = srcRect.min().y(); y < srcRect.max().y(); ++y) {
            for (int x = srcRect.min().x(); x < srcRect.max().x(); ++x) {
                int x0 = max(x - R, 0), x1 = min(w, x + R + 1);
                int y0 = max(y - R, 0), y1 = min(h, y + R + 1);
                int n = (x1 - x0)*(y1 - y0);

                // Build weight matrix (W), feature matrix (X) and RHS (Y)
                Eigen::VectorXf W(n);
                Eigen::MatrixXf Y(n, 3);
                Eigen::MatrixXf X(n, d);

                for (int iy = y0; iy < y1; ++iy) {
                    for (int ix = x0; ix < x1; ++ix) {
                        int idxP = ix + iy*w;
                        int idx = (ix - x0) + (iy - y0)*(x1 - x0);

                        for (int i = 0; i < 3; ++i)
                            Y(idx, i) = image[idxP][i];

                        X(idx, 0) = 1.0f;
                        X(idx, 1) = ix - x;
                        X(idx, 2) = iy - y;
                        for (size_t i = 0; i < features.size(); ++i)
                            X(idx, i + 3) = features[i][idxP] - features[i][x + y*w];

                        int idxW = (ix - x + R) + (iy - y + R)*(2*R + 1);
                        W[idx] = data.weights[idxW][Vec2i(x, y) - tile.pos];
                    }
                }

                // Solve least squares system
                Eigen::VectorXf wSqrt = W.cwiseSqrt();
                Eigen::MatrixXf denoised = X*(wSqrt.asDiagonal()*X).colPivHouseholderQr().solve(wSqrt.asDiagonal()*Y);

                // Accumulate denoised patch into image
                for (int iy = y0; iy < y1; ++iy) {
                    for (int ix = x0; ix < x1; ++ix) {
                        Vec2i p = Vec2i(ix, iy) - tile.dstRect.min();
                        int idx = (ix - x0) + (iy - y0)*(x1 - x0);
                        tile.result       [p] += W[idx]*Vec3f(denoised(idx, 0), denoised(idx, 1), denoised(idx, 2));
                        tile.resultWeights[p] += W[idx];
                    }
                }
            }
        }

    }, tiles.size())->wait();

    // Gather results from all threads and divide by weights
    Pixmap3f result(w, h);
    PixmapF resultWeights(w, h);
    for (const auto &tile : tiles) {
        for (int y  : tile.dstRect.range(1)) {
            for (int x  : tile.dstRect.range(0)) {
                Vec2i p(x, y);
                result       [p] += tile.result       [p - tile.dstRect.min()];
                resultWeights[p] += tile.resultWeights[p - tile.dstRect.min()];
            }
        }
    }
    for (int j = 0; j < w*h; ++j)
        result[j] /= resultWeights[j];

    printProgressBar(tiles.size(), tiles.size());

    return std::move(result);
}

}


/*
 * ElGammal.cpp
 *
 *  Created on: 03.10.2010
 *      Author: stephaniebayer
 */

#include "ElGammal.h"
#include "G_q.h"

#include <NTL/ZZ.h>
NTL_CLIENT

#include "Mod_p.h"
#include <stdio.h>
#include <time.h>
#include <vector>
#include <fstream>

ElGammal::ElGammal() {
	// TODO Auto-generated constructor stub

}

//Creates ElGammal with secret key s, public key p and group H
ElGammal::ElGammal(long s, Mod_p p, G_q H){
	G=H;
	sk = to_ZZ(s);
	pk = p;

}

//Creates ElGammal with secret key s, public key p and group H
ElGammal::ElGammal(ZZ s, Mod_p p, G_q H){
	G=H;
	sk = s;
	pk = p;

}

//Creates ElGammal with secret key s and group H, the public key is pk = gen^s , gen generator of H
ElGammal::ElGammal(long s, G_q H){
	Mod_p temp;
	G = H;
	sk = to_ZZ(s);
	temp = Mod_p(G.get_gen().get_val(), G.get_mod());
	pk = temp.expo(s);
}

//Creates ElGammal with secret key s and group H, the public key is pk = gen^s , gen generator of H
ElGammal::ElGammal(ZZ s, G_q H){
	Mod_p temp;
	G = H;
	sk = s;
	temp = Mod_p(G.get_gen().get_val(), G.get_mod());
	pk = temp.expo(s);
}

//Set the group to G_q with order o, G_q subset of G_mod_p and generator gen, secret key is s and pk = gen^s
ElGammal::ElGammal(Mod_p gen, long o, long  mod, long s){

	G = G_q(gen, o, mod);
	sk = to_ZZ(s);
	Mod_p temp;
	temp = Mod_p(gen.get_val(), mod);
	pk = temp.expo(s);
}

//Set the group to G_q with order o, G_q subset of G_mod_p and generator gen, secret key is s and pk = gen^s
ElGammal::ElGammal(Mod_p gen, long o, ZZ  mod, long s){

	G = G_q(gen, o, mod);
	sk = to_ZZ(s);
	Mod_p temp;
	temp = Mod_p(gen.get_val(), mod);
	pk = temp.expo(s);
}

//Set the group to G_q with order o, G_q subset of G_mod_p and generator gen, secret key is s and pk = gen^s
ElGammal::ElGammal(Mod_p gen, long o, ZZ  mod, ZZ s){

	G = G_q(gen, o, mod);
	sk = s;
	Mod_p temp;
	temp = Mod_p(gen.get_val(), mod);
	pk = temp.expo(s);
}

//Set the group to G_q with order o, G_q subset of G_mod_p and generator gen, secret key is s and pk = gen^s
ElGammal::ElGammal(Mod_p gen, ZZ o, ZZ  mod, long s){

	G = G_q(gen, o, mod);
	sk = to_ZZ(s);
	Mod_p temp;
	temp = Mod_p(gen.get_val(), mod);
	pk = temp.expo(s);
}

//Set the group to G_q with order o, G_q subset of G_mod_p and generator gen, secret key is s and pk = gen^s
ElGammal::ElGammal(Mod_p gen, ZZ o, ZZ  mod, ZZ s){

	G = G_q(gen, o, mod);
	sk = s;
	Mod_p temp;
	temp = Mod_p(gen.get_val(), mod);
	pk = temp.expo(s);
}

//Set the group to G_q with order o, G_q subset of G_mod_p and generator gen, secret key is s and public key p
ElGammal::ElGammal(Mod_p gen, long o, long  mod, long s, Mod_p p){

	G = G_q(gen, o, mod);
	sk = to_ZZ(s);
	pk = p;
}

//Set the group to G_q with order o, G_q subset of G_mod_p and generator gen, secret key is s and public key p
ElGammal::ElGammal(Mod_p gen, long o, ZZ  mod, long s, Mod_p p){

	G = G_q(gen, o, mod);
	sk = to_ZZ(s);
	pk = p;
}

//Set the group to G_q with order o, G_q subset of G_mod_p and generator gen, secret key is s and public key p
ElGammal::ElGammal(Mod_p gen, long o, ZZ  mod, ZZ s, Mod_p p){

	G = G_q(gen, o, mod);
	sk = s;
	pk = p;
}

//Set the group to G_q with order o, G_q subset of G_mod_p and generator gen, secret key is s and public key p
ElGammal::ElGammal(Mod_p gen, ZZ o, ZZ  mod, long s, Mod_p p){

	G = G_q(gen, o, mod);
	sk = to_ZZ(s);
	pk = p;
}

//Set the group to G_q with order o, G_q subset of G_mod_p and generator gen, secret key is s and public key p
ElGammal::ElGammal(Mod_p gen, ZZ o, ZZ  mod, ZZ s, Mod_p p){

	G = G_q(gen, o, mod);
	sk = s;
	pk = p;
}

//Set the group to G_q with order o, G_q subset of G_mod_p and generator gen, secret key is s and public key p
ElGammal::ElGammal( long o, long  mod, long s, Mod_p p){

	G = G_q( o, mod);
	sk = to_ZZ(s);
	pk = p;
}

//Set the group to G_q with order o and modular value mod, secret key is s and public key p
ElGammal::ElGammal( long o, ZZ  mod, long s, Mod_p p){

	G = G_q( o, mod);
	sk = to_ZZ(s);
	pk = p;
}

//Set the group to G_q with order o and modular value mod, secret key is s and public key p
ElGammal::ElGammal( long o, ZZ  mod, ZZ s, Mod_p p){

	G = G_q( o, mod);
	sk = s;
	pk = p;
}

//Set the group to G_q with order o and modular value mod, secret key is s and public key p
ElGammal::ElGammal( ZZ o, ZZ  mod, long s, Mod_p p){

	G = G_q( o, mod);
	sk = to_ZZ(s);
	pk = p;
}

//Set the group to G_q with order o and modular value mod, secret key is s and public key p
ElGammal::ElGammal( ZZ o, ZZ  mod, ZZ s, Mod_p p){

	G = G_q( o, mod);
	sk = to_ZZ(s);
	pk = p;
}

//Set the group to G_q with order o and modular value mod, secret key is s and public key pk = gen^s
ElGammal::ElGammal( long o, long  mod, long s){

	G = G_q( o, mod);
	sk = to_ZZ(s);
	Mod_p temp;
	temp = Mod_p(G.get_gen().get_val(), mod);
	pk = temp.expo(s);
}

//Set the group to G_q with order o and modular value mod, secret key is s and public key pk = gen^s
ElGammal::ElGammal( long o, ZZ  mod, long s){

	G = G_q( o, mod);
	sk = to_ZZ(s);
	Mod_p temp;
	temp = Mod_p(G.get_gen().get_val(), mod);
	pk = temp.expo(s);
}

//Set the group to G_q with order o and modular value mod, secret key is s and public key pk = gen^s
ElGammal::ElGammal( long o, ZZ  mod, ZZ s){

	G = G_q( o, mod);
	sk = s;
	Mod_p temp;
	temp = Mod_p(G.get_gen().get_val(), mod);
	pk = temp.expo(s);
}

//Set the group to G_q with order o and modular value mod, secret key is s and public key pk = gen^s
ElGammal::ElGammal( ZZ o, ZZ  mod, long s){

	G = G_q( o, mod);
	sk = to_ZZ(s);
	Mod_p temp;
	temp = Mod_p(G.get_gen().get_val(), mod);
	pk = temp.expo(s);
}

//Set the group to G_q with order o and modular value mod, secret key is s and public key pk = gen^s
ElGammal::ElGammal( ZZ o, ZZ  mod, ZZ s){

	G = G_q( o, mod);
	sk = s;
	Mod_p temp;
	temp = Mod_p(G.get_gen().get_val(), mod);
	pk = temp.expo(s);
}

ElGammal::~ElGammal() {
	// TODO Auto-generated destructor stub
}

//Access to the parameters
G_q ElGammal::get_group()const{
	return G;
}

Mod_p ElGammal::get_pk() const{

	return pk;
}

ZZ ElGammal::get_sk()const{

	return sk;
}

//functions to change parameters
void ElGammal::set_group(G_q H){

	G = H;
}

void ElGammal::set_sk(long s){

	sk = to_ZZ(s);
	pk = G.get_gen().expo( s);
}

void ElGammal::set_sk(ZZ s){

	sk = s;
	pk = G.get_gen().expo(s);
	string name = "example.txt";
	ofstream ost;
	ost.open(name.c_str(),ios::app);
	ost<<"private key and public key "<<sk<<" "<<pk<<endl;
}

//functions to encrypt value/element
Cipher_elg ElGammal::encrypt(Mod_p el){
	Cipher_elg c;
	Mod_p temp_1, temp_2;
	ZZ ran;
	SetSeed(to_ZZ(time(0)));
	ran = RandomBnd(G.get_ord());
	temp_1 = G.get_gen().expo(ran);
	temp_2 = pk.expo(ran)*el;
	c = Cipher_elg(temp_1,temp_2);
	return c;

}

Cipher_elg ElGammal::encrypt(ZZ m){
	Cipher_elg c;
	Mod_p temp_1, temp_2;
	ZZ ran;
	SetSeed(to_ZZ(time(0)));
	ran = RandomBnd(G.get_ord());
	cout<< ran << endl;
	temp_1 = G.get_gen().expo(ran);
	temp_2 = pk.expo(ran)*Mod_p(m,G.get_mod());
	c = Cipher_elg(temp_1,temp_2);
	return c;
}

Cipher_elg ElGammal::encrypt(long m){
	Cipher_elg c;
	Mod_p temp_1, temp_2;
	ZZ ran;
	SetSeed(to_ZZ(time(0)));
	ran = RandomBnd(G.get_ord());
	temp_1 = G.get_gen().expo(ran);
	temp_2 = pk.expo(ran)*Mod_p(m,G.get_mod());
	c = Cipher_elg(temp_1,temp_2);
	return c;
}

Cipher_elg ElGammal::encrypt(Mod_p el, long ran){
	Cipher_elg c;
	Mod_p temp_1, temp_2;
	temp_1 = G.get_gen().expo(ran);
	temp_2 = pk.expo(ran)*el;
	c = Cipher_elg(temp_1,temp_2);
	return c;

}

Cipher_elg ElGammal::encrypt(Mod_p el, ZZ ran){
	Cipher_elg c;
	Mod_p temp_1, temp_2;
	temp_1 = G.get_gen().expo(ran);
	temp_2 = pk.expo(ran)*el;
	c = Cipher_elg(temp_1,temp_2);
	return c;

}

Cipher_elg ElGammal::encrypt(ZZ m, long ran){
	Cipher_elg c;
	Mod_p temp_1, temp_2;
	temp_1 = G.get_gen().expo(ran);
	temp_2 = pk.expo(ran)*Mod_p(m,G.get_mod());
	c = Cipher_elg(temp_1,temp_2);
	return c;
}

Cipher_elg ElGammal::encrypt(ZZ m, ZZ ran){
	Cipher_elg c;
	Mod_p temp_1, temp_2;
	temp_1 = G.get_gen().expo(ran);
	temp_2 = pk.expo(ran)*Mod_p(m,G.get_mod());
	c = Cipher_elg(temp_1,temp_2);
	return c;
}

Cipher_elg ElGammal::encrypt(long m, long ran){
	Cipher_elg c;
	Mod_p temp_1, temp_2;
	temp_1 = G.get_gen().expo(ran);
	temp_2 = pk.expo(ran)*Mod_p(m,G.get_mod());
	c = Cipher_elg(temp_1,temp_2);
	return c;
}

Cipher_elg ElGammal::encrypt(long m, ZZ ran){
	Cipher_elg c;
	Mod_p temp_1, temp_2;
	temp_1 = G.get_gen().expo(ran);
	temp_2 = pk.expo(ran)*Mod_p(m,G.get_mod());
	c = Cipher_elg(temp_1,temp_2);
	return c;
}

//Decrypts the ciphertext c
Mod_p ElGammal::decrypt(Cipher_elg c){
	ZZ temp;
	ZZ mod = G.get_mod();
	temp = InvMod(c.get_u(),mod);
	temp = PowerMod(temp,sk, mod);
	temp = MulMod(temp,c.get_v(),mod);
	return temp;
}

//Assigment operator
void ElGammal::operator=(const ElGammal& el){

	G = el.get_group();
	sk = el.get_sk();
	pk = el.get_pk();
}