#include <Eigen/Sparse>
#include <Eigen/Eigenvalues>
#include <Eigen/StdVector>
#include <cstdlib>
#include <iostream>
#include <iomanip>
#include <mpi.h>
#include <mkl.h>
#include <mkl_spblas.h>
#include <fstream>
#include <iterator>
#include <vector>
#include <algorithm>
#include <complex>
#include <sstream>
#include <math.h>
#include <omp.h>
#include <numeric>
#include <chrono>
#include <ctime>
#include <limits>
#include <lyra/lyra.hpp>
#include <sys/stat.h>
#include <sys/types.h>
#include <cusparse.h>
#include <cusolverDn.h>
#include <cuComplex.h>
#include "cuda.h"
#include "cuda_runtime_api.h"
#include "CUDAaccelerated.h"
#include "bsplines.h"
#include "InputParameter.h"

#pragma GCC diagnostic ignored "-Wdeprecated-declarations"

#define watch(x) cout << (#x) << " is " << (x) << endl
#define pwatch(x) cout <<"rank " << rank <<" "<< (#x) << " is " << (x) << endl
#define bp(x) cout << "breakpoint" << (#x) << endl

using namespace Eigen;
using namespace std::complex_literals;
using namespace std::chrono;
using std::cout;
using std::endl;

int dim;
int rank,p,p_omp;
int64_t nnz_rank;
int states_chan_rank;
int states_chan_rank_begin;
int states_chan_rank_end;

//This is a COO matrix that holds a tupel of values for every matrix cell
class SpMatrix{
    //idea write a matrix class that has the functionality of init memory, write element,
    //multiply
    public:
        std::vector<int> nnz;
        int dim;
        int nmbrLEGO=3;

        //The COO matrix
        std::vector<std::vector<int> > Row;
        std::vector<std::vector<int> > Col;
        std::vector<std::vector<double> > Value;

        //The CSR matrix
        std::vector<std::vector<int> > RowCSR;
        std::vector<std::vector<int> > ColCSR;
        std::vector<std::vector<double> > ValueCSR;

        //GPU accelerated matrix vector
        //pointers to the device memory
        std::vector<int*> LEGO_d_cooCols;
        std::vector<int*> LEGO_d_rowPtr;
        std::vector<double* > LEGO_d_cooVals_sorted;
        std::vector<int*> LEGO_d_cscRows;
        std::vector<int*>  LEGO_d_cscColPtr;
        std::vector<double*> LEGO_d_cscVals;

    void initMatrix(int nnz_input){
        Row.resize(nmbrLEGO);
        Col.resize(nmbrLEGO);
        Value.resize(nmbrLEGO);

        RowCSR.resize(nmbrLEGO);
        ColCSR.resize(nmbrLEGO);
        ValueCSR.resize(nmbrLEGO);

        LEGO_d_cooVals_sorted.resize(nmbrLEGO);
        LEGO_d_cooCols.resize(nmbrLEGO);
        LEGO_d_rowPtr.resize(nmbrLEGO);
        LEGO_d_cscVals.resize(nmbrLEGO);
        LEGO_d_cscColPtr.resize(nmbrLEGO);
        LEGO_d_cscRows.resize(nmbrLEGO);

        for(int i = 0; i < nmbrLEGO; i++){
            Row[i].resize(nnz_input);
            Col[i].resize(nnz_input);
            Value[i].resize(nnz_input);
        }
        nnz.resize(nmbrLEGO);

    }

    void writeElement(int ind, int row, int col, int LEGOind,double value){
        //writes an element to the matrix
        Row[LEGOind][ind] = row;
        Col[LEGOind][ind] = col;
        Value[LEGOind][ind] = value;
    }

    void CSRMKL(){
        //allocate the memeory for the CSR matrices
        watch(nmbrLEGO);
        for(int i = 0; i < nmbrLEGO; i++){
            watch(i);
            watch(nnz[i]);
            RowCSR[i].resize(dim+2);
            ColCSR[i].resize(nnz[i]);
            ValueCSR[i].resize(nnz[i]);
        }
        bp("Allocate memory");

        //build the CSR matrix
        for (int i = 0; i< nmbrLEGO;i++){
            int job[5] = {2,0,0,nnz[i],0};
            int info = 0;
            mkl_dcsrcoo(job,
                    &dim,
                    ValueCSR[i].data(),
                    ColCSR[i].data(),
                    RowCSR[i].data(),
                    &nnz[i],
                    Value[i].data(),
                    Row[i].data(),
                    Col[i].data(),
                    &info
                );
        }

    }

    void CSRonDevice(){
        //converts the matrices stored in COO on the host to CSR on the device by
        //calling the accoarding CUDA function
        //iterate over the LEGO bricks
        if(rank==0){
             int i=0;
             makeCSR(nnz[i],dim,Row[i].data(),Col[i].data(),Value[i].data(),LEGO_d_cooCols[i],
                    LEGO_d_rowPtr[i], LEGO_d_cooVals_sorted[i],
                    LEGO_d_cscRows[i],
                    LEGO_d_cscColPtr[i],
                    LEGO_d_cscVals[i]);
        }
        for(int i=1; i < nmbrLEGO; i++){
            makeCSR(nnz[i],dim,Row[i].data(),Col[i].data(),Value[i].data(),LEGO_d_cooCols[i],
                    LEGO_d_rowPtr[i], LEGO_d_cooVals_sorted[i],
                    LEGO_d_cscRows[i],
                    LEGO_d_cscColPtr[i],
                    LEGO_d_cscVals[i]);

       }
    }

    void matrixVectorUpper(int LEGO,VectorXcd &x,VectorXcd &b,VectorXcd &ft){
       //This is a COO matrix vector product M * x = b
#pragma omp parallel
        {
        VectorXcd b_private= VectorXcd::Zero(dim);
#pragma omp for
        for(int  i = 0; i < nnz[LEGO]; i++)
        {
                b_private(Row[LEGO][i]) += Value[LEGO][i] * x(Col[LEGO][i]) * ft(LEGO);
            if (Row[LEGO][i] != Col[LEGO][i]){
                b_private(Col[LEGO][i]) += Value[LEGO][i] * x(Row[LEGO][i]) * std::conj(ft[LEGO]);
            }
        }
#pragma omp critical
        for(int i=0; i < dim; i++){
            b(i) += b_private(i);
        }
        }
    }

    //wrapper for MKL CSR SpMV
    void matrixVectorMKL(VectorXcd &x,VectorXcd &b,VectorXcd &ft){
        //as the matrix is real we seperate the vectors in real and iamg parts
        //this allows a real double multiplications
        VectorXd dydx_rank_real = b.real();
        VectorXd dydx_rank_imag = b.imag();

       char transa = 'N';
        char matdescra[] = {
            'G', // type of matrix
            ' ', // triangular indicator (ignored in multiplication)
            ' ', // diagonal indicator (ignored in multiplication)
            'C'  // type of indexing
        };

        double beta = 1.0;
        for(int j = 0; j < nmbrLEGO; j++){
            VectorXcd y_local = x * ft(j);
            VectorXd y_local_real = y_local.real();
            VectorXd y_local_imag = y_local.imag();

            //SpMVCuda
            double alpha=1.0;
            //non transpose real x, real y, alpha = 1
            mkl_dcsrmv(&transa,
                    &dim,
                    &dim,
                    &alpha,
                    matdescra,
                    ValueCSR[j].data(),
                    ColCSR[j].data(),
                    RowCSR[j].data(),
                    RowCSR[j].data()+1,
                    y_local_real.data(),
                    &beta,
                    dydx_rank_real.data());

            //non transpose imag x, imag y, alpha = 1
             mkl_dcsrmv(&transa,
                    &dim,
                    &dim,
                    &alpha,
                    matdescra,
                    ValueCSR[j].data(),
                    ColCSR[j].data(),
                    RowCSR[j].data(),
                    RowCSR[j].data()+1,
                    y_local_imag.data(),
                    &beta,
                    dydx_rank_imag.data());

            if (j != 0){
                alpha=-1.0;
                char transa = 'T';
                //transpose real x, real y , alpha = -1
                mkl_dcsrmv(&transa,
                    &dim,
                    &dim,
                    &alpha,
                    matdescra,
                    ValueCSR[j].data(),
                    ColCSR[j].data(),
                    RowCSR[j].data(),
                    RowCSR[j].data()+1,
                    y_local_real.data(),
                    &beta,
                    dydx_rank_real.data());

                //transpose imag x, imag y , alpha = -1
                mkl_dcsrmv(&transa,
                    &dim,
                    &dim,
                    &alpha,
                    matdescra,
                    ValueCSR[j].data(),
                    ColCSR[j].data(),
                    RowCSR[j].data(),
                    RowCSR[j].data()+1,
                    y_local_imag.data(),
                    &beta,
                    dydx_rank_imag.data());

            }
        }
        //join into b
        b.real()=dydx_rank_real;
        b.imag()=dydx_rank_imag;

    }
};

void Propergator(InputParameter &Inp,SpMatrix &H_rank,std::string PATH);

void BuildMatrix(InputParameter &Inp,SpMatrix &H_rank);

void getEigenStates(InputParameter Inp, int l,int states_chan,double dX,double &tmp_d,ArrayXd &weights,
        ArrayXd &x,ArrayXXd &B, ArrayXXd &d_B,ArrayXXd &diff_B,ArrayXXd &psi,ArrayXXd &d_psi
        ,ArrayXXd &diff_psi , MatrixXd &H_base, MatrixXd &S, ArrayXd &Energy);

void helperEigenStates(InputParameter Inp,int states_chan,double dX,double &tmp_d,
        ArrayXd &weights, ArrayXd &x,ArrayXXd &B,ArrayXXd &diff_B,MatrixXd &H_base,
        MatrixXd &S);

void polarization(SpMatrix &M,int states_chan, int &ind, int index,int l,int lm_index_L,int
        lm_index_R, double LM_faktor,double dX,ArrayXd &weights, ArrayXd &x, ArrayXXd &psi_1,
        ArrayXXd  &psi_2,ArrayXXd &d_psi_2);
void lanczos(InputParameter Inp,SpMatrix &A, int m,VectorXcd &y , MatrixXcd &Q, MatrixXcd &h,
        VectorXcd &FT);

auto stoportho = high_resolution_clock::now();
auto startortho = high_resolution_clock::now();
auto durationortho = duration_cast<milliseconds>(startortho-startortho);

int main(int argc, const char** argv){
    //Read in the sparse matrix representing the Hamiltonian
    //The matrix is in COO format(also called matrix market format)
    //The data is stored in 4 different files each of them
    //obtains a vector. 2 files hold the indices for our sparse matrix and
    //2 hold the values that belong to the to indices. One is the value
    //belonging to the Hamiltion the other one is an integer telling us which
    //ft to use
    MPI_Init(NULL,NULL);

    MPI_Comm_size(MPI_COMM_WORLD, &p);
    MPI_Comm_rank(MPI_COMM_WORLD, &rank);

    cout << "rank: " << rank << " of " << p << " is here" << endl;
    InputParameter Inp;

    //map the ranks to the devices(grpahic cards)
    //local rank within a node for multi node version
    //this has to be chagned
    int local_rank = rank;

    p_omp = omp_get_num_threads();
    watch(p_omp);

    //int id = omp_get_thread_num();
    std::setprecision (15);

    //Inp.ReadInput(argc,argv);
    auto cli
        =   lyra::opt(Inp.nnz, "nnz")
            ["-nnz"]["--nnz"]
            ("Number of max nnz")
        | lyra::opt(Inp.om, "om")
            ["-om"]["--om"]
            ("Photon Energy")
        | lyra::opt(Inp.E0, "E0")
            ["-E0"]["--E0"]
            ("Electric field strength")
        |lyra::opt(Inp.T_cycle, "T_cycle")
            ["-T"]["--T_cycle"]
            ("length of pulse in cycles")
        |lyra::opt(Inp.timesteps, "timesteps")
            ["-t"]["--timesteps"]
            ("Number of timesteps")
        |lyra::opt(Inp.b,"b")
            ["-b"]["--b"]
            ("Boxsize")
        |lyra::opt(Inp.N_b,"N_b")
            ["-N_b"]["--N_b"]
            ("Number of gridpoints")
        |lyra::opt(Inp.Lmax,"Lmax")
            ["-Lmax"]["--Lmax"]
            ("Maximum angular momentum")
        |lyra::opt(Inp.Mmax,"Mmax")
            ["-Mmax"]["--Mmax"]
            ("Maximum magnetic quantum number")
        |lyra::opt(Inp.Emax,"Emax")
            ["-Emax"]["--Emax"]
            ("Maximum Energy in simulation")
        |lyra::opt(Inp.I0,"I0")
            ["-I0"]["--I0"]
            ("Laser Intensity in 1e12")
        |lyra::opt(Inp.init_state,"Initial state")
            ["-init"]["--init"]
            ("Initial state of system")
        |lyra::opt(Inp.Accelerator, "Accelerator")
            ["-acc"]["--acc"]
            ("Use CPU or GPU as accelerator")
        |lyra::opt(Inp.zComp, "z-Comp")
            ["-zComp"]["--zComp"]
            ("z comp")
        |lyra::opt(Inp.xComp, "x-comp")
            ["-xComp"]["--xComp"]
            ("x Comp")
        |lyra::opt(Inp.xft, "x-ft")
            ["-xft"]["--xft"]
            ("x ft")
        |lyra::opt(Inp.zft, "z-ft")
            ["-zft"]["--zft"]
            ("z ft")
        |lyra::opt(Inp.l_quantum,"l Quantum number")
            ["-lq"]["--l_q"]
            ("l Quantum number of initial state")
        |lyra::opt(Inp.m_quantum,"n Quantum number")
            ["-mq"]["--m_q"]
            ("m Quantum number of initial state")
        |lyra::opt(Inp.n_quantum,"n Quantum number")
            ["-nq"]["--n_q"]
            ("n Quantum number of initial state");

        auto result = cli.parse({ argc, argv });
        if ( !result )
        {
                std::cerr << "Error in command line: " << result.errorMessage() << std::endl;
                    exit(1);
        }

    //number of accelerators in the system
    int num_devices = 0;
    if (Inp.Accelerator == "GPU" ||  Inp.Accelerator == "GPUFull"){
        setupGPUs(local_rank,num_devices);
        watch(num_devices);
    }

    SpMatrix H_rank;
    H_rank.nmbrLEGO = 1 + Inp.zComp +  Inp.xComp;
    BuildMatrix(Inp,H_rank);
    std::string PATH;
    Propergator(Inp,H_rank,PATH);

    if (Inp.Accelerator == "GPU" ||  Inp.Accelerator == "GPUFull"){
        CUDAFinalize(H_rank.nmbrLEGO,
                H_rank.LEGO_d_rowPtr,H_rank.LEGO_d_cooCols,H_rank.LEGO_d_cooVals_sorted,
                H_rank.LEGO_d_cscRows,H_rank.LEGO_d_cscColPtr,H_rank.LEGO_d_cscVals);
    }
    watch(rank);
    MPI_Finalize();
    cout << "MPI has been finalized" << endl;
    return 0;

}

void BuildMatrix(InputParameter &Inp,SpMatrix &H_rank){
    //recalculate variables that depend on command line input
        //spacing between breakpoints
    VectorXd x_break(Inp.N_b);
    double dX = (Inp.b - Inp.a)/ (Inp.N_b-1);
    x_break = ArrayXd::LinSpaced(Inp.N_b,Inp.a,Inp.b);

    int states_chan = (int) round(Inp.b/M_PI*sqrt(2*Inp.Emax));

    if (rank==0){
        //print the Input parameters
        watch(Inp.nnz);
        watch(Inp.a);
        watch(Inp.b);
        watch(Inp.Lmax);
        watch(Inp.Mmax);
        watch(Inp.Emax);
        watch(Inp.N_b);
        watch(states_chan);
        watch(Inp.zComp);
        watch(Inp.xComp);
        watch(Inp.zft);
        watch(Inp.xft);

        watch(Inp.Accelerator);
        watch(Inp.l_quantum);
        watch(Inp.m_quantum);
        watch(Inp.n_quantum);
        watch(Inp.init_state);
        watch(Inp.om);
        watch(Inp.I0);
        watch(Inp.E0);
        watch(Inp.T_puls);
        watch(Inp.T_cycle);
        watch(Inp.Tint);
        watch(Inp.timesteps);

    }

    auto start = system_clock::now();
      nnz_rank = Inp.nnz/p;
    states_chan_rank=states_chan/p;
    //distribute the remaining channels if states_cahn%p != 0
    if(rank < states_chan%p){
        states_chan_rank++;
    }
    //offset describes how many ranks have an extra channel
    int offset = std::min(states_chan%p,rank);
    states_chan_rank_begin=states_chan/p*rank+offset;
    states_chan_rank_end=states_chan/p*rank+states_chan_rank+offset;

    pwatch(states_chan_rank);
    pwatch(offset);
    pwatch(states_chan_rank_begin);
    pwatch(states_chan_rank_end);

    pwatch("matrix obj is created");
    bp("before memory is allocated");
    H_rank.initMatrix(Inp.nnz/p);
    H_rank.nnz[0] = nnz_rank;
    bp("after memory is allocated");

    //this is the ind on every rank his is a vector that holds the ind for all
    //different matrices linked to the different LEGO bricks
    std::vector< int>ind(H_rank.nmbrLEGO,0);

    //the matrix will be generated on different MPI ranks
    //generating the B-spliens does not take much time therefore all ranks generate
    //their own set. They also for now all hold their own set of psi

    //Gaus Legendre integration set up
    ArrayXd x((Inp.N_b-1)*Inp.n);
    ArrayXd weights((Inp.N_b-1)*Inp.n);
    bp("before gaus set up");
    gausLegendreSetup(Inp.n,Inp.N_b,x_break,x,weights);
    //gauss legendre returns same weights and x as python script
    //allocating a vector of Eigenvector matrices

    bp("after gaus set up");

    ArrayXXd B = ArrayXXd::Zero(x.size(),Inp.N_b-1+2*Inp.n);
    ArrayXXd d_B = ArrayXXd::Zero(x.size(),Inp.N_b-1+2*Inp.n);
    ArrayXXd diff_B = ArrayXXd::Zero(x.size(),Inp.N_b-1+2*Inp.n);

    int dim_H = Inp.N_b+Inp.n-3;
    generateBsplines(Inp,dX,x,B,d_B,diff_B);
    bp("bsplines generated");
    //Memory declaration
    ArrayXXd psi_1 = ArrayXXd::Zero(x.size(),states_chan);
    ArrayXXd d_psi_1 = ArrayXXd::Zero(x.size(),states_chan);
    ArrayXXd diff_psi_1 = ArrayXXd::Zero(x.size(),states_chan);
    ArrayXd E1 = ArrayXd::Zero(dim_H);

    ArrayXXd psi_2 = ArrayXXd::Zero(x.size(),states_chan);
    ArrayXXd d_psi_2 = ArrayXXd::Zero(x.size(),states_chan);
    ArrayXXd diff_psi_2 = ArrayXXd::Zero(x.size(),states_chan);
    ArrayXd E2 = ArrayXd::Zero(dim_H);

    ArrayXXd psi_3 = ArrayXXd::Zero(x.size(),states_chan);
    ArrayXXd d_psi_3 = ArrayXXd::Zero(x.size(),states_chan);
    ArrayXXd diff_psi_3 = ArrayXXd::Zero(x.size(),states_chan);
    ArrayXd E3 = ArrayXd::Zero(dim_H);

    MatrixXd H_base = MatrixXd::Zero(dim_H,dim_H);
    MatrixXd S = MatrixXd::Zero(dim_H,dim_H);
    double tmp_d;
    helperEigenStates(Inp,states_chan,dX,tmp_d,weights, x, B,diff_B,H_base,S);
     getEigenStates(Inp,Inp.Lmin,states_chan,dX,tmp_d,weights,x,B,d_B,diff_B,psi_1,d_psi_1,diff_psi_1,H_base,S,E1);
     bp("First set of psi");

    getEigenStates(Inp,Inp.Lmin+1,states_chan,dX,tmp_d,weights,x,B,d_B,diff_B,psi_2,d_psi_2,diff_psi_2,H_base,S,E2);

     bp("second set of psi");

     ArrayXd Input_Energies = ArrayXd::Zero(5000000);
    int ind2 =0;
   int lm_index_L =-1;
    for (int l=Inp.Lmin; l < Inp.Lmax+1; l++){
        pwatch(l);
       getEigenStates(Inp,l+2,states_chan,dX,tmp_d,weights,x,B,d_B,diff_B,psi_3,d_psi_3,diff_psi_3,H_base,S,E3);

        for(int m=-std::min(Inp.Mmax,l); m < std::min(Inp.Mmax,l)+1; m++){
            lm_index_L++;
            int lm_index_R=-1;
            for(int l2 = Inp.Lmin; l2 < Inp.Lmax+1; l2++){
                for(int m2 = -std::min(Inp.Mmax,l2);m2 < std::min(Inp.Mmax,l2)+1;m2++){
                    lm_index_R++;
                    if (l2==l+1 && m2 ==m){
                        //z-polarization
                        //< l m | x | l+1 m>:
                        if(Inp.zComp == 1){
                            double LM_faktor=sqrt(((l+1)*(l+1)-m*m)/(double) (4*((l+1)*(l+1))-1));
                            polarization(H_rank,states_chan, ind[Inp.zft],Inp.zft,l,lm_index_L,lm_index_R,
                                LM_faktor,dX,weights,x,psi_1,psi_2,d_psi_2);
                        }
                    }
                    if(l2== l+1 && m2 == m+1){
                        //x-polarization
                        //%< l m | p_x | l+1 m+1>:
                        if(Inp.xComp == 1){
                            double LM_faktor=-sqrt((l+m+1)*(l+m+2)/(double) (4*(2*l+1)*(2*l+3)));
                            polarization(H_rank,states_chan, ind[Inp.xft],Inp.xft,l,lm_index_L,lm_index_R,
                                LM_faktor,dX,weights,x,psi_1,psi_2,d_psi_2);
                        }
                    }
                    if(l2== l+1 && m2 == m-1){
                        //x-polarization
                        //< l m | p_x | l+1 m-1>:
                        if(Inp.xComp == 1){
                            double LM_faktor=sqrt((l-m+1)*(l-m+2)/(double) (4*(2*l+1)*(2*l+3)));
                            polarization(H_rank,states_chan, ind[Inp.xft],Inp.xft,l,lm_index_L,lm_index_R,
                                    LM_faktor,dX,weights,x,psi_1,psi_2,d_psi_2);
                        }
                    }
                }
            }

            //Lego brick 4 will unperturbated states will be only calculated on rank 0 as they are
            //not many. We also need the Input_Energies later as a txt file
            if(rank==0){
               //Lego brick 4
                for(int j= 0; j < states_chan;j++){
                    H_rank.writeElement(ind[0],lm_index_L*states_chan+j,lm_index_L*states_chan+j,0,E1(j));
                    Input_Energies(ind2)=std::real(E1(j));
                    ind[0]++;
                    ind2++;

                }
            }

        }
        //so we do not have to calculate the psi1,psi2 again
        //psi3 -> psi2, psi2 -> psi1
        E1=E2;
        psi_1=psi_2;
        d_psi_1=d_psi_2;
        diff_psi_1=diff_psi_2;

        E2=E3;
        psi_2=psi_3;
        d_psi_2=d_psi_3;
        diff_psi_2=diff_psi_3;
        if(rank==0){
            int nnzRank=0;
            for(int k=0; k < H_rank.nmbrLEGO; k++){
                watch(ind[k]);
                nnzRank+=ind[k];
            }
            watch(nnzRank);
            cout << "Total number of elements in 1e+6: " << (long long) p * nnzRank / 1.0e6 << endl;
        }
    }
    //the number of elements every matrix on every rank holds
    nnz_rank=ind[0];
    for(int i = 0; i < H_rank.nmbrLEGO; i++){
        H_rank.nnz[i]=ind[i];
    }
    //write to file
    if(rank ==0){
        dim=ind2;
        watch(ind2);
        std::ofstream f6 ("Energies_Input_shago.txt");
        cout << "Begin to write file" << endl;

              if (f6.is_open()){
               for (int i = 0; i < ind2; i++){
                    f6 << Input_Energies(i) << "\n";
               }
           }

         else cout << "Unable to open file";

        //End of timing
        auto stop = system_clock::now();
        auto duration = duration_cast<milliseconds>(stop - start);
        cout << "Time taken for generating the matrix: " << (double) duration.count()/1000 << "seconds" << endl;
    }

    //Broadcast the dimension of the vector;
    MPI_Bcast(&dim,1,MPI_INT,0,MPI_COMM_WORLD);

    cout << "rank" << " " << rank << " holds elemets: " << nnz_rank << endl;
    H_rank.dim=dim;
    pwatch(dim);
    pwatch(nnz_rank);

    //transform the matrix to CSR after the read in is done
    //we use less memory and obtain a better performance
    auto startSort = system_clock::now();
    if(Inp.Accelerator == "GPU" || Inp.Accelerator == "GPUFull"){
        bp("Before create CSR-GPU");
        H_rank.CSRonDevice();
        bp("After create CSR-GPU");
    }
    else if(Inp.Accelerator == "MKL"){
        bp("Before create CSR-MKL");
        H_rank.CSRMKL();
        bp("After create CSR-MKL");
    }
    else if(Inp.Accelerator != "CPU"){
        cout << "Inp.Accelerator is not a vaild accelerator [CPU,GPU,MKL]" << endl;
        exit(EXIT_FAILURE);

    }
    auto stopSort = system_clock::now();
    auto durationSort = duration_cast<milliseconds>(stopSort - startSort);
    cout << "Time taken for sorting the matrix: " << (double) durationSort.count()/1000 << "seconds" << endl;

}//End of matrix build up

void Propergator(InputParameter &Inp,SpMatrix &H_rank,std::string PATH){
    Inp.E0 = 5.338*sqrt(Inp.I0*1e12)*1e-9;
    Inp.T_puls = Inp.T_cycle * 2 * M_PI / Inp.om;

    int states_chan = (int) round(Inp.b/M_PI*sqrt(2*Inp.Emax));
    //the initial state mind that C++ index starts at 0
    //this depends on Lmax,Mmax & qunatum number n,l,m
    int l_states= std::min(Inp.l_quantum,Inp.Mmax);
    int m_states= std::min(Inp.m_quantum,Inp.Mmax);
    Inp.init_state = states_chan * (Inp.l_quantum * (l_states+1) + m_states)
         +  Inp.n_quantum-Inp.l_quantum-1;

   if(rank ==0){
        //Input param propagator
        watch(Inp.Accelerator);
        watch(Inp.l_quantum);
        watch(Inp.m_quantum);
        watch(Inp.n_quantum);
        watch(Inp.init_state);
        watch(Inp.om);
        watch(Inp.I0);
        watch(Inp.E0);
        watch(Inp.T_puls);
        watch(Inp.T_cycle);
        watch(Inp.Tint);
        watch(Inp.timesteps);
    }

    //initalize the state vector y
    VectorXcd y = VectorXcd::Zero(dim);
    //For the TDSE we start with the electron in the lowest possible state s1
    y(Inp.init_state) = 1.0;
    VectorXcd y_init(dim);
    VectorXcd p_init(Inp.timesteps);
    VectorXcd p_excited_1(Inp.timesteps);
    VectorXcd p_excited_2(Inp.timesteps);
    y_init << y;
    //Time propagator
    double dt = (Inp.T_puls-Inp.Tint)/Inp.timesteps;

    //Only needed on rank 0 but they are small
    MatrixXcd Q = MatrixXcd::Zero(dim,Inp.dim_kryl);  //Eigenvectors
    MatrixXcd h = MatrixXcd::Zero(Inp.dim_kryl,Inp.dim_kryl); //Eigenvalues
    VectorXcd d = VectorXcd::Zero(Inp.dim_kryl);
    VectorXd od = VectorXd::Zero(Inp.dim_kryl-1);

    if(rank==0)
        cout << "decleration of variables finished" << endl;

    //Start timing the propagation
    auto start2 = system_clock::now();

    if(rank==0)
        cout << "start timepropagation" << endl;

    MPI_Barrier(MPI_COMM_WORLD);
    if (Inp.Accelerator == "GPUFull"){
        cuDoubleComplex *dQ;
        cuDoubleComplex *dd;
        cuDoubleComplex *dy;
        double *od;
         cusparseHandle_t handle_sparse;
         cublasHandle_t handle_blas;
        cusolverDnHandle_t handle_solver;
        allocateMemoryPropagator(Inp,H_rank.dim,dQ,dd,od,dy,y.data());
        for(int i = 1; i < Inp.timesteps+1;i++){
            //normalize y

            double normStart;

            double t = dt *i;
            //get ft
            VectorXcd ft =  VectorXcd(116);
            ft_calc(Inp,t, ft.data());

            lanczosCUDA(
                        H_rank.nnz,H_rank.dim,Inp.dim_kryl, H_rank.nmbrLEGO,
                        normStart,
                            H_rank.LEGO_d_cooCols,
                            H_rank.LEGO_d_rowPtr,
                            H_rank.LEGO_d_cooVals_sorted,
                            H_rank.LEGO_d_cscRows,
                            H_rank.LEGO_d_cscColPtr,
                            H_rank.LEGO_d_cscVals,
                            dy,ft.data(),
                            od,dQ,dd);

            //physical propergator
            if (rank ==0){
                physPropagator(Inp,Inp.dim_kryl,dim,normStart,dt,dQ,dd,od,dy,y.data());

                if(i%20 == 0){
                    cout << "Timesteps: " << i << endl;
                    double norm = y.norm();
                watch(y(Inp.init_state));
                watch(y(Inp.init_state+1));
                watch(y(Inp.init_state+2));
                cout << "norm: " << norm << endl;
                cout << "Pinit: " <<  abs(y(Inp.init_state))*abs(y(Inp.init_state)) << endl;

                }

            }
        }

        freeMemoryPropagator(Inp,H_rank.dim,dQ,dd,od,dy,y.data());
                watch(y(Inp.init_state));
                watch(y(Inp.init_state+1));
                watch(y(Inp.init_state+2));
                cout << "Pinit: " <<  abs(y(Inp.init_state))*abs(y(Inp.init_state)) << endl;

    }
    else{
        for(int i = 1; i < Inp.timesteps+1; i++){
            double norm;
            if (rank==0){
                norm = y.norm();
                y = y/norm;
            }

            //get the ft values
            VectorXcd ft =  VectorXcd(116);
            double t = dt * i;
            ft_calc(Inp,t, ft.data());

            //Lanzcos with GPU for SpMV
            if(Inp.Accelerator == "CPU" || Inp.Accelerator == "MKL"){
                lanczos(Inp, H_rank,  dim, y , Q, h, ft);
            }
            cout << "h" << endl;
            cout << h << endl;
            //To get the matrix exponential first diaganolize the matrix the exponate its
            //elements.
            if(rank == 0){
                 startortho = high_resolution_clock::now();

                SelfAdjointEigenSolver<MatrixXcd> ces;
                ces.compute(h);
                MatrixXcd D = (-1i * dt * ces.eigenvalues()).array().exp().matrix().asDiagonal();
                MatrixXcd P = ces.eigenvectors();
                y = norm * Q * (P * D * P.inverse()).col(0);

                stoportho = high_resolution_clock::now();
                durationortho += duration_cast<milliseconds>(stoportho - startortho);

                //write data
                p_init(i-1)=abs(y(Inp.init_state))*abs(y(Inp.init_state));
                p_excited_1(i-1)=abs(y(Inp.init_state+1))*abs(y(Inp.init_state+1));
                p_excited_2(i-1)=abs(y(Inp.init_state+2))*abs(y(Inp.init_state+2));
               if(i%20 == 0){
                    cout << "Timesteps: " << i << endl;
                    norm = y.norm();
                    watch(y(Inp.init_state));
                    watch(y(Inp.init_state+1));
                    watch(y(Inp.init_state+2));
                    cout << "norm: " << norm << endl;
                    cout << "Pinit: " <<  abs(y(Inp.init_state))*abs(y(Inp.init_state)) << endl;

                }
            }
        }
    }
    watch(rank);
    //End of timing
    auto stop2 = system_clock::now();
    auto duration2 = duration_cast<milliseconds>(stop2 - start2);
    if (rank==0){
        cout << "Time taken by propagating the matrix: " << (double) duration2.count()/1000 << "seconds" << endl;

        cout << "Time taken by the actual propergator " << (double) durationortho.count()/1000
            <<"seconds" << endl;

       //write the output
        std::ofstream f1 (PATH + "fort.88");
        f1.precision(15);
        if (f1.is_open()){
            for (int i = 0; i < dim; i++){
                f1 << abs(y(i)) * abs(y(i)) << "\n";
            }
        }
        else cout << "Unable to open file";
        cout << "fort.88 is written" << endl;
        std::ofstream f2 (PATH +"fort.89");
        if (f2.is_open()){
            for (int i = 0; i < dim; i++){
                f2 << i << " " <<  y(i).real() << " " << y(i).imag() << "\n";
            }
        }
        else cout << "Unable to open file";
        cout << "fort.89 is written" << endl;
        //write the output
        std::ofstream f3 (PATH + "pinit.dat");
        f3.precision(15);
        if (f3.is_open()){
            for (int i = 0; i < Inp.timesteps; i++){
                double time = (i+1) * Inp.T_puls/Inp.timesteps;
                f3 << time <<" " << p_init(i).real() << "\n";
            }
        }
        else cout << "Unable to open file";

        std::ofstream f4 (PATH +"pexcited_1.dat");
        f4.precision(15);
        if (f4.is_open()){
            for (int i = 0; i < Inp.timesteps; i++){
                double time = (i+1) * Inp.T_puls/Inp.timesteps;
                f4 << time <<" " << p_excited_1(i).real() << "\n";
            }
        }
        else cout << "Unable to open file";

        std::ofstream f5 (PATH +"pexcited_2.dat");
        f5.precision(15);
        if (f5.is_open()){
            for (int i = 0; i < Inp.timesteps; i++){
                double time = (i+1) * Inp.T_puls/Inp.timesteps;
                f5 << time <<" " << p_excited_2(i).real() << "\n";
            }
        }
        else cout << "Unable to open file";
        cout << "pinit,pexcited are written" << endl;

    }
}

void lanczos(InputParameter Inp,SpMatrix &H_rank, int m,VectorXcd &y ,MatrixXcd &Q, MatrixXcd &h,
        VectorXcd &FT){
    /* Input
    H_rank: mxm matrix
    b: initial vector
    Inp.drim_kryl: dimension of Krylov subspace L in morten code
    m: dimesion of matrix H_rank
    Output:
    Q: orthogonal Krylov space
    h: tridiagonal matrix
    */
    VectorXcd d = VectorXcd::Zero(Inp.dim_kryl);
    VectorXcd od = VectorXcd::Zero(Inp.dim_kryl-1); //L-1 from morten
    VectorXcd dydx_rank = VectorXcd::Zero(dim);
    //will only be used on rank0
    VectorXcd dydx = VectorXcd::Zero(dim);

    Q.setZero();
    Q.col(0) = y;
    h.setZero();

    for (int i = 0; i < Inp.dim_kryl; i++){
        dydx_rank.setZero(dim);
        //broadcast changeing vector y to all ranks
        MPI_Bcast(y.data(),dim*2,MPI_DOUBLE,0,MPI_COMM_WORLD);

        //Matrix vector product this is the only loop that get executed in parallel
        //iterate over all LEGO bricks

        //the cpu COO version
        if(Inp.Accelerator == "CPU"){
            for(int j = 0; j < H_rank.nmbrLEGO; j++){
                H_rank.matrixVectorUpper(j,y,dydx_rank,FT);
            }
        }

        if(Inp.Accelerator == "MKL"){
        // the cpu MKL CSR version
            bp("Before Matrix Vector");
            H_rank.matrixVectorMKL(y,dydx_rank,FT);
            bp("After matrix Vector");
        }
      //sum the local dydx_rank into dydx on rank 0
        MPI_Reduce(dydx_rank.data(),dydx.data(),dim*2,MPI_DOUBLE,MPI_SUM, 0,MPI_COMM_WORLD);

        if(rank==0){
            d(i) = y.adjoint() * dydx;
               //Full reorthogonalization (x2):
            dydx = dydx - Q.block(0,0,dim,i+1) * ( Q.block(0,0,dim,i+1).adjoint() * dydx);
            dydx = dydx - Q.block(0,0,dim,i+1) * ( Q.block(0,0,dim,i+1).adjoint() * dydx);
            if (i < Inp.dim_kryl-1){
                od(i) = dydx.norm();
                if (abs(od(i)) == 0){
                    cout << "devision by 0 during renormalizing" <<endl;
                    break;
                }
                y = dydx/od(i);
                Q.col(i+1) = y;
                //adaptivly control the dimension of the Krylov space. At most we do dim_kryl
                //iterations if we converge earlier we terminate
                if((Q.col(i+1)-Q.col(i)).norm()  < Inp.eps){
                    for(int j = i+1; j < Inp.dim_kryl-1 ; j++){
                        //fill the remaining vectors
                        Q.col(j+1) = y;
                        d(j) = d(i);
                        od(j) = od(i);

                    }
                    d(Inp.dim_kryl-1)=d(i);
                    cout << "Krylov space converged after: " << i << " iterations\n";
                    break;

                }
            }
        }

    }
    //set the h Matrix mortens TT matrix
    for (int i = 1; i < Inp.dim_kryl; i++){
        h(i,i-1) = od(i-1);
        h(i-1,i) = od(i-1);
    }
    for (int i = 0; i < Inp.dim_kryl; i++){
        h(i,i) = d(i);
    }
   return;
}

void polarization(SpMatrix &M,int states_chan, int &ind, int index,int l,int lm_index_L,int lm_index_R,
        double LM_faktor,double dX,ArrayXd &weights, ArrayXd &x, ArrayXXd &psi_1,ArrayXXd
        &psi_2,ArrayXXd &d_psi_2){
    //Lego brick 1 + 2
#pragma omp parallel for collapse(2)
    for(int j =  states_chan_rank_begin; j < states_chan_rank_end; j++){
        for(int i = 0; i < states_chan; i++){
                M.writeElement(ind+j*states_chan+i,
                        lm_index_L*states_chan+j,
                        lm_index_R*states_chan+i,
                        index,
                        dX/2*(weights*psi_1.col(j)*psi_2.col(i)/x).sum()*LM_faktor*(l+1)
                        +dX/2*(weights*psi_1.col(j)*d_psi_2.col(i)).sum()*LM_faktor);
          }

    }

    //number of elements written on one rank over all threads
    ind = ind + (states_chan_rank_end-states_chan_rank_begin) * states_chan;
}

void helperEigenStates(InputParameter Inp, int states_chan,double dX,double &tmp_d,
        ArrayXd &weights, ArrayXd &x,ArrayXXd &B,ArrayXXd &diff_B,MatrixXd &H_base,
        MatrixXd &S){
    //This function should speed up the calculation
    //S and large aprts of H are always the same, by storing them and passing them to
    //getEigenStates redundant computation can be avoided
    //S and the kinetic Energy are the same for all l,

    int dim_H = Inp.N_b+Inp.n-3;
    for (int i = 0; i < dim_H ; i++){
        for (int j = 0; j < dim_H; j++){
            H_base(i,j) = - 0.5 * dX * 0.5 * (weights * B.col(i+1)*diff_B.col(j+1)).sum();
            ArrayXd tmp = dX * 0.5 * (weights * B.col(i+1)*B.col(j+1));
            H_base(i,j) += (tmp * -1/x).sum();
            S(i,j) = (tmp).sum();
        }
    }

}

void getEigenStates(InputParameter Inp, int l,int states_chan,double dX,double &tmp_d,ArrayXd &weights,
        ArrayXd &x,ArrayXXd &B, ArrayXXd &d_B,ArrayXXd &diff_B,ArrayXXd &psi,ArrayXXd &d_psi
        ,ArrayXXd &diff_psi , MatrixXd &H_base, MatrixXd &S, ArrayXd &Energy ){
    //The H and S matrix
    int dim_H = Inp.N_b+Inp.n-3;  //this equals to s_H
    MatrixXd H = MatrixXd::Zero(dim_H,dim_H);
    ArrayXd tmp_l = (l*(1+l)/(2*x*x));

//only works on rect loops
#pragma omp parallel for collapse(2)
    for (int i = 0; i < dim_H ; i++){
        for (int j = 0; j < dim_H; j++){
             //Kinetic Engergy
            H(i,j) = H_base(i,j);
            //potential Energy V
            ArrayXd tmp =  dX * 0.5 * weights * B.col(i+1)*B.col(j+1);
            H(i,j) += (tmp * tmp_l).sum();
       }
    }

    //Start timing the solver
    auto start2 = system_clock::now();
    //The Gen Eigenvalue solver on CPU
    MatrixXd Eigenvectors =  MatrixXd::Zero(dim_H,dim_H);
    if(Inp.Accelerator == "CPU" || Inp.Accelerator== "MKL")
    {
        GeneralizedSelfAdjointEigenSolver<MatrixXd> ges(H,S);
        Eigenvectors = ges.eigenvectors();
        Energy = ges.eigenvalues();
    }
   //Gen Eigenvalue solver on GPU
    //H,S are in coloum major, 0 index format
    //Cusolver assumes col-major no conversion needed
    //lda,ldb,m are all the same as dim W,V = Eigenvalues,Eigenvectors
    if(Inp.Accelerator == "GPU" || Inp.Accelerator =="GPUFull")
    {
        //VectorXd Eigenvalues =  VectorXd::Zero(dim_H);
        genEigenSolverCUDA(dim_H,H.data(),S.data(),Energy.data(),Eigenvectors.data());
    }

    auto stop2 = system_clock::now();
    auto duration2 = duration_cast<milliseconds>(stop2 - start2);
    if (rank==0){
        cout << "Time taken by the solver in l: "<< l << " : "  << (double) duration2.count()/1000
            << "seconds" << endl;}

    //set the psi to 0
    psi.setZero();
    d_psi.setZero();
    diff_psi.setZero();

    //watch(Energy);
    //watch(Eigenvectors);
#pragma omp parallel for collapse(2)
    for(int i = 0; i < states_chan; i++){
        for(int j=0; j < dim_H; j++){
            psi.col(i) = psi.col(i) + Eigenvectors(j,i) * B.col(j+1);
            d_psi.col(i) = d_psi.col(i) + Eigenvectors(j,i) * d_B.col(j+1);
            diff_psi.col(i) = diff_psi.col(i) + Eigenvectors(j,i) * diff_B.col(j+1);
       }
 }
    }


/*    Copyright (c) 2010-2018, Delft University of Technology
 *    All rigths reserved
 *
 *    This file is part of the Tudat. Redistribution and use in source and
 *    binary forms, with or without modification, are permitted exclusively
 *    under the terms of the Modified BSD license. You should have received
 *    a copy of the license with this file. If not, please or visit:
 *    http://tudat.tudelft.nl/LICENSE.
 *
 */

#define BOOST_TEST_MAIN

#include <boost/make_shared.hpp>
#include <boost/test/unit_test.hpp>

#include "Tudat/Mathematics/BasicMathematics/mathematicalConstants.h"

#include "Tudat/Mathematics/Interpolators/lagrangeInterpolator.h"

namespace tudat
{
namespace unit_tests
{

//! Function to evaluate polynomial
/*!
 *  Function to evaluate polynomial with coefficients and independent variable as input.
 *  \param coefficients Polynomial coefficients with the coefficient as map value and order as key.
 *  \param evaluationPoint Independent variable at which polynomial is to be evaluated.
 *  \return Polynomial value.
 */
double evaluatePolynomial( const std::map< int, double >& coefficients,
                           const double evaluationPoint )
{
    double polynomialValue = 0.0;
    for( std::map< int, double >::const_iterator it = coefficients.begin( );
         it != coefficients.end( ) ; it++ )
    {
        polynomialValue += it->second * std::pow( evaluationPoint, it->first );
    }
    return polynomialValue;
}

//! Function to retrieve polynomial coefficients
/*!
 *  Function to retrieve quasi-random polynomial coefficients, up to a given maximum order.
 *  \param polynomialOrder Order of polynomial.
 *  \return Polynomial coefficients with the coefficient as map value and order as key.
 */
std::map< int, double > getPolynomialCoefficients( const int polynomialOrder)
{
    std::map< int, double > allCoefficients;
    allCoefficients[ 0 ] = 8.05425;
    allCoefficients[ 1 ] = 2.540;
    allCoefficients[ 2 ] = -0.454;
    allCoefficients[ 3 ] = 1.1224;
    allCoefficients[ 4 ] = 0.03545;
    allCoefficients[ 5 ] = -0.004;
    allCoefficients[ 6 ] = 0.0784;
    allCoefficients[ 7 ] = -0.000334;
    allCoefficients[ 8 ] = -0.00004743;
    allCoefficients[ 9 ] = 0.000007284;
    allCoefficients[ 10 ] = 0.00000134;
    allCoefficients[ 11 ] = -0.000000324;

    // Copy subset of allCoefficients map into currentCoefficients.
    std::map< int, double > currentCoefficients( allCoefficients.begin(),
        boost::next( allCoefficients.begin(), polynomialOrder + 1 ) );
    return currentCoefficients;

}

//! Create quasi-random vector of non-uniform independent variables
/*!
 *  Create quasi-random vector of non-uniform independent variables
 *  \return Non-uniform, but continuously increasing, set of independent variables.
 */
std::vector< double > getIndependentVariableVector( )
{
    std::vector< double > independentVariableVector;
    independentVariableVector.push_back( 0.0 );
    independentVariableVector.push_back( 0.1 );
    independentVariableVector.push_back( 0.2 );
    independentVariableVector.push_back( 0.3 );
    independentVariableVector.push_back( 0.45 );
    independentVariableVector.push_back( 0.7 );
    independentVariableVector.push_back( 1.0 );
    independentVariableVector.push_back( 1.4 );
    independentVariableVector.push_back( 2.0 );
    independentVariableVector.push_back( 2.1 );
    independentVariableVector.push_back( 2.5 );
    independentVariableVector.push_back( 4.1 );
    independentVariableVector.push_back( 5.7 );
    independentVariableVector.push_back( 6.3 );
    independentVariableVector.push_back( 8.9 );
    independentVariableVector.push_back( 10.2 );
    independentVariableVector.push_back( 11.8 );
    independentVariableVector.push_back( 12.4 );
    independentVariableVector.push_back( 15.5 );
    independentVariableVector.push_back( 16.4 );
    independentVariableVector.push_back( 22.0 );
    independentVariableVector.push_back( 25.0 );
    independentVariableVector.push_back( 30.89 );
    independentVariableVector.push_back( 35.21 );
    independentVariableVector.push_back( 40.38 );
    independentVariableVector.push_back( 43.23 );
    independentVariableVector.push_back( 52.3 );
    independentVariableVector.push_back( 72.0 );
    independentVariableVector.push_back( 89.0 );
    independentVariableVector.push_back( 104.0 );
    return independentVariableVector;
}

BOOST_AUTO_TEST_SUITE( test_lagrange_interpolation )

// Test whetehr Lagrange interpolator can properly reproduce polynomial interpolation
// Since Lagrange interpolation uses a unique (n-1)th order polynomial to fit n data points,
// using an (n-1)th order polynomial as depedent variables should yield an exact reporduction
// of the original polynomial (barring numerical losses).
BOOST_AUTO_TEST_CASE( test_lagrange_interpolation_polynomials )
{
    std::map< double, double > dataMap;
    std::map< int, double > coefficients;
    std::vector< double > independentVariableVector = getIndependentVariableVector( );

    // Test interpolator for 4;6;8;10 data points per interpolant
    // (i.e. 3rd, 5th, 7th and 9th order polynomial)
    for( unsigned int stages = 4; stages < 11; stages += 2 )
    {
        dataMap.clear( );

        // Get polynomial coefficients for current number of points
        coefficients = getPolynomialCoefficients( stages - 1 );

        // Generate dependent variables
        for( unsigned int i = 0; i < independentVariableVector.size( ); i++ )
        {
            dataMap[ independentVariableVector.at( i ) ] =
                    evaluatePolynomial( coefficients, independentVariableVector.at( i ) );
        }

        // Create interpolator
        interpolators::LagrangeInterpolator< double, double > interpolator =
                interpolators::LagrangeInterpolator< double, double >(
                    dataMap, stages, interpolators::huntingAlgorithm,
                    interpolators::lagrange_no_boundary_interpolation );

        // Iterate over all data points inside allowed (i.e. non-boundary) range
        int offsetEntries = stages / 2 - 1;
        for( unsigned int i = offsetEntries;
             i < independentVariableVector.size( ) - ( offsetEntries + 2 ); i++ )
        {
            // Test current interval at 10 equispaced points
            double currentStepSize =  ( independentVariableVector.at( i + 1 ) -
                                        independentVariableVector.at( i ) ) / 10.0;
            for( unsigned j = 0; j < 10; j ++ )
            {
                double currentDataPoint = independentVariableVector.at( i ) +
                        static_cast< double >( j ) * currentStepSize;

                // Check interpolated value against theoretical polynomial
                if( j < 9 )
                {
                    BOOST_CHECK_CLOSE_FRACTION( interpolator.interpolate( currentDataPoint ),
                                                evaluatePolynomial( coefficients, currentDataPoint ),
                                                5.0E-15 );
                }
                else
                {
                    BOOST_CHECK_CLOSE_FRACTION( interpolator.interpolate( currentDataPoint ),
                                                evaluatePolynomial( coefficients, currentDataPoint ),
                                                2.0E-14 );
                }

            }
        }
    }
}

// Test to check whether the various boundary handling methopds are properly implemented
BOOST_AUTO_TEST_CASE( test_lagrange_interpolation_boundary )
{
    std::vector< double > dataVector;
    std::map< int, double > coefficients;
    std::vector< double > independentVariableVector = getIndependentVariableVector( );

    unsigned int independentVariableVectorSize = independentVariableVector.size( );

    {
        // Test interpolator for 4;6;8;10 data points per interpolant
        // (i.e. 3rd, 5th, 7th and 9th order polynomial)
        for( unsigned int stages = 4; stages < 11; stages += 2 )
        {
            dataVector.clear( );

            // Get polynomial coefficients for current number of points
            coefficients = getPolynomialCoefficients( stages - 1 );

            // Generate dependent variables
            for( unsigned int i = 0; i < independentVariableVectorSize; i++ )
            {
                dataVector.push_back( evaluatePolynomial(
                                          coefficients, independentVariableVector.at( i ) ) );
            }

            // Create interpolator with cubic spline interpolation at boundaries
            interpolators::LagrangeInterpolator< double, double > lagrangeInterpolator =
                    interpolators::LagrangeInterpolator< double, double >(
                        independentVariableVector, dataVector, stages,
                        interpolators::huntingAlgorithm,
                        interpolators::lagrange_cubic_spline_boundary_interpolation );

            // Create spline interpolator from edge points at lower bound.
            int dataPointsForSpline = ( stages / 2 > 4 ) ? ( stages / 2 ) : 4;
            std::map< double, double > boundaryMap;
            for( int i = 0; i < dataPointsForSpline; i++ )
            {
                boundaryMap[ independentVariableVector.at( i ) ] = dataVector.at( i );
            }
            interpolators::CubicSplineInterpolator< double, double > lowerBoundInterpolator =
                    interpolators::CubicSplineInterpolator< double, double >( boundaryMap );

            // Test whether the Lagrange interpolators correctly evaluate the cubic spline
            // polynomials at the lower edges.
            for( unsigned int i = 0; i < stages / 2 - 1; i++ )
            {
                double currentStepSize = ( independentVariableVector.at( i + 1 ) -
                                           independentVariableVector.at( i ) ) / 10.0;
                for( unsigned int j = 0; j < 10; j ++ )
                {
                    double currentTestIndependentVariable = independentVariableVector.at( i ) +
                            static_cast< double >( i ) * currentStepSize;
                    BOOST_CHECK_EQUAL(
                                lagrangeInterpolator.interpolate( currentTestIndependentVariable ),
                                lowerBoundInterpolator.interpolate(
                                    currentTestIndependentVariable ) );
                }
            }

            // Create spline interpolator from edge points at upper bound.
            boundaryMap.clear( );
            for( unsigned int i = independentVariableVectorSize - dataPointsForSpline;
                 i < independentVariableVectorSize; i++ )
            {
                boundaryMap[ independentVariableVector.at( i ) ] = dataVector.at( i );
            }
            interpolators::CubicSplineInterpolator< double, double > upperBoundInterpolator =
                    interpolators::CubicSplineInterpolator< double, double >( boundaryMap );

            // Test whether the Lagrange interpolators correctly evaluate the cubic spline
            // polynomials at the uper edges.
            for( unsigned int i = 0; i < stages / 2 - 1; i++ )
            {
                double currentStepSize =
                        ( independentVariableVector.at( independentVariableVectorSize - i - 1 ) -
                          independentVariableVector.at( independentVariableVectorSize - i - 2 ) ) /
                        10.0;
                for( unsigned int j = 0; j < 10; j ++ )
                {
                    double currentTestIndependentVariable = independentVariableVector.at(
                                independentVariableVectorSize - i - 2 ) +
                            static_cast< double >( i ) * currentStepSize;
                    BOOST_CHECK_EQUAL( lagrangeInterpolator.interpolate(
                                           currentTestIndependentVariable ),
                                       upperBoundInterpolator.interpolate(
                                           currentTestIndependentVariable ) );
                }
            }
        }
    }

    // Test whether an error is thrown if lagrange_no_boundary_interpolation is
    // selected an interpolation at the boundaries is requested.
    bool runtimeErrorOccurred;
    {
        // Test interpolators with various number of stages
        for( unsigned int stages = 4; stages < 11; stages += 2 )
        {
            dataVector.clear( );

            // Get polynomial coefficients for current number of points
            coefficients = getPolynomialCoefficients( stages - 1 );

            // Generate dependent variables
            for( unsigned int i = 0; i < independentVariableVectorSize; i++ )
            {
                dataVector.push_back( evaluatePolynomial(
                                          coefficients, independentVariableVector.at( i ) ) );
            }

            // Create interpolator lagrange_no_boundary_interpolation
            interpolators::LagrangeInterpolator< double, double > lagrangeInterpolator =
                    interpolators::LagrangeInterpolator< double, double >(
                        independentVariableVector, dataVector, stages,
                        interpolators::huntingAlgorithm,
                        interpolators::lagrange_no_boundary_interpolation );

            // Test for each whether the interpolator correctly throws an exception if
            // interpolation at the lower boundary is requested.
            for( unsigned int i = 0; i < stages / 2 - 1; i++ )
            {
                double currentStepSize =
                        ( independentVariableVector.at( i + 1 ) -
                          independentVariableVector.at( i ) ) / 3.0;

                for( unsigned int j = 0; j < 3; j ++ )
                {
                    try
                    {
                        double currentTestIndependentVariable =
                                independentVariableVector.at( i ) +
                                static_cast< double >( j ) * currentStepSize;
                        lagrangeInterpolator.interpolate( currentTestIndependentVariable );

                    }
                    catch( std::runtime_error )
                    {
                        runtimeErrorOccurred = 1;
                    }
                    BOOST_CHECK_EQUAL( runtimeErrorOccurred, 1 );
                    runtimeErrorOccurred = 0;
                }
            }

            // Test for each whether the interpolator correctly throws an exception if
            // interpolation at the upper boundary is requested.
            for( unsigned int i = 0; i < stages / 2 - 1; i++ )
            {
                double currentStepSize = ( independentVariableVector.at(
                                               independentVariableVectorSize - i - 1 ) -
                                           independentVariableVector.at(
                                               independentVariableVectorSize - i - 2 ) ) / 3.0;
                for( unsigned int j = 0; j < 3; j ++ )
                {
                    try
                    {
                        double currentTestIndependentVariable =
                                independentVariableVector.at(
                                    independentVariableVectorSize - i - 2 ) +
                                static_cast< double >( j ) * currentStepSize;
                        lagrangeInterpolator.interpolate( currentTestIndependentVariable );

                    }
                    catch( std::runtime_error )
                    {
                        runtimeErrorOccurred = true;
                    }
                    BOOST_CHECK_EQUAL( runtimeErrorOccurred, 1 );
                    runtimeErrorOccurred = false;
                }
            }
        }
    }
}

// Test to check whether the various error handling methods are correctly implemented
BOOST_AUTO_TEST_CASE( test_lagrange_error_checks )
{
    std::map< double, double > dataMap;
    std::vector< double > dataVector;
    std::map< int, double > coefficients;
    std::vector< double > independentVariableVector = getIndependentVariableVector( );

    bool runtimeErrorOccurred;

    {
        // Create interpolator with empty data map
        runtimeErrorOccurred = false;
        try
        {
            interpolators::LagrangeInterpolator< double, double > interpolator =
                    interpolators::LagrangeInterpolator< double, double >(
                        dataMap, 8, interpolators::huntingAlgorithm,
                        interpolators::lagrange_no_boundary_interpolation );
        }
        catch( std::runtime_error )
        {
            runtimeErrorOccurred = true;
        }
        BOOST_CHECK_EQUAL( runtimeErrorOccurred, 1 );
        runtimeErrorOccurred = false;

        // Create interpolator with empty independent variable vector
        runtimeErrorOccurred = false;
        try
        {
            dataVector.push_back( 1.0 );
            interpolators::LagrangeInterpolator< double, double > interpolator =
                    interpolators::LagrangeInterpolator< double, double >(
                        independentVariableVector, dataVector, 8,
                        interpolators::huntingAlgorithm,
                        interpolators::lagrange_no_boundary_interpolation );
        }
        catch( std::runtime_error )
        {
            runtimeErrorOccurred = true;
        }
        BOOST_CHECK_EQUAL( runtimeErrorOccurred, 1 );
        runtimeErrorOccurred = false;
        dataVector.clear( );

        // Create interpolator with empty dependent variable vector
        bool runtimeErrorOccurred = false;
        try
        {
            independentVariableVector.push_back( 1.0 );
            interpolators::LagrangeInterpolator< double, double > interpolator =
                    interpolators::LagrangeInterpolator< double, double >(
                        independentVariableVector, dataVector, 8,
                        interpolators::huntingAlgorithm,
                        interpolators::lagrange_no_boundary_interpolation );
        }
        catch( std::runtime_error )
        {
            runtimeErrorOccurred = true;
        }
        BOOST_CHECK_EQUAL( runtimeErrorOccurred, 1 );
        runtimeErrorOccurred = false;
        independentVariableVector = getIndependentVariableVector( );
    }

    // Create interpolator with NaN first entry (cannot make zero value)
    {
        // Create interpolator with NaN first entry from map constructor
        runtimeErrorOccurred = false;
        try
        {
            coefficients = getPolynomialCoefficients( 7 );
            dataMap[ independentVariableVector.at( 0 ) ] = TUDAT_NAN;
            for( unsigned int i = 1; i < independentVariableVector.size( ); i++ )
            {
                dataMap[ independentVariableVector.at( i ) ] = evaluatePolynomial(
                            coefficients, independentVariableVector.at( i ) );
            }
            interpolators::LagrangeInterpolator< double, double > interpolator =
                    interpolators::LagrangeInterpolator< double, double >(
                        dataMap, 8, interpolators::huntingAlgorithm,
                        interpolators::lagrange_no_boundary_interpolation );
        }
        catch( std::runtime_error )
        {
            runtimeErrorOccurred = true;
        }
        BOOST_CHECK_EQUAL( runtimeErrorOccurred, 1 );
        runtimeErrorOccurred = false;
        coefficients.clear( );
        dataMap.clear( );

        // Create interpolator with NaN first entry from vectors constructor
        runtimeErrorOccurred = false;
        try
        {
            coefficients = getPolynomialCoefficients( 7 );
            dataMap[ independentVariableVector.at( 0 ) ] = TUDAT_NAN;
            for( unsigned int i = 1; i < independentVariableVector.size( ); i++ )
            {
                dataVector.push_back( evaluatePolynomial(
                                          coefficients, independentVariableVector.at( i ) ) );
            }
            interpolators::LagrangeInterpolator< double, double > interpolator =
                    interpolators::LagrangeInterpolator< double, double >(
                        independentVariableVector, dataVector, 8,
                        interpolators::huntingAlgorithm,
                        interpolators::lagrange_no_boundary_interpolation );
        }
        catch( std::runtime_error )
        {
            runtimeErrorOccurred = true;
        }
        BOOST_CHECK_EQUAL( runtimeErrorOccurred, 1 );
        runtimeErrorOccurred = false;
        coefficients.clear( );
        dataVector.clear( );
    }

    // Test error throwing when making Lagrange Interpolator with odd number of stages
    {
        // Test for a range of odd number of stages
        for( unsigned int numberOfStages = 1; numberOfStages < 12; numberOfStages+= 2 )
        {
            // Create interpolator with odd number of stages for map constructor
            runtimeErrorOccurred = false;
            try
            {
                coefficients = getPolynomialCoefficients( numberOfStages - 1 );
                for( unsigned int i = 0; i < independentVariableVector.size( ); i++ )
                {
                    dataMap[ independentVariableVector.at( i ) ] = evaluatePolynomial(
                                coefficients, independentVariableVector.at( i ) );
                }
                interpolators::LagrangeInterpolator< double, double > interpolator =
                        interpolators::LagrangeInterpolator< double, double >(
                            dataMap, numberOfStages, interpolators::huntingAlgorithm,
                            interpolators::lagrange_no_boundary_interpolation );
            }
            catch( std::runtime_error )
            {
                runtimeErrorOccurred = true;
            }
            BOOST_CHECK_EQUAL( runtimeErrorOccurred, 1 );
            runtimeErrorOccurred = false;
            coefficients.clear( );
            dataMap.clear( );

            // Create interpolator with odd number of stages for vectors constructor
            runtimeErrorOccurred = false;
            try
            {
                coefficients = getPolynomialCoefficients( numberOfStages - 1 );
                for( unsigned int i = 0; i < independentVariableVector.size( ); i++ )
                {
                    dataVector.push_back( evaluatePolynomial(
                                              coefficients, independentVariableVector.at( i ) ) );
                }
                interpolators::LagrangeInterpolator< double, double > interpolator =
                        interpolators::LagrangeInterpolator< double, double >(
                            independentVariableVector, dataVector, numberOfStages,
                            interpolators::huntingAlgorithm,
                            interpolators::lagrange_no_boundary_interpolation );
            }
            catch( std::runtime_error )
            {
                runtimeErrorOccurred = true;
            }
            BOOST_CHECK_EQUAL( runtimeErrorOccurred, 1 );
            runtimeErrorOccurred = false;
            coefficients.clear( );
            dataVector.clear( );
        }
    }

    // Test error throwing when making Lagrange Interpolator from vectors constructor with
    // differently sized (in)dependent variable vectors
    {
        int numberOfStages = 8;
        runtimeErrorOccurred = false;
        try
        {
            coefficients = getPolynomialCoefficients( numberOfStages - 1 );
            for( unsigned int i = 0; i < independentVariableVector.size( ) - 1; i++ )
            {
                dataVector.push_back( evaluatePolynomial(
                                          coefficients, independentVariableVector.at( i ) ) );
            }
            interpolators::LagrangeInterpolator< double, double > interpolator =
                    interpolators::LagrangeInterpolator< double, double >(
                        independentVariableVector, dataVector, numberOfStages,
                        interpolators::huntingAlgorithm,
                        interpolators::lagrange_no_boundary_interpolation );
        }
        catch( std::runtime_error )
        {
            runtimeErrorOccurred = true;
        }
        BOOST_CHECK_EQUAL( runtimeErrorOccurred, 1 );
        runtimeErrorOccurred = false;
        coefficients.clear( );
        dataVector.clear( );
    }
}

BOOST_AUTO_TEST_SUITE_END( )

}

}


#include <boost/format.hpp>
#include <boost/math/constants/constants.hpp>
#include <chrono>
#include <fstream>
#include <iostream>
#include <random>
#include <utility>
#include <vector>
#include "planner/planner.h"

using std::milli;
using boost::format;
using boost::math::constants::sixth_pi;
using boost::math::constants::two_pi;

planner::polygon get_cone(planner::point pos, double cone_radius = 0.3) {
  planner::polygon p;
  for (double i = 0; i < two_pi<double>(); i += sixth_pi<double>()) {
    double x = pos.x() + cone_radius * cos(i);
    double y = pos.y() + cone_radius * sin(i);
    p.outer().push_back(planner::point(x, y));
  }
  return p;
}

double bench_obstacles(int samples, int nr_obstacles) {
  double L = 1000;
  std::vector<std::pair<double, double>> wp = {};

  for (double i = 0; i < L; ++i) {
    wp.push_back(std::make_pair(i, 0));
  }

  double s = 0;

  planner::Planner p(wp);
  planner::pose pos{p.baseframe().P(s), p.baseframe().theta(s), s, 0.0};

  std::vector<planner::polygon> obstacles;
  std::random_device rd;
  std::mt19937 gen(rd());
  std::uniform_real_distribution<> length_dist(0.0, p.baseframe().length());
  std::uniform_real_distribution<> distance_dist(3.5, 10);
  std::discrete_distribution<> side_dist({-1, 1});

  for (int i = 0; i < nr_obstacles; ++i) {
    planner::point pos(length_dist(gen), side_dist(gen) * distance_dist(gen));
    obstacles.push_back(get_cone(pos));
  }

  double time_milli = 0.0;
  int idx;

  for (idx = 0; idx <= samples && s <= p.baseframe().length() - 1; idx++) {
    auto start = std::chrono::steady_clock::now();
    auto m_opt = p.update(pos.position, pos.theta, 5, 1, obstacles);
    auto end = std::chrono::steady_clock::now();
    auto diff = end - start;
    time_milli += std::chrono::duration<double, milli>(diff).count();
    pos = m_opt.path()[1];
    s = m_opt.path()[1].s;
  }
  double time_milli_mean = time_milli / idx;
  return time_milli_mean;
}

int main(int argc, char *argv[]) {
  std::ofstream os("bench_obstacles.tex");

  os << "% Autor: Roman C. Podolski <roman.podolski@tum.de>\n";
  os << "% this is a generated file, please do not alter\n\n";
  os << "\\documentclass{standalone}[preview]\n";
  os << "\\usepackage{tikz}\n";
  os << "\\usetikzlibrary{decorations.markings,calc}\n";
  os << "\\usepackage{pgfplots}\n";
  os << "\\usepackage{pgfplotstable}\n";
  os << "\\usepackage[svgnames]{xcolor}\n";
  os << "% Settings for pgfplots\n";
  os << "\\pgfplotsset{compat=1.9}\n";
  os << "\\pgfplotsset{cycle list={black}}\n";
  os << "\\begin{document}\n";
  os << "\\begin{tikzpicture}[scale=1.0]\n";
  os << boost::format("\\begin{axis}[%s]\n") % "";
  os << boost::format("\\addplot[%s] coordinates {\n") % "";
  for (int i = 0; i < 5000; i += 10) {
    double time = bench_obstacles(100, i);
    std::cout << i << " : " << time << '\n';
    os << boost::format("(%d,%d)\n") % i % time;
  }
  os << "};";
  os << "\\end{axis}\n";
  os << "\\end{tikzpicture}\n";
  os << "\\end{document}\n";
}


///////////////////////////////////////////////////////////////////////////////////////
// distribution::survival::models::importance_sampling::log_ratio_pdf_statistics.hpp //
//                                                                           		 //
//  Copyright 2009 Erwann Rogard. Distributed under the Boost                		 //
//  Software License, Version 1.0. (See accompanying file                    		 //
//  LICENSE_1_0.txt or copy at http://www.boost.org/LICENSE_1_0.txt)         		 //
///////////////////////////////////////////////////////////////////////////////////////
#ifndef BOOST_STATISTICS_DETAIL_DISTRIBUTION_SURVIVAL_MODELS_IMPORTANCE_SAMPLING_LOG_RATIO_PDF_STATISTICS_HPP_ER_2009
#define BOOST_STATISTICS_DETAIL_DISTRIBUTION_SURVIVAL_MODELS_IMPORTANCE_SAMPLING_LOG_RATIO_PDF_STATISTICS_HPP_ER_2009
#include <boost/fusion/include/make_map.hpp>
#include <boost/statistics/detail/importance_sampling/statistics/percentage_effective_sample_size.hpp> //
#include <boost/statistics/detail/math/constant/shortcut.hpp>

namespace boost{
namespace statistics{
namespace detail{
namespace distribution{
namespace survival{

	template<typename T>
	struct log_ratio_pdf_statistics{
        typedef accumulator::tag::percentage_effective_sample_size tag_ess_;
        typedef accumulator::tag::proportion_less_than tag_plt_;
        typedef boost::accumulators::stats<tag_ess_,tag_plt_> stats_;
        typedef boost::accumulators::accumulator_set<T,stats_> acc_;

		public:

		log_ratio_pdf_statistics(){}

    	template<typename It_w,typename T1>
        void operator()(It_w b_w,It_w e_w,const T1& eps)
    	{
    		typedef T val_;
        	acc_ a = std::for_each(
            	b_w,
            	e_w,
            	acc_(( accumulator::keyword::threshold = eps ))
        	);
            this->ess_ = boost::accumulators::extract_result<tag_ess_>(a);
            this->plt_ = boost::accumulators::extract_result<tag_plt_>(a);
		}

    	template<typename It_w,typename T1>
        void operator()(It_w b_w,It_w e_w){
			return (*this)(b_w,e_w,boost::math::tools::epsilon<T>());
		}

		T ess()const{ return this->ess_; }
		T percent_lt_eps()const{ return this->plt_; }

        private:
        T ess_,plt_;
    };

}// survival
}// distribution
}// detail
}// statistics
}// boost

#endif


/**
 * @file
 * @license BSD 3-clause
 * @copyright Copyright (c) 2020, New York University and Max Planck
 * Gesellschaft
 *
 * @brief Implements a PD controller at the center of mass.
 *
 */

#pragma once

#include <Eigen/Dense>

namespace blmc_controllers
{
typedef Eigen::Array<double, 6, 1> Array6d;
typedef Eigen::Matrix<double, 6, 1> Vector6d;

/**
 * @brief Impedance controller between any two frames of the robot.
 */
class CentroidalPDController
{
public:
    /**
     * @brief Construct a new ImpedanceController object.
     */
    CentroidalPDController();

    /**
     * @brief Initialize the internal data. None real-time safe method.
     *
     * @param inertia Inertia of the base.
     */
    void initialize(double& mass, Eigen::Ref<const Eigen::Vector3d> inertia);

    /**
     * Computes the centroidal wrench using a PD controller.
     */
    void run(Eigen::Ref<const Eigen::Vector3d> kc,
             Eigen::Ref<const Eigen::Vector3d> dc,
             Eigen::Ref<const Eigen::Vector3d> kb,
             Eigen::Ref<const Eigen::Vector3d> db,
             Eigen::Ref<const Eigen::Vector3d> com,
             Eigen::Ref<const Eigen::Vector3d> com_des,
             Eigen::Ref<const Eigen::Vector3d> vcom,
             Eigen::Ref<const Eigen::Vector3d> vcom_des,
             Eigen::Ref<const Eigen::Vector4d> ori,
             Eigen::Ref<const Eigen::Vector4d> ori_des,
             Eigen::Ref<const Eigen::Vector3d> angvel,
             Eigen::Ref<const Eigen::Vector3d> angvel_des);

    /**
     * @brief Get the computed wrench from the PD controller.
     *
     * @return Eigen::VectorXd&
     */
    Vector6d& get_wrench();

private:  // attributes
    /** @brief Output wrench */
    Vector6d wrench_;

    double mass_;
    Eigen::Vector3d inertia_;

    Eigen::Vector3d pos_error_;
    Eigen::Vector3d vel_error_;
    Eigen::Vector3d ori_error_;

    Eigen::Quaternion<double> ori_quat_;
    Eigen::Quaternion<double> des_ori_quat_;
    Eigen::Quaternion<double> ori_error_quat_;

    Eigen::Matrix<double, 3, 3> ori_se3_;
    Eigen::Matrix<double, 3, 3> des_ori_se3_;
    Eigen::Matrix<double, 3, 3> ori_error_se3_; // refer to christian ott paper for definitions (Rdb)

};

}  // namespace blmc_controllers


// -*- mode: c++; fill-column: 80; indent-tabs-mode: nil; -*-

#ifndef ARMA_LAPACK_EXTRA_HPP
#define ARMA_LAPACK_EXTRA_HPP

#include <armadillo>

#if !defined(ARMA_BLAS_CAPITALS)

#define arma_sbdsqr sbdsqr
#define arma_dbdsqr dbdsqr

/* Least-square problem */
#define arma_sgglse sgglse
#define arma_dgglse dgglse
#define arma_cgglse cgglse
#define arma_zgglse zgglse

/* QR factorization with column pivoting */
#define arma_sgeqp3 sgeqp3
#define arma_dgeqp3 dgeqp3
#define arma_cgeqp3 cgeqp3
#define arma_zgeqp3 zgeqp3

/* Eigenvalue decomposition of upper Hessenberg matrix */
#define arma_shseqr shseqr
#define arma_dhseqr dhseqr
#define arma_chseqr chseqr
#define arma_zhseqr zhseqr

/* Jacobi SVD */
#define arma_sgesvj sgesvj
#define arma_dgesvj dgesvj
#define arma_cgesvj cgesvj
#define arma_zgesvj zgesvj

/* Another Jacobi SVD */
#define arma_sgejsv sgejsv
#define arma_dgejsv dgejsv
#define arma_cgejsv cgejsv
#define arma_zgejsv zgejsv

#else

#define arma_sbdsqr SBDSQR
#define arma_dbdsqr DBDSQR

/* Least-square problem */
#define arma_sgglse SGGLSE
#define arma_dgglse DGGLSE
#define arma_cgglse CGGLSE
#define arma_zgglse ZGGLSE

/* QR factorization with column pivoting */
#define arma_sgeqp3 SGEQP3
#define arma_dgeqp3 DGEQP3
#define arma_cgeqp3 CGEQP3
#define arma_zgeqp3 ZGEQP3

/* Eigenvalue decomposition of upper Hessenberg matrix */
#define arma_shseqr SHSEQR
#define arma_dhseqr DHSEQR
#define arma_chseqr CHSEQR
#define arma_zhseqr ZHSEQR

/* Jacobi SVD */
#define arma_sgesvj SGESVJ
#define arma_dgesvj DGESVJ
#define arma_cgesvj CGESVJ
#define arma_zgesvj ZGESVJ

/* Another Jacobi SVD */
#define arma_sgejsv SGEJSV
#define arma_dgejsv DGEJSV
#define arma_cgejsv CGEJSV
#define arma_zgejsv ZGEJSV

#endif /* ARMA_BLAS_CAPITALS */

namespace arma
{

extern "C" {
// SVD of real bidiagonal matrix
void arma_fortran_noprefix(arma_sbdsqr)(char* uplo, blas_int* n, blas_int* ncvt,
                                        blas_int* nru, blas_int* ncc, float* d,
                                        float* e, float* vt, blas_int* ldvt,
                                        float* u, blas_int* ldu, float* c,
                                        blas_int* ldc, float* work,
                                        blas_int* info);

void arma_fortran_noprefix(arma_dbdsqr)(char* uplo, blas_int* n, blas_int* ncvt,
                                        blas_int* nru, blas_int* ncc, double* d,
                                        double* e, double* vt, blas_int* ldvt,
                                        double* u, blas_int* ldu, double* c,
                                        blas_int* ldc, double* work,
                                        blas_int* info);

// Solve generalized eigenvalue problem
void arma_fortran_noprefix(arma_sggev)(char* jobl, char* jobr, blas_int* n,
                                       float* a, blas_int* lda, float* b,
                                       blas_int* ldb, float* alphar,
                                       float* alphai, float* beta, float* vl,
                                       blas_int* ldvl, float* vr,
                                       blas_int* ldvr, float* work,
                                       blas_int* lwork, blas_int* info);

void arma_fortran_noprefix(arma_dggev)(char* jobl, char* jobr, blas_int* n,
                                       double* a, blas_int* lda, double* b,
                                       blas_int* ldb, double* alphar,
                                       double* alphai, double* beta, double* vl,
                                       blas_int* ldvl, double* vr,
                                       blas_int* ldvr, double* work,
                                       blas_int* lwork, blas_int* info);

void arma_fortran_noprefix(arma_cggev)(
    char* jobl, char* jobr, blas_int* n, void* a, blas_int* lda, void* b,
    blas_int* ldb, void* alpha, void* beta, void* vl, blas_int* ldvl, void* vr,
    blas_int* ldvr, void* work, blas_int* lwork, float* rwork, blas_int* info);

void arma_fortran_noprefix(arma_zggev)(
    char* jobl, char* jobr, blas_int* n, void* a, blas_int* lda, void* b,
    blas_int* ldb, void* alpha, void* beta, void* vl, blas_int* ldvl, void* vr,
    blas_int* ldvr, void* work, blas_int* lwork, double* rwork, blas_int* info);

// linear equality-constrained least squares problem (LSE)
void arma_fortran_noprefix(arma_sgglse)(blas_int* m, blas_int* n, blas_int* p,
                                        float* a, blas_int* lda, float* b,
                                        blas_int* ldb, float* c, float* d,
                                        float* x, float* work, blas_int* lwork,
                                        blas_int* info);

void arma_fortran_noprefix(arma_dgglse)(blas_int* m, blas_int* n, blas_int* p,
                                        double* a, blas_int* lda, double* b,
                                        blas_int* ldb, double* c, double* d,
                                        double* x, double* work,
                                        blas_int* lwork, blas_int* info);

void arma_fortran_noprefix(arma_cgglse)(blas_int* m, blas_int* n, blas_int* p,
                                        void* a, blas_int* lda, void* b,
                                        blas_int* ldb, void* c, void* d,
                                        void* x, void* work, blas_int* lwork,
                                        blas_int* info);

void arma_fortran_noprefix(arma_zgglse)(blas_int* m, blas_int* n, blas_int* p,
                                        void* a, blas_int* lda, void* b,
                                        blas_int* ldb, void* c, void* d,
                                        void* x, void* work, blas_int* lwork,
                                        blas_int* info);
// QR factorization with column pivoting
void arma_fortran_noprefix(arma_sgeqp3)(blas_int* m, blas_int* n, float* a,
                                        blas_int* lda, blas_int* jpiv,
                                        float* tau, float* work,
                                        blas_int* lwork, blas_int* info);

void arma_fortran_noprefix(arma_dgeqp3)(blas_int* m, blas_int* n, double* a,
                                        blas_int* lda, blas_int* jpiv,
                                        double* tau, double* work,
                                        blas_int* lwork, blas_int* info);

void arma_fortran_noprefix(arma_cgeqp3)(blas_int* m, blas_int* n, void* a,
                                        blas_int* lda, blas_int* jpiv,
                                        void* tau, void* work, blas_int* lwork,
                                        float* rwork, blas_int* info);
void arma_fortran_noprefix(arma_zgeqp3)(blas_int* m, blas_int* n, void* a,
                                        blas_int* lda, blas_int* jpiv,
                                        void* tau, void* work, blas_int* lwork,
                                        double* rwork, blas_int* info);

// xHSEQR --- Eigendecomposition of upper Hessenberg matrix by mutishift QR
void arma_fortran_noprefix(arma_shseqr)(char* job, char* compz, blas_int* n,
                                        blas_int* ilo, blas_int* ihi, float* h,
                                        blas_int* ldh, float* wr, float* wi,
                                        float* z, blas_int* ldz, float* work,
                                        blas_int* lwork, blas_int* info);

void arma_fortran_noprefix(arma_dhseqr)(char* job, char* compz, blas_int* n,
                                        blas_int* ilo, blas_int* ihi, double* h,
                                        blas_int* ldh, double* wr, double* wi,
                                        double* z, blas_int* ldz, double* work,
                                        blas_int* lwork, blas_int* info);

void arma_fortran_noprefix(arma_chseqr)(char* job, char* compz, blas_int* n,
                                        blas_int* ilo, blas_int* ihi, void* h,
                                        blas_int* ldh, void* w, void* z,
                                        blas_int* ldz, void* work,
                                        blas_int* lwork, blas_int* info);

void arma_fortran_noprefix(arma_zhseqr)(char* job, char* compz, blas_int* n,
                                        blas_int* ilo, blas_int* ihi, void* h,
                                        blas_int* ldh, void* w, void* z,
                                        blas_int* ldz, void* work,
                                        blas_int* lwork, blas_int* info);

// xGESVJ --- Jacobi SVD
void arma_fortran_noprefix(arma_sgesvj)(char* joba, char* jobu, char* jobv,
                                        blas_int* m, blas_int* n, float* a,
                                        blas_int* lda, float* sva, blas_int* mv,
                                        float* v, blas_int* ldv, float* work,
                                        blas_int* lwork, blas_int* info);
void arma_fortran_noprefix(arma_dgesvj)(char* joba, char* jobu, char* jobv,
                                        blas_int* m, blas_int* n, double* a,
                                        blas_int* lda, double* sva,
                                        blas_int* mv, double* v, blas_int* ldv,
                                        double* work, blas_int* lwork,
                                        blas_int* info);
void arma_fortran_noprefix(arma_cgesvj)(char* joba, char* jobu, char* jobv,
                                        blas_int* m, blas_int* n, void* a,
                                        blas_int* lda, float* sva, blas_int* mv,
                                        void* v, blas_int* ldv, void* cwork,
                                        blas_int* lwork, float* rwork,
                                        blas_int* lrwork, blas_int* info);
void arma_fortran_noprefix(arma_zgesvj)(char* joba, char* jobu, char* jobv,
                                        blas_int* m, blas_int* n, void* a,
                                        blas_int* lda, double* sva,
                                        blas_int* mv, void* v, blas_int* ldv,
                                        void* cwork, blas_int* lwork,
                                        double* rwork, blas_int* lrwork,
                                        blas_int* info);

// xGEJSV --- Jacobi SVD
void arma_fortran_noprefix(arma_sgejsv)(char* joba, char* jobu, char* jobv,
                                        char* jobr, char* jobt, char* jobp,
                                        blas_int* m, blas_int* n, float* a,
                                        blas_int* lda, float* sva, float* u,
                                        blas_int* ldu, float* v, blas_int* ldv,
                                        float* work, blas_int* lwork,
                                        blas_int* iwork, blas_int* info);

void arma_fortran_noprefix(arma_dgejsv)(char* joba, char* jobu, char* jobv,
                                        char* jobr, char* jobt, char* jobp,
                                        blas_int* m, blas_int* n, double* a,
                                        blas_int* lda, double* sva, double* u,
                                        blas_int* ldu, double* v, blas_int* ldv,
                                        double* work, blas_int* lwork,
                                        blas_int* iwork, blas_int* info);

void arma_fortran_noprefix(arma_cgejsv)(
    char* joba, char* jobu, char* jobv, char* jobr, char* jobt, char* jobp,
    blas_int* m, blas_int* n, void* a, blas_int* lda, float* sva, void* u,
    blas_int* ldu, void* v, blas_int* ldv, void* cwork, blas_int* lwork,
    float* rwork, blas_int* lrwork, blas_int* iwork, blas_int* info);

void arma_fortran_noprefix(arma_zgejsv)(
    char* joba, char* jobu, char* jobv, char* jobr, char* jobt, char* jobp,
    blas_int* m, blas_int* n, void* a, blas_int* lda, double* sva, void* u,
    blas_int* ldu, void* v, blas_int* ldv, void* cwork, blas_int* lwork,
    double* rwork, blas_int* lrwork, blas_int* iwork, blas_int* info);
}

namespace lapack
{

// xBDSQR
inline void bdsqr(char* uplo, blas_int* n, blas_int* ncvt, blas_int* nru,
                  blas_int* ncc, float* d, float* e, float* vt, blas_int* ldvt,
                  float* u, blas_int* ldu, float* c, blas_int* ldc, float* work,
                  blas_int* info)
{
    arma_fortran_noprefix(arma_sbdsqr)(uplo, n, ncvt, nru, ncc, d, e, vt, ldvt,
                                       u, ldu, c, ldc, work, info);
}

inline void bdsqr(char* uplo, blas_int* n, blas_int* ncvt, blas_int* nru,
                  blas_int* ncc, double* d, double* e, double* vt,
                  blas_int* ldvt, double* u, blas_int* ldu, double* c,
                  blas_int* ldc, double* work, blas_int* info)
{
    arma_fortran_noprefix(arma_dbdsqr)(uplo, n, ncvt, nru, ncc, d, e, vt, ldvt,
                                       u, ldu, c, ldc, work, info);
}

// xGGLSE
inline void gglse(blas_int* m, blas_int* n, blas_int* p, float* a,
                  blas_int* lda, float* b, blas_int* ldb, float* c, float* d,
                  float* x, float* work, blas_int* lwork, blas_int* info)
{
    arma_fortran_noprefix(arma_sgglse)(m, n, p, a, lda, b, ldb, c, d, x, work,
                                       lwork, info);
}

inline void gglse(blas_int* m, blas_int* n, blas_int* p, double* a,
                  blas_int* lda, double* b, blas_int* ldb, double* c, double* d,
                  double* x, double* work, blas_int* lwork, blas_int* info)
{
    arma_fortran_noprefix(arma_dgglse)(m, n, p, a, lda, b, ldb, c, d, x, work,
                                       lwork, info);
}

inline void gglse(blas_int* m, blas_int* n, blas_int* p, std::complex<float>* a,
                  blas_int* lda, std::complex<float>* b, blas_int* ldb,
                  std::complex<float>* c, std::complex<float>* d,
                  std::complex<float>* x, std::complex<float>* work,
                  blas_int* lwork, blas_int* info)
{
    arma_fortran_noprefix(arma_cgglse)(m, n, p, a, lda, b, ldb, c, d, x, work,
                                       lwork, info);
}

inline void gglse(blas_int* m, blas_int* n, blas_int* p,
                  std::complex<double>* a, blas_int* lda,
                  std::complex<double>* b, blas_int* ldb,
                  std::complex<double>* c, std::complex<double>* d,
                  std::complex<double>* x, std::complex<double>* work,
                  blas_int* lwork, blas_int* info)
{
    arma_fortran_noprefix(arma_zgglse)(m, n, p, a, lda, b, ldb, c, d, x, work,
                                       lwork, info);
}

// xGEQP3
inline void geqp3(blas_int* m, blas_int* n, float* a, blas_int* lda,
                  blas_int* jpiv, float* tau, float* work, blas_int* lwork,
                  blas_int* info)
{
    arma_fortran_noprefix(sgeqp3)(m, n, a, lda, jpiv, tau, work, lwork, info);
}

inline void geqp3(blas_int* m, blas_int* n, double* a, blas_int* lda,
                  blas_int* jpiv, double* tau, double* work, blas_int* lwork,
                  blas_int* info)
{
    arma_fortran_noprefix(dgeqp3)(m, n, a, lda, jpiv, tau, work, lwork, info);
}

inline void geqp3(blas_int* m, blas_int* n, std::complex<float>* a,
                  blas_int* lda, blas_int* jpiv, std::complex<float>* tau,
                  std::complex<float>* work, blas_int* lwork, float* rwork,
                  blas_int* info)
{
    arma_fortran_noprefix(cgeqp3)(m, n, a, lda, jpiv, tau, work, lwork, rwork,
                                  info);
}

inline void geqp3(blas_int* m, blas_int* n, std::complex<double>* a,
                  blas_int* lda, blas_int* jpiv, std::complex<double>* tau,
                  std::complex<double>* work, blas_int* lwork, double* rwork,
                  blas_int* info)
{
    arma_fortran_noprefix(zgeqp3)(m, n, a, lda, jpiv, tau, work, lwork, rwork,
                                  info);
}

// xHSEQR
inline void hseqr(char* job, char* compz, blas_int* n, blas_int* ilo,
                  blas_int* ihi, float* h, blas_int* ldh, float* wr, float* wi,
                  float* z, blas_int* ldz, float* work, blas_int* lwork,
                  blas_int* info)
{
    arma_fortran_noprefix(arma_shseqr)(job, compz, n, ilo, ihi, h, ldh, wr, wi,
                                       z, ldz, work, lwork, info);
}

inline void hseqr(char* job, char* compz, blas_int* n, blas_int* ilo,
                  blas_int* ihi, double* h, blas_int* ldh, double* wr,
                  double* wi, double* z, blas_int* ldz, double* work,
                  blas_int* lwork, blas_int* info)
{
    arma_fortran_noprefix(arma_dhseqr)(job, compz, n, ilo, ihi, h, ldh, wr, wi,
                                       z, ldz, work, lwork, info);
}

inline void hseqr(char* job, char* compz, blas_int* n, blas_int* ilo,
                  blas_int* ihi, std::complex<float>* h, blas_int* ldh,
                  std::complex<float>* w, std::complex<float>* z, blas_int* ldz,
                  std::complex<float>* work, blas_int* lwork, blas_int* info)
{
    arma_fortran_noprefix(arma_chseqr)(job, compz, n, ilo, ihi, h, ldh, w, z,
                                       ldz, work, lwork, info);
}

inline void hseqr(char* job, char* compz, blas_int* n, blas_int* ilo,
                  blas_int* ihi, std::complex<double>* h, blas_int* ldh,
                  std::complex<double>* w, std::complex<double>* z,
                  blas_int* ldz, std::complex<double>* work, blas_int* lwork,
                  blas_int* info)
{
    arma_fortran_noprefix(arma_zhseqr)(job, compz, n, ilo, ihi, h, ldh, w, z,
                                       ldz, work, lwork, info);
}

// Jacobi SVD
inline void gesvj(char* joba, char* jobu, char* jobv, blas_int* m, blas_int* n,
                  float* a, blas_int* lda, float* sva, blas_int* mv, float* v,
                  blas_int* ldv, float* work, blas_int* lwork, blas_int* info)
{
    arma_fortran_noprefix(arma_sgesvj)(joba, jobu, jobv, m, n, a, lda, sva, mv,
                                       v, ldv, work, lwork, info);
}

inline void gesvj(char* joba, char* jobu, char* jobv, blas_int* m, blas_int* n,
                  double* a, blas_int* lda, double* sva, blas_int* mv,
                  double* v, blas_int* ldv, double* work, blas_int* lwork,
                  blas_int* info)
{
    arma_fortran_noprefix(arma_dgesvj)(joba, jobu, jobv, m, n, a, lda, sva, mv,
                                       v, ldv, work, lwork, info);
}

inline void gesvj(char* joba, char* jobu, char* jobv, blas_int* m, blas_int* n,
                  std::complex<float>* a, blas_int* lda, float* sva,
                  blas_int* mv, std::complex<float>* v, blas_int* ldv,
                  std::complex<float>* cwork, blas_int* lwork, float* rwork,
                  blas_int* lrwork, blas_int* info)
{
    using complex_t = std::complex<float>;
    arma_fortran_noprefix(arma_cgesvj)(
        joba, jobu, jobv, m, n, (complex_t*)a, lda, sva, mv, (complex_t*)v, ldv,
        (complex_t*)cwork, lwork, rwork, lrwork, info);
}

inline void gesvj(char* joba, char* jobu, char* jobv, blas_int* m, blas_int* n,
                  std::complex<double>* a, blas_int* lda, double* sva,
                  blas_int* mv, std::complex<double>* v, blas_int* ldv,
                  std::complex<double>* cwork, blas_int* lwork, double* rwork,
                  blas_int* lrwork, blas_int* info)
{
    using complex_t = std::complex<float>;
    arma_fortran_noprefix(arma_zgesvj)(
        joba, jobu, jobv, m, n, (complex_t*)a, lda, sva, mv, (complex_t*)v, ldv,
        (complex_t*)cwork, lwork, rwork, lrwork, info);
}

// Another Jacobi SVD
inline void gejsv(char* joba, char* jobu, char* jobv, char* jobr, char* jobt,
                  char* jobp, blas_int* m, blas_int* n, float* a, blas_int* lda,
                  float* sva, float* u, blas_int* ldu, float* v, blas_int* ldv,
                  float* work, blas_int* lwork, blas_int* iwork, blas_int* info)
{
    arma_fortran_noprefix(arma_sgejsv)(joba, jobu, jobv, jobr, jobt, jobp, m, n,
                                       a, lda, sva, u, ldu, v, ldv, work, lwork,
                                       iwork, info);
}

inline void gejsv(char* joba, char* jobu, char* jobv, char* jobr, char* jobt,
                  char* jobp, blas_int* m, blas_int* n, double* a,
                  blas_int* lda, double* sva, double* u, blas_int* ldu,
                  double* v, blas_int* ldv, double* work, blas_int* lwork,
                  blas_int* iwork, blas_int* info)
{
    arma_fortran_noprefix(arma_dgejsv)(joba, jobu, jobv, jobr, jobt, jobp, m, n,
                                       a, lda, sva, u, ldu, v, ldv, work, lwork,
                                       iwork, info);
}

inline void gejsv(char* joba, char* jobu, char* jobv, char* jobr, char* jobt,
                  char* jobp, blas_int* m, blas_int* n, std::complex<float>* a,
                  blas_int* lda, float* sva, std::complex<float>* u,
                  blas_int* ldu, std::complex<float>* v, blas_int* ldv,
                  std::complex<float>* cwork, blas_int* lwork, float* rwork,
                  blas_int* lrwork, blas_int* iwork, blas_int* info)
{
    arma_fortran_noprefix(arma_cgejsv)(joba, jobu, jobv, jobr, jobt, jobp, m, n,
                                       a, lda, sva, u, ldu, v, ldv, cwork,
                                       lwork, rwork, lrwork, iwork, info);
}

inline void gejsv(char* joba, char* jobu, char* jobv, char* jobr, char* jobt,
                  char* jobp, blas_int* m, blas_int* n, std::complex<double>* a,
                  blas_int* lda, double* sva, std::complex<double>* u,
                  blas_int* ldu, std::complex<double>* v, blas_int* ldv,
                  std::complex<double>* cwork, blas_int* lwork, double* rwork,
                  blas_int* lrwork, blas_int* iwork, blas_int* info)
{
    arma_fortran_noprefix(arma_zgejsv)(joba, jobu, jobv, jobr, jobt, jobp, m, n,
                                       a, lda, sva, u, ldu, v, ldv, cwork,
                                       lwork, rwork, lrwork, iwork, info);
}

} // namespace: lapack
} // namespace: arma

#endif /* ARMA_LAPACK_EXTRA_HPP */


#include <iostream>
#include <vector>
#include <random>
#include <fstream>
#include <omp.h>
#include <boost/numeric/odeint.hpp>
#include "cahnhilliard_thermal.h"
#include "utils_ch.h"

#include <petscvec.h>
#include <petscksp.h>
#include <petscsnes.h>
#include <petscts.h>

  /*
  Cahn-Hilliard:

  dc/dt = laplacian( u*c^3 - b*c ) - eps_2*biharm(c) - sigma*(c - m) + sigma_noise * N(0,1^2)

  expanding out RHS into individual differentials:
  D*laplacian( u*c^3 - b*c) - D*eps_2*biharm(c)
  assuming constant eps_2.

  need a d^4 and a d^2 operator.
  */

CahnHilliard2DRHS_thermal::CahnHilliard2DRHS_thermal(CHparamsScalar& chp , SimInfo& info)
  : noise_dist_(0.0,1.0) , info_(info) , petsc_context_(*this)
  {
    chpV_.eps_2    = std::vector<double>( info_.nx*info_.ny , chp.eps_2     );
    chpV_.b        = std::vector<double>( info_.nx*info_.ny , chp.b         );
    chpV_.u        = std::vector<double>( info_.nx*info_.ny , chp.u         );
    chpV_.sigma    = std::vector<double>( info_.nx*info_.ny , chp.sigma     );
    chpV_.m        = std::vector<double>( info_.nx*info_.ny , chp.m  );
    chpV_.DT       = std::vector<double>( info_.nx*info_.ny , chp.DT  );
    chpV_.f_T      = std::vector<double>( info_.nx*info_.ny , chp.f_T  );
    chpV_.sigma_noise    = chp.sigma_noise;

    if ( info.bc.compare("dirichlet") == 0) {
      ch_rhs_ = &compute_ch_nonlocal_stationary_boundaries;
      std::cout << "Initialized Cahn-Hilliard equation: scalar parameters, dirichlet BCs, thermal coefficient dependence, thermal diffusion" << std::endl;
    }
    else if ( info.bc.compare("neumann") == 0) {
      ch_rhs_ = &compute_ch_nonlocal_neumannBC;
      std::cout << "Initialized Cahn-Hilliard equation: scalar parameters, neumann BCs, thermal coefficient dependence, thermal diffusion" << std::endl;
    }
    else {
      ch_rhs_ = &compute_ch_nonlocal;
      std::cout << "Initialized Cahn-Hilliard equation: scalar parameters, periodic BCs, thermal coefficient dependence, thermal diffusion" << std::endl;
    }

  }

CahnHilliard2DRHS_thermal::CahnHilliard2DRHS_thermal(CHparamsVector& chp , SimInfo& info)
  : noise_dist_(0.0,1.0) , chpV_(chp) , info_(info) , petsc_context_(*this)
  {

    if ( info.bc.compare("dirichlet") == 0) {
      ch_rhs_ = &compute_ch_nonlocal_stationary_boundaries;
      std::cout << "Initialized Cahn-Hilliard equation: spatial-field parameters, dirichlet BCs, thermal coefficient dependence, thermal diffusion" << std::endl;
    }
    if ( info.bc.compare("neumann") == 0) {
      ch_rhs_ = &compute_ch_nonlocal_neumannBC;
      std::cout << "Initialized Cahn-Hilliard equation: spatial-field parameters, neumann BCs, thermal coefficient dependence, thermal diffusion" << std::endl;
    }
    else {
      ch_rhs_ = &compute_ch_nonlocal;
      std::cout << "Initialized Cahn-Hilliard equation: spatial-field parameters, periodic BCs, thermal coefficient dependence, thermal diffusion" << std::endl;
    }

  }

CahnHilliard2DRHS_thermal::~CahnHilliard2DRHS_thermal() { };

void CahnHilliard2DRHS_thermal::rhs(const std::vector<double> &ct, std::vector<double> &dcTdt, const double t)
  {
    dcTdt.resize(2 * info_.nx * info_.ny);
    std::vector<double> c = std::vector<double>( ct.begin() , ct.begin() + info_.nx*info_.ny );
    std::vector<double> T = std::vector<double>( ct.begin() + info_.nx*info_.ny , ct.end() );

    // enforce thermal BC: dT/dnormal = 0
    # pragma omp parallel for
    for (int i = 0; i < info_.nx; ++i) {
      T[info_.idx2d(0,i)]          = T[info_.idx2d(1,i)];
      T[info_.idx2d(info_.ny-1,i)] = T[info_.idx2d(info_.ny-2,i)];
    }

    # pragma omp parallel for
    for (int i = 0; i < info_.ny; ++i) {
      T[info_.idx2d(i,0)]          = T[info_.idx2d(i,1)];
      T[info_.idx2d(i,info_.nx-1)] = T[info_.idx2d(i,info_.nx-2)];
    }

    // evaluate CH parameter dependencies on temperature
    //chpV_ = compute_chparams_using_temperature( chpV_ , info_ , T );
    chpV_ = compute_eps2_and_sigma_from_polymer_params( chpV_ , info_ , T );

    // evaluate deterministic nonlocal dynamics
    compute_ch_nonlocal(c, dcTdt, t, chpV_, info_);

    // evaluate thermal diffusion
    # pragma omp parallel for
    for (int i = 0; i < info_.ny; ++i) {
      for (int j = 0; j < info_.nx; ++j) {

        const double T_i   = T[info_.idx2d(i, j)];
        const double T_im1 = T[info_.idx2d(i - 1, j)];
        const double T_ip1 = T[info_.idx2d(i + 1, j)];
        const double T_jm1 = T[info_.idx2d(i, j - 1)];
        const double T_jp1 = T[info_.idx2d(i, j + 1)];

        double dxx = 1.0 / (info_.dx * info_.dx) * (T_jm1 + T_jp1 - 2.0 * T_i);
        double dyy = 1.0 / (info_.dy * info_.dy) * (T_im1 + T_ip1 - 2.0 * T_i);

        dcTdt[info_.idx2d(i, j) + info_.nx*info_.ny]  = chpV_.DT[info_.idx2d(i, j)] * (dxx + dyy) + chpV_.f_T[info_.idx2d(i, j)];

      }
    }

    // enforce thermal BC: dT/dnormal = 0
    # pragma omp parallel for
    for (int i = 0; i < info_.nx; ++i) {
      dcTdt[info_.idx2d(0,i) + info_.nx*info_.ny]          = 0;
      dcTdt[info_.idx2d(info_.ny-1,i) + info_.nx*info_.ny] = 0;
    }

    # pragma omp parallel for
    for (int i = 0; i < info_.ny; ++i) {
      dcTdt[info_.idx2d(i,0) + info_.nx*info_.ny]          = 0;
      dcTdt[info_.idx2d(i,info_.nx-1) + info_.nx*info_.ny] = 0;
    }

  }

void CahnHilliard2DRHS_thermal::setInitialConditions(std::vector<double> &x)
  {
    x.resize(2 * info_.nx * info_.ny);

    std::default_random_engine generator;
    std::uniform_real_distribution<double> distribution(-1.0,1.0);

    for (int i = 0; i < info_.ny; ++i) {
      for (int j = 0; j < info_.nx; ++j) {
        x[info_.idx2d(i,j)]                     = distribution(generator) * 0.005;
	x[info_.idx2d(i,j) + info_.nx*info_.ny] = chpV_.T_min;
      }
    }

    // Set BCs if needed
    if ( info_.bc.compare("dirichlet") == 0) {
      x = apply_dirichlet_bc( x , info_ );
    }
    else if ( info_.bc.compare("neumann") == 0 ) {
      x = apply_neumann_bc( x , info_ );
    }

  }

void CahnHilliard2DRHS_thermal::write_state(const std::vector<double> &x , const int idx , const int nx , const int ny , std::string& outdir)
{
  if ( outdir.back() != '/' )
    outdir += '/';
  std::ofstream outC;
  std::ofstream outT;
  outC.open( outdir + "C_" + std::to_string(idx) + ".out" );
  outT.open( outdir + "T_" + std::to_string(idx) + ".out" );
  outC.precision(16);
  outT.precision(16);

  for (int i = 0; i < ny; ++i){
    for (int j = 0; j < nx; ++j){
      outC << x[i * ny + j] << " ";
      outT << x[i * ny + j + nx*ny] << " ";
    }
  }

  outC.close();
  outT.close();
};