/* Copyright 2017 Battelle Energy Alliance, LLC

   Licensed under the Apache License, Version 2.0 (the "License");
   you may not use this file except in compliance with the License.
   You may obtain a copy of the License at

   http://www.apache.org/licenses/LICENSE-2.0

   Unless required by applicable law or agreed to in writing, software
   distributed under the License is distributed on an "AS IS" BASIS,
   WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
   See the License for the specific language governing permissions and
   limitations under the License.
*/
/*
 *
 *  Created on: April 28, 2012
 *      Author: MANDD
 *
 *      Tests  : None for the custom
 *
 *      Problems : None
 *      Issues  : None
 *      Complaints : None
 *      Compliments : None
 *
 *      source: Numerical Recipes in C++ 3rd edition
 *
 */

#include <sstream>
#include <fstream>
#include <ctime>
#include <cstdlib>
#include <vector>
#include <iostream>
#include <string>
#include <iostream>
#include <cmath> // to use erfc error function
#include <ctime> // for rand() and srand()
#include <cstdio>

#include "distribution_1D.h"
#include "distributionFunctions.h"

#include <boost/numeric/ublas/matrix.hpp>
#include <boost/numeric/ublas/lu.hpp>

//#include <Eigen/Dense>
#include <Eigen/SVD>
#define throwError(msg) { std::cerr << "\n\n" << msg << "\n\n"; throw std::runtime_error("Error"); }

#define _USE_MATH_DEFINES

/*
extern "C" {
    // LU decomoposition of a general matrix
    void dgetrf_(int* M, int *N, double* A, int* lda, int* IPIV, int* INFO);

    // generate inverse of a matrix given its LU decomposition
    void dgetri_(int* N, double* A, int* lda, int* IPIV, double* WORK, int* lwork, int* INFO);
}
*/

typedef boost::numeric::ublas::matrix<double> matrixDouble;

void matrixConversionBoost(const std::vector<std::vector<double> > & original, matrixDouble & converted)
{
  converted.resize(original.size(),original.at(0).size());
  for(unsigned int r=0; r < original.size(); r++)
  {
    for(unsigned int c=0; c < original.at(r).size(); c++)
    {
      converted(r,c)  = original.at(r).at(c);
    }
  }
}

/*
void matrixConversion(std::vector<std::vector<double> > original, double converted[]){
        if (original.size() == original[0].size()){
                int dimensions = original.size();
                for (int r=0; r<dimensions; r++)
                        for (int c=0; c<dimensions; c++){
                                converted[r*dimensions+c] = original[r][c];
                        }
        }else
                throwError("Error in matrixConversion: matrix is not squared.");
}
*/

void matrixBackConversionBoost(const matrixDouble & original, std::vector<std::vector<double> > & converted)
{
  for(unsigned int r=0; r < original.size1(); r++)
  {
    for(unsigned int c = 0; c < original.size2(); c++)
    {
      converted.at(r).at(c) = original(r,c);
    }
  }
}

/*
void matrixBackConversion(double original[], std::vector<std::vector<double> > converted){
        int dimensions = int(sizeof(original)/sizeof(double));
        dimensions = sqrt(dimensions);

        for (int r=0; r<dimensions; r++)
                for (int c=0; c<dimensions; c++)
                        converted[r][c] = original[r*dimensions+c];
}
*/

/*
//http://stackoverflow.com/questions/3519959/computing-the-inverse-of-a-matrix-using-lapack-in-c
void inverseMatrix(double* a, int n)
{
    int *IPIV = new int[n+1];
    int LWORK = n*n;
    double *WORK = new double[LWORK];
    int INFO;

    dgetrf_(&n,&n,a,&n,IPIV,&INFO);
    dgetri_(&n,a,&n,IPIV,WORK,&LWORK,&INFO);

    delete IPIV;
    delete WORK;
}
*/

//Roughly based on http://savingyoutime.wordpress.com/2009/09/21/c-matrix-inversion-boostublas/ and libs/numeric/ublas/test/test_lu.cpp
void invertMatrixBoost(matrixDouble & a, matrixDouble & aInverted)
{
  boost::numeric::ublas::permutation_matrix<std::size_t> pm(a.size1());

  boost::numeric::ublas::lu_factorize<matrixDouble,
                                      boost::numeric::ublas::permutation_matrix<> >(a, pm);

  aInverted.assign(boost::numeric::ublas::identity_matrix<double>(a.size1()));

  boost::numeric::ublas::lu_substitute(a, pm, aInverted);
}

void computeInverse(const std::vector<std::vector<double> > & matrix, std::vector<std::vector<double> > & inverse){
        int dimensions = matrix.size();
        matrixDouble A(dimensions,dimensions),inverted(dimensions,dimensions);
        matrixConversionBoost(matrix, A);
        invertMatrixBoost(A,inverted);
        matrixBackConversionBoost(inverted, inverse);
}

// Convert the vector of covariance to vector of vector of covariance
void  vectorToMatrix(unsigned int &rows,unsigned int &columns,std::vector<double> &vec_matrix, std::vector<std::vector<double> > &cov_matrix) {
        /** Input Parameter
         * vec_matrix: covariance matrix stored in a vector
         * Output Parameter
         * rows: the first dimension of the covariance matrix
         * columns: the second dimension of the covariance matrix
         * cov_matrix: covariance matrix stored in vector<vector<double> >
         */
        unsigned int dimensions = vec_matrix.size();
        dimensions = std::lround(std::sqrt(dimensions));
        rows = dimensions;
        columns = dimensions;
        if(rows*columns != vec_matrix.size())
                      throwError("MultivariateNormal error: covariance matrix in is not a square matrix.");
        for (unsigned int row = 0; row < rows; ++row) {
                std::vector<double> temp;
                for (unsigned int colm = 0; colm < columns; ++colm) {
                        temp.push_back(vec_matrix.at(colm+row*columns));
                }
                cov_matrix.push_back(temp);
        }
}

//See for example http://en.wikipedia.org/wiki/LU_decomposition or
// http://programmingexamples.net/wiki/CPP/Boost/Math/uBLAS/determinant
double getDeterminantBoost(matrixDouble & a)
{
  boost::numeric::ublas::permutation_matrix<std::size_t> pm(a.size1());
  double determinant = 1.0;

  int result = boost::numeric::ublas::lu_factorize<matrixDouble,boost::numeric::ublas::permutation_matrix<> >(a, pm);

  if(result)
  {
    return 0.0;
  }
  //det(a) = det(P^-1)det(L)det(U) = (-1)^S*(product(u_ii))
  //Where S is the number of row exchanges
  for(unsigned int i = 0; i < a.size1(); i++)
  {
    //Multiple by current diagonal entry
    determinant *= a(i,i);
    if(i != pm(i))
    {
      //Found a row exchange
      determinant *= -1;
    }

  }
  return determinant;
}

double getDeterminant(std::vector<std::vector<double> > matrix){
        int dimensions = matrix.size();
        matrixDouble A(dimensions,dimensions);

        matrixConversionBoost(matrix, A);

        return getDeterminantBoost(A);

}

/*
double getDeterminant(std::vector<std::vector<double> > matrix){
        int dimensions = matrix.size();
        double A [dimensions*dimensions];

        matrixConversion(matrix, A);

    int *IPIV = new int[dimensions+1];
    int LWORK = dimensions*dimensions;
    double *WORK = new double[LWORK];
    int INFO;

    dgetrf_(&dimensions,&dimensions,A,&dimensions,IPIV,&INFO);

    double determinant =1;

    for(int index=0; index<dimensions; index++)
        determinant *= A[index*dimensions];

    delete IPIV;
    delete WORK;

        return determinant;
}
*/

void svdDecomposition(const std::vector<std::vector<double> > &matrix, std::vector<std::vector<double> > &left_singular_vectors, std::vector<std::vector<double> > &right_singular_vectors, std::vector<double> &singular_values, std::vector<std::vector<double> > &transformed_matrix) {
  /**
   * This function compute the singular value decomposition for given matrix
   * Input Parameters
   * matrix: provided data
   * Output Parameters
   * left_singular_vectors: stores the left singular vectors for given matrix
   * right_singular_vectors: stores the right singular vectors for given matrix
   * singular_values: stores the singular values for given matrix
   * transformed_matrix: stores the transformation matrix
   */
  unsigned int row = matrix.size();
  unsigned int col = matrix.at(0).size();
  unsigned int dim = 0;
  if(row > col) {
    dim = col;
  }else {
    dim = row;
  }
  Eigen::MatrixXd A(row,col);
  Eigen::MatrixXd U(row,row);
  Eigen::MatrixXd V(col,col);
  Eigen::MatrixXd X(row,dim);
  Eigen::VectorXd S(dim);
  matrixConversionToEigenType(matrix,A);
  Eigen::JacobiSVD<Eigen::MatrixXd> svd(A,Eigen::ComputeFullU | Eigen::ComputeFullV);
  U = svd.matrixU();
  V = svd.matrixV();
  S = svd.singularValues();
  for(unsigned int i = 0; i < dim; ++i) {
    X.col(i) = U.col(i)*sqrt(S(i));
  }
  matrixConversionToCxxVVectorType(U,left_singular_vectors);
  matrixConversionToCxxVVectorType(V,right_singular_vectors);
  vectorConversionToCxxVectorType(S,singular_values);
  matrixConversionToCxxVVectorType(X,transformed_matrix);
}

void getInverseTransformedMatrix(const std::vector<std::vector<double> > &left_singular_vectors, std::vector<double> &singular_values, std::vector<std::vector<double> > &inverse_transformed_matrix) {
  /**
   * This function compute the inverse transformation matrix
   * Input Parameters
   * left_singular_vectors: stores the left singular vectors for given matrix
   * singular_values: stores the singular values for given matrix
   * Output Parameters
   * inverse_transformed_matrix: stores the inverse transformation matrix
   */
  unsigned int row = left_singular_vectors.size();
  unsigned int col = left_singular_vectors.at(0).size();
  unsigned int dim = singular_values.size();
  Eigen::MatrixXd U(row,col);
  Eigen::MatrixXd inverseX(row,dim);
  matrixConversionToEigenType(left_singular_vectors,U);
  for (unsigned int i = 0; i < dim; ++i) {
    if (singular_values.at(i) == 0) {
      inverseX.col(i) = U.col(i) * 0.0;
    } else {
      inverseX.col(i) = U.col(i) * (1.0/sqrt(singular_values.at(i)));
    }
  }
  matrixConversionToCxxVVectorType(inverseX.transpose(),inverse_transformed_matrix);
}

void svdDecomposition(const std::vector<std::vector<double> > &matrix, std::vector<std::vector<double> > &left_singular_vectors, std::vector<std::vector<double> > &right_singular_vectors, std::vector<double> &singular_values, std::vector<std::vector<double> > &transformed_matrix, unsigned int rank) {
  /**
   * This function compute the singular value decomposition for given matrix
   * Input Parameters
   * matrix: provided data
   * Output Parameters
   * left_singular_vectors: stores the left singular vectors for given matrix
   * right_singular_vectors: stores the right singular vectors for given matrix
   * singular_values: stores the singular values for given matrix
   * rank: used for truncated svd, the number of singular values that will be kept for truncated svd
   */
  unsigned int row = matrix.size();
  unsigned int col = matrix.at(0).size();
  unsigned int dim = 0;
  if(row > col) {
    dim = col;
  }else {
    dim = row;
  }
  Eigen::MatrixXd A(row,col);
  Eigen::MatrixXd U(row,row);
  Eigen::MatrixXd V(col,col);
  Eigen::MatrixXd X(row,rank);
  Eigen::VectorXd S(dim);
  matrixConversionToEigenType(matrix,A);
  Eigen::JacobiSVD<Eigen::MatrixXd> svd(A,Eigen::ComputeFullU | Eigen::ComputeFullV);
  U = svd.matrixU();
  V = svd.matrixV();
  S = svd.singularValues();
  for(unsigned int i = 0; i < rank; ++i) {
    X.col(i) = U.col(i)*sqrt(S(i));
  }
  // transform and store the matrix for the truncated svd
  matrixConversionToCxxVVectorType(U.block(0,0,row,rank),left_singular_vectors);
  matrixConversionToCxxVVectorType(V.block(0,0,col,rank),right_singular_vectors);
  vectorConversionToCxxVectorType(S.head(rank),singular_values);
  matrixConversionToCxxVVectorType(X,transformed_matrix);
}

void  computeNearestSymmetricMatrix(const std::vector<std::vector<double> > &matrix, std::vector<std::vector<double> > &symmetric_matrix) {
  /**
   * This function used to compute the nearest symmetric matrix
   * Input Parameters
   * matrix: std::vector<std::vector<double> >, given matrix
   * Output Parameters
   * symmetric_matrix: std::vector<std::vector<double> >, the computed symmetric matrix
   */
  unsigned int row = matrix.size();
  unsigned int col = matrix.at(0).size();
  if (row != col) {
    throwError("The provided matrix is not a square matrix!" );
  }
  Eigen::MatrixXd A(row,col);
  Eigen::MatrixXd B(col,row);
  matrixConversionToEigenType(matrix,A);
  B = A.transpose();
  A = (A + B)*0.5;
  matrixConversionToCxxVVectorType(A,symmetric_matrix);
}

void resetSingularValues(std::vector<std::vector<double> > &left_singular_vectors, std::vector<std::vector<double> > &right_singular_vectors, std::vector<double> &singular_values,std::vector<std::vector<double> > &transformed_matrix) {
  /**
   * used to reset singular values
   * Input Parameters
   * left_singular_vectors: std::vector<std::vector<double> >, the left singular vectors
   * right_singular_vectors: std::vector<std::vector<double> >, the right singular vectors
   * singular_values: std::vector<double>, the singular values
   * transformed_matrix: std::vector<vector<double> >, the transformation matrix
   * Output Parameters
   * singular_values: std::vector<double>, the modified singular values
   * transformed_matrix: std::vector<vector<double> >, the updated transformation matrix
   */
  unsigned int row1 = left_singular_vectors.size();
  unsigned int col1 = left_singular_vectors.at(0).size();
  unsigned int row2 = right_singular_vectors.size();
  unsigned int col2 = right_singular_vectors.at(0).size();
  unsigned int rank = transformed_matrix.at(0).size();
  if (row1 != row2 && col1 != col2) {
    throwError("The provided matrices should have the same shape!" );
  }
  Eigen::MatrixXd A(row1,col1);
  Eigen::MatrixXd B(row2,col2);
  Eigen::MatrixXd X(row1,rank);
  matrixConversionToEigenType(left_singular_vectors,A);
  matrixConversionToEigenType(right_singular_vectors,B);
  double tol;
  tol = 1.0E-15;
  if (singular_values.at(0) == 0.0) {
    throwError("The provided covariance matrix is zero matrix!")
  }
  for (unsigned int i = 0; i < singular_values.size(); ++i) {
    if ((A.col(i) + B.col(i)).lpNorm<1>()/row1 < tol) {
      singular_values.at(i) = 0.0;
    }
    if (singular_values.at(i)/singular_values.at(0) < tol) {
      singular_values.at(i) = 0.0;
    }
  }
  for(unsigned int i = 0; i < rank; ++i) {
    X.col(i) = A.col(i)*sqrt(singular_values.at(i));
  }
  transformed_matrix.clear();
  matrixConversionToCxxVVectorType(X,transformed_matrix);
}

void matrixConversionToEigenType(std::vector<std::vector<double> > original, Eigen::MatrixXd &converted) {
  /**
   * This function convert the data from type std::vector<std::vector<double> > to Eigen::MatrixXd
   * Input Parameters
   * original: provided data with the type of std::vector<std::vector<double> >
   * Output Parameters
   * converted: output data with the type of Eigen::MatrixXd
   */
  converted.resize(original.size(),original.at(0).size());
  for(unsigned int row = 0; row < original.size(); ++row) {
    for(unsigned int col = 0; col < original.at(0).size(); ++col) {
      if(original.at(0).size() != original.at(row).size()) {
        throwError("The matrix stored in the C++ vector container with different lenght of columns");
      }
      converted(row,col) = original.at(row).at(col);
    }
  }
}

void matrixConversionToCxxVVectorType(const Eigen::MatrixXd & original, std::vector<std::vector<double> > &converted) {
  /**
   * This function convert the data from type Eigen::MatrixXd to type std::vector<double>
   * Input Parameters
   * original: provided data with the type of Eigen::MatrixXd
   * Output Parameters
   * converted: output data with the type of std::vector<std::vector<double> >
   */
  for(unsigned int row = 0; row < original.rows(); ++row) {
    std::vector<double> temp;
    for(unsigned int col = 0; col < original.cols(); ++col) {
      temp.push_back(original(row,col));
    }
    converted.push_back(temp);
  }
}

void vectorConversionToCxxVectorType(const Eigen::VectorXd & original, std::vector<double> &converted) {
  /**
   * This function convert the data from type Eigen::VectorXd to type std::vector<double>
   * Input Parameters
   * original: provided data with the type of Eigen::VectorXd
   * Output Parameters
   * converted: output data with the type of std::vector<double>
   */
  for(unsigned int dim = 0; dim < original.rows(); ++dim) {
    converted.push_back(original(dim));
  }
}

// void nrerror(const char error_text[]){     // added const to avoid "warning: deprecated conversion from string constant to *char
// /* Numerical Recipes standard error handler */
//  fprintf(stderr,"Numerical Recipes run-time error...\n");
//  fprintf(stderr,"%s\n",error_text);
//  fprintf(stderr,"...now exiting to system...\n");
// }
//
// double gammp(double a, double x){
// /* high level function for incomplete gamma function */
//         void gcf(double *gammcf,double a,double x,double *gln);
//         void gser(double *gamser,double a,double x,double *gln);
//         double gamser,gammcf,gln;
//         if(x < 0.0 || a <= 0.0) nrerror("Invalid arg in gammp");
//         if(x < (a+1.0)){
// /* here I change routine so that it returns \gamma(a,x)
//         or P(a,x)-just take out comments to get P(a,x) vs \gamma(a,x)-
//         to get latter use the exp(log(.)+gln) expression */
//                gser(&gamser,a,x,&gln);
// //      return exp(log(gamser)+gln);
//                return gamser;
//         }
//         else{
//                gcf(&gammcf,a,x,&gln);
// //      return exp(log(1.0-gammcf)+gln);
//                return 1.0-gammcf;
//         }
// }
//
// double loggam(double xx)
// {
//         double x,y,tmp,ser;
//         static double cof[6]={76.18009172947146, -86.50532032941677,
//                24.01409824083091,-1.231739572450155, 0.001208650973866179,
//                -5.395239384953e-006};
//         int j;
//         y=x=xx;
//         tmp=x+5.5;
//         tmp -= (x+0.5)*log(tmp);
//         ser=1.000000000190015;
//         for(j=0;j<=5;j++) ser += cof[j]/++y;
//         return -tmp+log(2.506628274631*ser/x);
// }
//
// #define ITMAX 100
// #define EPSW 3.0e-7
//
// void gser(double *gamser,double a,double x,double *gln){
//         int n;
//         double sum,del,ap;
//         *gln=loggam(a);
//         if(x <= 0.0){
//                if(x < 0.0) nrerror("x less than 0 in routine gser");
//                *gamser=0.0;
//                return;
//         }
//         else{
//                ap=a;
//                del=sum=1.0/a;
//                for(n=1;n<=ITMAX;n++){
//                       ++ap;
//                       del *= x/ap;
//                       sum += del;
//                       if(fabs(del) < fabs(sum)*EPSW){
//    *gamser=sum*exp(-x+a*log(x)-(*gln));
//    return;
//                       }
//                }
//                nrerror("a too large, ITMAX too small in routine gser");
//                return;
//         }
// }
//
//
// #define FPMIN 1.0e-30
//
// void gcf(double *gammcf,double a,double x,double *gln){
//         int i;
//         double an,b,c,d,del,h;
//         *gln=loggam(a);
//         b=x+1.0-a;
//         c=1.0/FPMIN;
//         d=1.0/b;
//         h=d;
//         for(i=1;i<=ITMAX;i++){
//                an = -i*(i-a);
//                b += 2.0;
//                d=an*d+b;
//                if(fabs(d) < FPMIN) d=FPMIN;
//                c=b+an/c;
//                if(fabs(c) < FPMIN) c=FPMIN;
//                d=1.0/d;
//                del=d*c;
//                h *= del;
//                if(fabs(del-1.0) < EPSW) break;
//         }
//         if(i > ITMAX) nrerror("a too large, ITMAX too small in gcf");
//         *gammcf=exp(-x+a*log(x)-(*gln))*h;
// }
//
// // Gamma function
// // source http://www.crbond.com/math.htm
// double gammaFunc(double x){
//  int i,k,m;
//  double ga,gr,r,z;
//
//  static double g[] = {
//   1.0,
//   0.5772156649015329,
//                 -0.6558780715202538,
//                 -0.420026350340952e-1,
//   0.1665386113822915,
//                 -0.421977345555443e-1,
//                 -0.9621971527877e-2,
//   0.7218943246663e-2,
//                 -0.11651675918591e-2,
//                 -0.2152416741149e-3,
//   0.1280502823882e-3,
//                 -0.201348547807e-4,
//                 -0.12504934821e-5,
//   0.1133027232e-5,
//                 -0.2056338417e-6,
//   0.6116095e-8,
//   0.50020075e-8,
//                 -0.11812746e-8,
//   0.1043427e-9,
//   0.77823e-11,
//                 -0.36968e-11,
//   0.51e-12,
//                 -0.206e-13,
//                 -0.54e-14,
//   0.14e-14};
//
//  if (x > 171.0) return 1e308;    // This value is an overflow flag.
//  if (x == (int)x) {
//   if (x > 0.0) {
//    ga = 1.0;               // use factorial
//    for (i=2;i<x;i++) {
//                                 ga *= i;
//    }
//                       }
//                       else
//    ga = 1e308;
//               }
//               else {
//   if (fabs(x) > 1.0) {
//    z = fabs(x);
//    m = (int)z;
//    r = 1.0;
//    for (k=1;k<=m;k++)
//     r *= (z-k);
//    z -= m;
//   }
//   else
//    z = x;
//   gr = g[24];
//   for (k=23;k>=0;k--)
//    gr = gr*z+g[k];
//
//   ga = 1.0/(gr*z);
//   if (fabs(x) > 1.0) {
//    ga *= r;
//    if (x < 0.0)
//     ga = -M_PI/(x*ga*sin(M_PI*x));
//   }
//  }
//  return ga;
// }
//
//
// // Beta function
// double betaFunc(double alpha, double beta){
//  double value=gammaFunc(alpha)*gammaFunc(beta)/gammaFunc(alpha+beta);
//  return value;
// }
//
//
// // log gamma using the Lanczos approximation
// double logGamma(double x) {
// const double c[8] = { 676.5203681218851, -1259.1392167224028,
//                       771.32342877765313, -176.61502916214059,
//                       12.507343278686905, -0.13857109526572012,
//                       9.9843695780195716e-6, 1.5056327351493116e-7 };
// double sum = 0.99999999999980993;
// double y = x;
// for (int j = 0; j < 8; j++)
//        sum += c[j] / ++y;
// return log(sqrt(2*3.14159) * sum / x) - (x + 7.5) + (x + 0.5) * log(x + 7.5);
// }
//
// // helper function for incomplete beta
// // computes continued fraction
//
// double betaContFrac(double a, double b, double x) {
//  const int MAXIT = 1000;
//  const double EPS = 3e-7;
//  double qab = a + b;
//  double qap = a + 1;
//  double qam = a - 1;
//  double c = 1;
//  double d = 1 - qab * x / qap;
//  if (fabs(d) < FPMIN) d = FPMIN;
//  d = 1 / d;
//  double h = d;
//  int m;
//  for (m = 1; m <= MAXIT; m++) {
//                int m2 = 2 * m;
//                double aa = m * (b-m) * x / ((qam + m2) * (a + m2));
//                d = 1 + aa * d;
//                if (fabs(d) < FPMIN) d = FPMIN;
//                c = 1 + aa / c;
//                if (fabs(c) < FPMIN) c = FPMIN;
//                d = 1 / d;
//                h *= (d * c);
//                aa = -(a+m) * (qab+m) * x / ((a+m2) * (qap+m2));
//                d = 1 + aa * d;
//                if (fabs(d) < FPMIN) d = FPMIN;
//                c = 1 + aa / c;
//                if (fabs(c) < FPMIN) c = FPMIN;
//                d = 1 / d;
//                double del = d*c;
//                h *= del;
//                if (fabs(del - 1) < EPS) break;
//  }
//  if (m > MAXIT) {
//                cerr << "betaContFrac: too many iterations\n";
//  }
//  return h;
// }
//
// // incomplete beta function
// // must have 0 <= x <= 1
// double betaInc(double a, double b, double x) {
//        if (x == 0)
//  return 0;
//        else if (x == 1)
//  return 1;
//        else {
//  double logBeta = logGamma(a+b) - logGamma(a) - logGamma(b)
//                + a * log(x) + b * log(1-x);
//  if (x < (a+1) / (a+b+2))
//                return exp(logBeta) * betaContFrac(a, b, x) / a;
//  else
//                return 1 - exp(logBeta) * betaContFrac(b, a, 1-x) / b;
//        }
// }
//
// double normRNG(double mu, double sigma, double RNG) {
//  static bool deviateAvailable=false;                        //        flag
//  static float storedDeviate;                        //        deviate from previous calculation
//  double polar, rsquared, var1, var2;
//  //srand(time(NULL));
//  //srand((unsigned)time(0));
//  //srand(time(0));
//  //Ran ran(time(0));
//  //        If no deviate has been stored, the polar Box-Muller transformation is
//  //        performed, producing two independent normally-distributed random
//  //        deviates.  One is stored for the next round, and one is returned.
//  if (!deviateAvailable) {
//
//   //        choose pairs of uniformly distributed deviates, discarding those
//   //        that don't fall within the unit circle
//   do {
//    var1 = 2.0*( RNG ) - 1.0;
//    var2 = 2.0*( RNG ) - 1.0;
//
//    //var1=2.0*( ran.doub() ) - 1.0;
//    //var2=2.0*( ran.doub() ) - 1.0;
//
//    rsquared=var1*var1+var2*var2;
//   } while ( rsquared>=1.0 || rsquared == 0.0);
//
//   //        calculate polar transformation for each deviate
//   polar=sqrt(-2.0*log(rsquared)/rsquared);
//
//   //        store first deviate and set flag
//   storedDeviate=var1*polar;
//   deviateAvailable=true;
//
//   //        return second deviate
//   return var2*polar*sigma + mu;
//  }
//
//  //        If a deviate is available from a previous call to this function, it is
//  //        returned, and the flag is set to false.
//  else {
//   deviateAvailable=false;
//   return storedDeviate*sigma + mu;
//  }
// }
//
//
// void loadData(double** data, int dimensionality, int cardinality, string filename) {
//        int x, y;
//
//        ifstream in(filename.c_str());
//
//        if (!in) {
//          cout << "Cannot open file.\n";
//          return;
//        }
//
//        for (y = 0; y < cardinality; y++) {
//          for (x = 0; x < dimensionality; x++) {
//            in >> data[y][x];
//          }
//        }
//
//        in.close();
// }
//
// double calculateCustomPdf(double position, double fitting, double** data_set, int number_samples){
//  double value=-1;
//  double min;
//  double max;
//
//  for (int i=1; i<number_samples; i++){
//   max=data_set[i][1];
//   min=data_set[i-1][1];
//
//   if((position>min)&(position<max)){
//    if (fitting==1)
//     value=data_set[i-1][2];
//    else
//     value=data_set[i-1][2]+(data_set[i][2]-data_set[i-1][2])/(data_set[i][1]-data_set[i-1][1])*(position-data_set[i-1][1]);
//   }
//   else
//    perror ("The following error occurred: distribution sampled out of its boundaries");
//  }
//
//  return value;
// }
//
// double calculateCustomCDF(double position, double fitting, double** data_set, int number_samples){
//  double value=-1;
//  double min;
//  double max;
//  double cumulative=0;
//
//  for (int i=1; i<number_samples; i++){
//   max=data_set[i][1];
//   min=data_set[i-1][1];
//
//   if((position>min)&(position<max)){
//    if (fitting==1)
//     value=cumulative+data_set[i-1][2]*(position-data_set[i-1][1]);
//    else{
//     double pdfValueInPosition =data_set[i-1][2]+(data_set[i][2]-data_set[i-1][2])/(data_set[i][1]-data_set[i-1][1])*(position-data_set[i-1][1]);
//     value=cumulative + (pdfValueInPosition+data_set[i-1][2])*(position-data_set[i-1][1])/2;
//    }
//   }
//   else
//    perror ("The following error occurred: distribution sampled out of its boundaries");
//
//   if (fitting==1)
//    cumulative=cumulative+data_set[i-1][2]*(data_set[i][1]-data_set[i-1][1]);
//   else
//    cumulative=cumulative+(data_set[i][2]+data_set[i-1][2])*(data_set[i][1]-data_set[i-1][1])/2;
//  }
//
//  return value;
// }
//
//        double stdGammaRNG(double shape);
//
// double gammaRNG(double shape, double scale){
//  double value=scale * stdGammaRNG(shape);
//  return value;
// }
//
//        double rkGauss();
//
// double stdGammaRNG(double shape)
// {
//          double b, c;
//          double U, V, X, Y;
//
//          if (shape == 1.0)
//          {
//              return -log(1.0 - rand());
//          }
//          else if (shape < 1.0)
//          {
//              for (;;)
//              {
//                  U = rand();
//                  V = -log(1.0 - rand());
//                  if (U <= 1.0 - shape)
//                  {
//                      X = pow(U, 1./shape);
//                      if (X <= V)
//                      {
//                          return X;
//                      }
//                  }
//                  else
//                  {
//                      Y = -log((1-U)/shape);
//                      X = pow(1.0 - shape + shape*Y, 1./shape);
//                      if (X <= (V + Y))
//                      {
//                          return X;
//                      }
//                  }
//              }
//          }
//          else
//          {
//              b = shape - 1./3.;
//              c = 1./sqrt(9*b);
//              for (;;)
//              {
//                  do
//                  {
//                      X = rkGauss();
//                      V = 1.0 + c*X;
//                  } while (V <= 0.0);
//
//                  V = V*V*V;
//                  U = rand();
//                  if (U < 1.0 - 0.0331*(X*X)*(X*X))
//                      return (b*V);
//                  if (log(U) < 0.5*X*X + b*(1. - V + log(V)))
//                      return (b*V);
//              }
//          }
// }
//
// double rkGauss() {
//  double f, x1, x2, r2;
//
//  do {
//   x1 = 2.0*rand() - 1.0;
//   x2 = 2.0*rand() - 1.0;
//   r2 = x1*x1 + x2*x2;
//  }
//  while (r2 >= 1.0 || r2 == 0.0);
//
//  /* Box-Muller transform */
//  f = sqrt(-2.0*log(r2)/r2);
//  return f*x2;
//
// }
//
//
// double betaRNG(double alpha, double beta){
//  // To be updated
//  double value=0;
//  return value;
// }
//
// double modifiedLogFunction(double x){
//          if (x <= -1.0)
//          {
//              std::stringstream os;
//              os << "Invalid input argument (" << x << "); must be greater than -1.0";
//              throw std::invalid_argument( os.str() );
//          }
//
//          if (fabs(x) > 1e-4)
//              return log(1.0 + x);
//          else
//              return (-0.5*x + 1.0)*x;
// }

        double abramStegunApproximation(double t)
        {
            // Abramowitz and Stegun formula

            double c[] = {2.515517, 0.802853, 0.010328};
            double d[] = {1.432788, 0.189269, 0.001308};
            return t - ((c[2]*t + c[1])*t + c[0]) /
                       (((d[2]*t + d[1])*t + d[0])*t + 1.0);
        }


/**
 * \file LMSFilter.cpp
 */

#include <ATK/Adaptive/LMSFilter.h>

#include <complex>
#include <cstdint>
#include <stdexcept>

#include <Eigen/Core>

#include <ATK/Core/TypeTraits.h>
#include <ATK/Core/Utilities.h>

namespace ATK
{
  template<typename DataType_>
  class LMSFilter<DataType_>::LMSFilterImpl
  {
  public:
    using wType = Eigen::Matrix<DataType_, Eigen::Dynamic, 1>;
    using xType = Eigen::Map<const wType>;

    wType w;
    /// Memory factor
    double alpha = 0.99;
    /// line search
    double mu = 0.05;

    explicit LMSFilterImpl(gsl::index size)
    :w(wType::Zero(size))
    {
    }

    using UpdateFunction = void (LMSFilterImpl::*)(const xType& x, DataType error);

    void update(const xType& x, DataType error)
    {
      w = static_cast<DataType>(alpha) * w + static_cast<DataType>(mu) * error * x;
    }

    void update_normalized(const xType& x, DataType error)
    {
      w = static_cast<DataType>(alpha) * w + static_cast<DataType>(mu) * error * x / (std::numeric_limits<DataType>::epsilon() + static_cast<DataType>(x.squaredNorm()));
    }

    void update_signerror(const xType& x, DataType error)
    {
      w = static_cast<DataType>(alpha) * w + static_cast<DataType>(mu) * error / (std::numeric_limits<DataType>::epsilon() + std::abs(error)) * x;
    }

    void update_signdata(const xType& x, DataType error)
    {
      w = static_cast<DataType>(alpha) * w.array() + static_cast<DataType>(mu) * error * x.array() / (x.cwiseAbs().template cast<DataType>().array() + static_cast<DataType>(std::numeric_limits<DataType>::epsilon()));
    }

    void update_signsign(const xType& x, DataType error)
    {
      w = static_cast<DataType>(alpha) * w.array() + static_cast<DataType>(mu) * error / (std::numeric_limits<DataType>::epsilon() + std::abs(error)) * x.array() / (x.cwiseAbs().template cast<DataType>().array() + static_cast<DataType>(std::numeric_limits<DataType>::epsilon()));
    }

    UpdateFunction select(Mode mode)
    {
      switch (mode)
      {
      case Mode::NORMAL:
        return &LMSFilterImpl::update;
      case Mode::NORMALIZED:
        return &LMSFilterImpl::update_normalized;
      case Mode::SIGNERROR:
        return &LMSFilterImpl::update_signerror;
      case Mode::SIGNDATA:
        return &LMSFilterImpl::update_signdata;
      case Mode::SIGNSIGN:
        return &LMSFilterImpl::update_signsign;
      default:
          throw std::range_error("Wrong mode for LMS filter");
      }
    }
  };

  template<typename DataType_>
  LMSFilter<DataType_>::LMSFilter(gsl::index size)
  :Parent(2, 1), impl(std::make_unique<LMSFilterImpl>(size))
  {
    input_delay = size - 1;
  }

  template<typename DataType_>
  LMSFilter<DataType_>::~LMSFilter()
  {
  }

  template<typename DataType_>
  void LMSFilter<DataType_>::set_size(gsl::index size)
  {
    if(size == 0)
    {
      throw RuntimeError("Size must be strictly positive");
    }

    input_delay = size - 1;
    impl = std::make_unique<LMSFilterImpl>(size);
  }

  template<typename DataType_>
  gsl::index LMSFilter<DataType_>::get_size() const
  {
    return input_delay + 1;
  }

  template<typename DataType_>
  void LMSFilter<DataType_>::set_memory(double memory)
  {
    if (memory >= 1)
    {
      throw ATK::RuntimeError("Memory must be less than 1");
    }
    if (memory <= 0)
    {
      throw ATK::RuntimeError("Memory must be strictly positive");
    }

    impl->alpha = memory;
  }

  template<typename DataType_>
  double LMSFilter<DataType_>::get_memory() const
  {
    return impl->alpha;
  }

  template<typename DataType_>
  void LMSFilter<DataType_>::set_mu(double mu)
  {
    if (mu >= 1)
    {
      throw ATK::RuntimeError("Mu must be less than 1");
    }
    if (mu <= 0)
    {
      throw ATK::RuntimeError("Mu must be strictly positive");
    }

    impl->mu = mu;
  }

  template<typename DataType_>
  double LMSFilter<DataType_>::get_mu() const
  {
    return impl->mu;
  }

  template<typename DataType_>
  void LMSFilter<DataType_>::set_mode(Mode mode)
  {
    this->mode = mode;
  }

  template<typename DataType_>
  typename LMSFilter<DataType_>::Mode LMSFilter<DataType_>::get_mode() const
  {
    return mode;
  }

  template<typename DataType_>
  void LMSFilter<DataType_>::process_impl(gsl::index size) const
  {
    const DataType* ATK_RESTRICT input = converted_inputs[0];
    const DataType* ATK_RESTRICT ref = converted_inputs[1];
    DataType* ATK_RESTRICT output = outputs[0];

    auto update_function = impl->select(mode);

    for(gsl::index i = 0; i < size; ++i)
    {
      typename LMSFilterImpl::xType x(input - input_delay + i, input_delay + 1, 1);
      output[i] = impl->w.conjugate().dot(x);
      if(learning)
      {
        (impl.get()->*update_function)(x, TypeTraits<DataType>::conj(ref[i] - output[i]));
      }
    }
  }

  template<typename DataType_>
  const DataType_* LMSFilter<DataType_>::get_w() const
  {
    return impl->w.data();
  }

  template<typename DataType_>
  void LMSFilter<DataType_>::set_w(gsl::not_null<const DataType_*> w)
  {
    impl->w = Eigen::Map<const typename LMSFilterImpl::wType>(w.get(), get_size());
  }

  template<typename DataType_>
  void LMSFilter<DataType_>::set_learning(bool learning)
  {
    this->learning = learning;
  }

  template<typename DataType_>
  bool LMSFilter<DataType_>::get_learning() const
  {
    return learning;
  }

  template class LMSFilter<double>;
#if ATK_ENABLE_INSTANTIATION
  template class LMSFilter<float>;
  template class LMSFilter<std::complex<float>>;
  template class LMSFilter<std::complex<double>>;
#endif
}


#include <map>
#include <numeric>

#include <Eigen/Dense>

#include "BingoCpp/acyclic_graph.h"
#include "BingoCpp/backend.h"
#include "BingoCpp/backend_nodes.h"

const int NODE_IDX = 0;
const int OP_1 = 1;
const int OP_2 = 2;

namespace bingo {
namespace {

Eigen::ArrayXXd reverse_eval(const std::pair<int, int>& deriv_shape,
                             const int deriv_wrt_node,
                             const std::vector<Eigen::ArrayXXd>& forward_eval,
                             const Eigen::ArrayX3i& stack) {
  int num_samples = deriv_shape.first;
  int num_features = deriv_shape.second;
  int stack_depth = stack.rows();

  Eigen::ArrayXXd derivative = Eigen::ArrayXXd::Zero(num_samples, num_features);
  std::vector<Eigen::ArrayXXd> reverse_eval(stack_depth);
  for (int row = 0; row < stack_depth; row++) {
      reverse_eval[row] = Eigen::ArrayXd::Zero(num_samples);
  }

  reverse_eval[stack_depth-1] = Eigen::ArrayXd::Ones(num_samples);
  for (int i = stack_depth - 1; i >= 0; i--) {
    int node = stack(i, NODE_IDX);
    int param1 = stack(i, OP_1);
    int param2 = stack(i, OP_2);
    if (node == deriv_wrt_node) {
      derivative.col(param1) += reverse_eval[i];
    } else {
      reverse_eval_function(node, i, param1, param2, forward_eval, reverse_eval);
    }
  }
  return derivative;
}

Eigen::ArrayXXd reverse_eval_with_mask(const std::pair<int, int>& deriv_shape,
                                       const int deriv_wrt_node,
                                       const std::vector<Eigen::ArrayXXd>& forward_eval,
                                       const Eigen::ArrayX3i& stack,
                                       const std::vector<bool>& mask) {
  int num_samples = deriv_shape.first;
  int num_features = deriv_shape.second;
  int stack_depth = stack.rows();

  Eigen::ArrayXXd derivative = Eigen::ArrayXXd::Zero(num_samples, num_features);
  std::vector<Eigen::ArrayXXd> reverse_eval(stack_depth);
  for (int row = 0; row < stack_depth; row++) {
    if (mask[row]) {
      reverse_eval[row] = Eigen::ArrayXd::Zero(num_samples);
    }
  }

  reverse_eval[stack_depth-1] = Eigen::ArrayXd::Ones(num_samples);
  for (int i = stack_depth - 1; i >= 0; i--) {
    if (mask[i]) {
      int node = stack(i, NODE_IDX);
      int param1 = stack(i, OP_1);
      int param2 = stack(i, OP_2);
      if (node == deriv_wrt_node) {
        derivative.col(param1) += reverse_eval[i];
      } else {
        reverse_eval_function(node, i, param1, param2, forward_eval, reverse_eval);
      }
    }
  }
  return derivative;
}

std::vector<Eigen::ArrayXXd> forward_eval(
    const Eigen::ArrayX3i& stack,
    const Eigen::ArrayXXd& x,
    const Eigen::VectorXd& constants) {
  std::vector<Eigen::ArrayXXd> _forward_eval(stack.rows());

  for (int i = 0; i < stack.rows(); ++i) {
    int node = stack(i, NODE_IDX);
    int op1 = stack(i, OP_1);
    int op2 = stack(i, OP_2);
    _forward_eval[i] = forward_eval_function(
      node, op1, op2, x, constants, _forward_eval);
  }
  return _forward_eval;
}

std::vector<Eigen::ArrayXXd> forward_eval_with_mask(
    const Eigen::ArrayX3i& stack,
    const Eigen::ArrayXXd& x,
    const Eigen::VectorXd& constants,
    const std::vector<bool>& mask) {
  std::vector<Eigen::ArrayXXd> _forward_eval(stack.rows());

  for (int i = 0; i < stack.rows(); ++i) {
    if (mask[i]) {
      int node = stack(i, NODE_IDX);
      int op1 = stack(i, OP_1);
      int op2 = stack(i, OP_2);
      _forward_eval[i] = forward_eval_function(
        node, op1, op2, x, constants, _forward_eval);
    }
  }
  return _forward_eval;
}

std::pair<Eigen::ArrayXXd, Eigen::ArrayXXd> _evaluate_with_derivative(
    const Eigen::ArrayX3i& stack,
    const Eigen::ArrayXXd& x,
    const Eigen::VectorXd& constants,
    const bool param_x_or_c) {
  std::vector<Eigen::ArrayXXd> _forward_eval = forward_eval(
      stack, x, constants);

  std::pair<int, int> deriv_shape;
  int deriv_wrt_node;
  if (param_x_or_c) {  // true = x
    deriv_shape = std::make_pair(x.rows(), x.cols());
    deriv_wrt_node = 0;
  } else {  // false = c
    deriv_shape = std::make_pair(x.rows(), constants.size());
    deriv_wrt_node = 1;
  }

  Eigen::ArrayXXd derivative = reverse_eval(
      deriv_shape, deriv_wrt_node, _forward_eval, stack);
  return std::make_pair(_forward_eval.back(), derivative);
}

std::pair<Eigen::ArrayXXd, Eigen::ArrayXXd> evaluate_with_derivative_and_mask(
    const Eigen::ArrayX3i& stack,
    const Eigen::ArrayXXd& x,
    const Eigen::VectorXd& constants,
    const std::vector<bool>& mask,
    const bool param_x_or_c) {
  std::vector<Eigen::ArrayXXd> forward_eval = forward_eval_with_mask(
      stack, x, constants, mask);

  std::pair<int, int> deriv_shape;
  int deriv_wrt_node;
  if (param_x_or_c) {  // true = x
    deriv_shape = std::make_pair(x.rows(), x.cols());
    deriv_wrt_node = 0;
  } else {  // false = c
    deriv_shape = std::make_pair(x.rows(), constants.size());
    deriv_wrt_node = 1;
  }

  Eigen::ArrayXXd derivative = reverse_eval_with_mask(
      deriv_shape, deriv_wrt_node, forward_eval, stack, mask);
  return std::make_pair(forward_eval.back(), derivative);
}
} // namespace

bool is_cpp() {
    return true;
}

Eigen::ArrayXXd evaluate(const Eigen::ArrayX3i& stack,
                         const Eigen::ArrayXXd& x,
                         const Eigen::VectorXd& constants) {
  std::vector<Eigen::ArrayXXd> _forward_eval = forward_eval(
      stack, x, constants);
  return _forward_eval.back();
}

std::pair<Eigen::ArrayXXd, Eigen::ArrayXXd> evaluate_with_derivative(
    const Eigen::ArrayX3i& stack,
    const Eigen::ArrayXXd& x,
    const Eigen::VectorXd& constants,
    const bool param_x_or_c) {
  return _evaluate_with_derivative(stack, x, constants, param_x_or_c);
}

Eigen::ArrayXXd simplify_and_evaluate(const Eigen::ArrayX3i& stack,
                                    const Eigen::ArrayXXd& x,
                                    const Eigen::VectorXd& constants) {
  std::vector<bool> mask = get_utilized_commands(stack);
  std::vector<Eigen::ArrayXXd> forward_eval = forward_eval_with_mask(
      stack, x, constants, mask);
  return forward_eval.back();
}

std::pair<Eigen::ArrayXXd, Eigen::ArrayXXd> simplify_and_evaluate_with_derivative(
    const Eigen::ArrayX3i& stack,
    const Eigen::ArrayXXd& x,
    const Eigen::VectorXd& constants,
    const bool param_x_or_c) {
  std::vector<bool> mask = get_utilized_commands(stack);
  return evaluate_with_derivative_and_mask(stack, x, constants, mask, param_x_or_c);
}

std::vector<bool> get_utilized_commands(const Eigen::ArrayX3i& stack) {
  std::vector<bool> used_commands(stack.rows());
  used_commands.back() = true;
  int stack_size = stack.rows();
  for (int i = 1; i < stack_size; i++) {
    int row = stack_size - i;
    int node = stack(row, NODE_IDX);
    int param1 = stack(row, OP_1);
    int param2 = stack(row, OP_2);
    if (used_commands[row] && node > 1) {
      used_commands[param1] = true;
      if (AcyclicGraph::has_arity_two(node)) {
        used_commands[param2] = true;
      }
    }
  }
  return used_commands;
}

Eigen::ArrayX3i simplify_stack(const Eigen::ArrayX3i& stack) {
  std::vector<bool> used_command = get_utilized_commands(stack);
  std::map<int, int> reduced_param_map;
  int num_commands = 0;
  num_commands = std::accumulate(used_command.begin(), used_command.end(), 0);
  Eigen::ArrayX3i new_stack(num_commands, 3);

  for (int i = 0, j = 0; i < stack.rows(); ++i) {
    if (used_command[i]) {
      new_stack(j, 0) = stack(i, 0);
      if (AcyclicGraph::is_terminal(new_stack(j, 0))) {
        new_stack(j, 1) = stack(i, 1);
        new_stack(j, 2) = stack(i, 2);
      } else {
        new_stack(j, 1) = reduced_param_map[stack(i, 1)];
        if (AcyclicGraph::has_arity_two(new_stack(j, 0))) {
          new_stack(j, 2) = reduced_param_map[stack(i, 2)];
        } else {
          new_stack(j, 2) = new_stack(j, 1);
        }
      }
      reduced_param_map[i] = j;
      ++j;
    }
  }
  return new_stack;
}

int get_arity(int node) {
  if (AcyclicGraph::is_terminal(node)) return 0;
  return AcyclicGraph::has_arity_two(node) ? 2 : 1;
}
} // namespace bingo


#include <Sphere.h>

#include <Eigen/Dense>

#include <Util.h>

#include <cmath>
#include <cstdio>

void Sphere::generate_geometry(int width, int height) {
    int  gNumVertices  = 0;
    int  gNumTriangles = 0;
    int* gIndexBuffer  = nullptr;

    float theta, phi;
    int t;

    gNumVertices    = (height - 2) * width + 2;
    gNumTriangles   = (height - 2) * (width - 1) * 2;

    // TODO: Allocate an array for gNumVertices vertices.
    Eigen::Vector3f* vertices = new Eigen::Vector3f[gNumVertices];

    gIndexBuffer = new int[3*gNumTriangles];

    t = 0;
    for (int j = 1; j < height-1; ++j) {
        for (int i = 0; i < width; ++i) {
            theta = (float) j / (height-1) * M_PI;
            phi   = (float) i / (width-1)  * M_PI * 2;

            float   x   = sinf(theta) * cosf(phi);
            float   y   = cosf(theta);
            float   z   = -sinf(theta) * sinf(phi);

            // TODO: Set vertex t in the vertex array to {x, y, z}.
            vertices[t] = Eigen::Vector3f(x, y, z);
            t++;
        }
    }

    // TODO: Set vertex t in the vertex array to {0, 1, 0}.
    vertices[t] = Eigen::Vector3f(0, 1, 0);
    t++;

    // TODO: Set vertex t in the vertex array to {0, -1, 0}.
    vertices[t] = Eigen::Vector3f(0, -1, 0);
    t++;

    t = 0;
    for (int j = 0; j < height-3; ++j) {
        for (int i = 0; i < width-1; ++i) {
            gIndexBuffer[t++] = j*width + i;
            gIndexBuffer[t++] = (j+1)*width + (i+1);
            gIndexBuffer[t++] = j*width + (i+1);
            gIndexBuffer[t++] = j*width + i;
            gIndexBuffer[t++] = (j+1)*width + i;
            gIndexBuffer[t++] = (j+1)*width + (i+1);
        }
    }
    for (int i = 0; i < width-1; ++i) {
        gIndexBuffer[t++] = (height-2)*width;
        gIndexBuffer[t++] = i;
        gIndexBuffer[t++] = i + 1;
        gIndexBuffer[t++] = (height-2)*width + 1;
        gIndexBuffer[t++] = (height-3)*width + (i+1);
        gIndexBuffer[t++] = (height-3)*width + i;
    }

    // Create triangles
    for (int i = 0; i < gNumTriangles; i++) {
        Eigen::Vector3f a = vertices[gIndexBuffer[3*i + 0]];
        Eigen::Vector3f b = vertices[gIndexBuffer[3*i + 1]];
        Eigen::Vector3f c = vertices[gIndexBuffer[3*i + 2]];

        triangles.push_back(Triangle(a, b, c));
    }

    delete[] gIndexBuffer;
    delete[] vertices;
}

void Sphere::transform_geometry(const Eigen::Matrix4f& xform) {
    for (auto& tri : triangles) {
        // Transformed Cartesian vertex coordinates
        auto a = vec4to3(xform * Eigen::Vector4f(tri.a[0], tri.a[1], tri.a[2], 1.0f));
        auto b = vec4to3(xform * Eigen::Vector4f(tri.b[0], tri.b[1], tri.b[2], 1.0f));
        auto c = vec4to3(xform * Eigen::Vector4f(tri.c[0], tri.c[1], tri.c[2], 1.0f));

        // Transform vertices
        tri.a = a;
        tri.b = b;
        tri.c = c;

        // Transform normal
        tri.n = ((b-a).cross(c-a)).normalized();

        // Transform vertex normals
        tri.an = (a-center).normalized();
        tri.bn = (b-center).normalized();
        tri.cn = (c-center).normalized();
    }
}


/*! @file pitts_tensortrain_from_dense_twosided.hpp
* @brief conversion of a dense tensor to the tensor-train format (based on a hopefully faster TSQR algorithm, two-sided variant)
* @author Melven Roehrig-Zoellner <Melven.Roehrig-Zoellner@DLR.de>
* @date 2020-08-08
* @copyright Deutsches Zentrum fuer Luft- und Raumfahrt e. V. (DLR), German Aerospace Center
*
**/

// include guard
#ifndef PITTS_TENSORTRAIN_FROM_DENSE_TWOSIDED_HPP
#define PITTS_TENSORTRAIN_FROM_DENSE_TWOSIDED_HPP

// includes
#include <limits>
#include <numeric>
#pragma GCC push_options
#pragma GCC optimize("no-unsafe-math-optimizations")
#include <Eigen/Dense>
#pragma GCC pop_options
#include "pitts_tensortrain.hpp"
#include "pitts_multivector.hpp"
#include "pitts_multivector_tsqr.hpp"
#include "pitts_multivector_transform.hpp"
#include "pitts_multivector_transpose.hpp"
#include "pitts_multivector_eigen_adaptor.hpp"
#include "pitts_tensor2.hpp"
#include "pitts_tensor2_eigen_adaptor.hpp"
#include "pitts_timer.hpp"

//! namespace for the library PITTS (parallel iterative tensor train solvers)
namespace PITTS
{
  //! namespace for helper functionality
  namespace internal
  {
    //! helper namespace for high-order SVD functionality (e.g. TensorTrain_fromDense)
    namespace HOSVD
    {
      template<typename T>
      void split(const MultiVector<T>& X, MultiVector<T>& Y, Tensor2<T>& M, int nextDim, T rankTolerance, int maxRank)
      {
        using EigenMatrix = Eigen::Matrix<T, Eigen::Dynamic, Eigen::Dynamic>;
        Eigen::JacobiSVD<EigenMatrix> svd;
std::cout << "HOSVD::split  matrix dimensions: " << X.rows() << " x " << X.cols() << "\n";
        if( X.rows() > 10*X.cols() )
        {
          // calculate QR decomposition (QR-trick: X=QR,SVD(R))
          block_TSQR(X, M, 0, false);
          svd.compute(ConstEigenMap(M), Eigen::ComputeThinU | Eigen::ComputeThinV);
        }
        else
        {
          svd.compute(ConstEigenMap(X), Eigen::ComputeThinU | Eigen::ComputeThinV);
        }

std::cout << "singular values: " << svd.singularValues().transpose() << "\n";

        // truncate svd
        svd.setThreshold(rankTolerance);
        int rank = svd.rank();
        if( maxRank > 0 )
          rank = std::min(maxRank, rank);

        // copy right singular vectors
        M.resize(X.cols(), rank);
        EigenMap(M) = svd.matrixV().leftCols(rank);

        // transform input, s.t. Y \approx X M
        transform(X, M, Y, {X.rows()/nextDim, rank*nextDim});
      }
    }
  }

  //! calculate tensor-train decomposition of a tensor stored in fully dense format
  //!
  //! Passing a large enough buffer in work helps to avoid costly reallocations + later page-faults for large data.
  //!
  //! @warning To reduce memory overhead, this function will overwrite the input arguments with temporary data.
  //!          Please pass a copy of the data if you still need it!
  //!
  //! @tparam T         underlying data type (double, complex, ...)
  //!
  //! @param X              input tensor, overwritten and modified output, dimension must be (size/lastDim, lastDim) where lastDim = dimensions.back()
  //! @param dimensions     tensor dimensions, input is interpreted in Fortran storage order (first index changes the fastest)
  //! @param work           buffer for temporary data, will be resized and modified
  //! @param rankTolerance  approximation accuracy, used to reduce the TTranks of the resulting tensor train
  //! @param maxRank        maximal TTrank (bond dimension), unbounded by default
  //! @return               resulting tensor train
  //!
  template<typename T>
  TensorTrain<T> fromDense_twoSided(MultiVector<T>& X, MultiVector<T>& work, const std::vector<int>& dimensions, T rankTolerance = std::sqrt(std::numeric_limits<T>::epsilon()), int maxRank = -1)
  {
    // timer
    const auto timer = PITTS::timing::createScopedTimer<TensorTrain<T>>();

    // abort early for zero dimensions
    if( dimensions.size() == 0 )
    {
      if( X.rows()*X.cols() !=  0 )
        throw std::out_of_range("Mismatching dimensions in TensorTrain<T>::fromDense");
      return TensorTrain<T>{dimensions};
    }

    const auto totalSize = std::accumulate(begin(dimensions), end(dimensions), (std::ptrdiff_t)1, std::multiplies<std::ptrdiff_t>());
    const auto nDims = dimensions.size();
    if( X.rows() != totalSize/dimensions[nDims-1] || X.cols() != dimensions[nDims-1] )
      throw std::out_of_range("Mismatching dimensions in TensorTrain<T>::fromDense");

    TensorTrain<T> result(dimensions);

    // actually convert to tensor train format
    Tensor2<T> M;
    for(int ii = 0; ii < nDims; ii++)
    {
      if( ii % 2 == 0 )
      {
        // right part
        const auto iDim = nDims - 1 - ii/2;
        if( ii > 0 )
        {
          const auto r1 = result.subTensors()[iDim+1].r1();
          const auto n = dimensions[iDim];
          transpose(work, X, {(work.rows()*work.cols())/(n*r1), n*r1}, true);
        }
        if( ii != nDims-1 )
        {
          internal::HOSVD::split(X, work, M, 1, rankTolerance, maxRank);
        }
        else
        {
          M.resize(X.cols(), X.rows());
          EigenMap(M) = ConstEigenMap(X).transpose();
        }

        auto& subT = result.editableSubTensors()[iDim];
        const int rank = M.r2();
        subT.resize(rank, dimensions[iDim], X.cols()/dimensions[iDim]);
        for(int i = 0; i < rank; i++)
          for(int j = 0; j < subT.n(); j++)
            for(int k = 0; k < subT.r2(); k++)
              subT(i,j,k) = M(j+subT.n()*k, i);
      }
      else
      {
        // left part
        const auto iDim = ii / 2;
        {
          const auto r2 = (iDim == 0 ) ? 1 : result.subTensors()[iDim-1].r2();
          const auto n = dimensions[iDim];
          transpose(work, X, {(work.cols()*work.rows())/(n*r2), n*r2}, false);
        }
        if( ii != nDims - 1 )
        {
          internal::HOSVD::split(X, work, M, 1, rankTolerance, maxRank);
        }
        else
        {
          M.resize(X.cols(), X.rows());
          EigenMap(M) = ConstEigenMap(X).transpose();
        }

        auto& subT = result.editableSubTensors()[iDim];
        const int rank = M.r2();
        subT.resize(X.cols()/dimensions[iDim], dimensions[iDim], rank);
        for(int i = 0; i < rank; i++)
          for(int j = 0; j < subT.n(); j++)
            for(int k = 0; k < subT.r1(); k++)
              subT(k,j,i) = M(k+j*subT.r1(), i);
      }
    }

    return result;
  }

}

#endif // PITTS_TENSORTRAIN_FROM_DENSE_TWOSIDED_HPP


#ifndef PYBNESIAN_LEARNING_INDEPENDENCES_CONTINUOUS_RCOT_HPP
#define PYBNESIAN_LEARNING_INDEPENDENCES_CONTINUOUS_RCOT_HPP

#include <random>
#include <Eigen/Eigenvalues>
#include <learning/independences/independence.hpp>
#include <util/math_constants.hpp>
#include <util/basic_eigen_ops.hpp>
#include <util/chisquaresum.hpp>

using learning::independences::IndependenceTest;

namespace learning::independences::continuous {

template <typename MatrixType>
typename MatrixType::Scalar rf_sigma_impl(MatrixType& m) {
    using Scalar = typename MatrixType::Scalar;
    using VectorType = Matrix<Scalar, Dynamic, 1>;
    auto r = std::min(static_cast<decltype(m.rows())>(500), m.rows());
    VectorType distances(r * (r - 1) / 2);

    for (int i = r - 1, j = 0; i > 0; --i, j += i) {
        distances.segment(j, i) = (m.topRows(i).rowwise() - m.row(i)).matrix().rowwise().norm();
    }

    // Compute median
    int median_index = distances.rows() / 2;

    double median;
    if (distances.rows() % 2 == 1) {
        std::nth_element(distances.data(), distances.data() + median_index, distances.data() + distances.rows());
        median = distances[median_index];
    } else {
        std::sort(distances.data(), distances.data() + distances.rows());
        median = 0.5 * (distances[median_index - 1] + distances[median_index]);
    }

    if (median == 0) median = 1;

    return median;
}

class RCoT : public IndependenceTest {
public:
    RCoT(const DataFrame& df, int random_fourier_xy = 5, int random_fourier_z = 100)
        : m_df(df.normalize()),
          m_num_random_fourier_xy(random_fourier_xy),
          m_num_random_fourier_z(random_fourier_z),
          m_dfourier_x(),
          m_dfourier_y(),
          m_dfourier_z(),
          m_dsigma(),
          m_ffourier_x(),
          m_ffourier_y(),
          m_ffourier_z(),
          m_fsigma() {
        auto continuous_indices = df.continuous_columns();

        if (continuous_indices.size() < 2) {
            throw std::invalid_argument("DataFrame does not contain enough continuous columns.");
        }

        auto type = m_df.same_type(continuous_indices);

        switch (type->id()) {
            case Type::DOUBLE: {
                m_dfourier_x = MatrixXd(df.num_rows(), m_num_random_fourier_xy);
                m_dfourier_y = MatrixXd(df.num_rows(), m_num_random_fourier_xy);
                m_dfourier_z = MatrixXd(df.num_rows(), m_num_random_fourier_z);
                m_tmp_dcov = MatrixXd(df.num_rows(), m_num_random_fourier_xy * m_num_random_fourier_xy);
                m_dsigma = VectorXd(df->num_columns());

                for (auto c : continuous_indices) {
                    if (m_df.null_count(c) == 0) {
                        auto x_vec = m_df.to_eigen<false, arrow::DoubleType, false>(c);
                        m_dsigma(c) = rf_sigma_impl(*x_vec);
                    }
                }

                break;
            }
            case Type::FLOAT: {
                m_ffourier_x = MatrixXf(df.num_rows(), m_num_random_fourier_xy);
                m_ffourier_y = MatrixXf(df.num_rows(), m_num_random_fourier_xy);
                m_ffourier_z = MatrixXf(df.num_rows(), m_num_random_fourier_z);
                m_tmp_fcov = MatrixXf(df.num_rows(), m_num_random_fourier_xy * m_num_random_fourier_xy);
                m_fsigma = VectorXf(df->num_columns());

                for (auto c : continuous_indices) {
                    if (m_df.null_count(c) == 0) {
                        auto x_vec = m_df.to_eigen<false, arrow::FloatType, false>(c);
                        m_fsigma(c) = rf_sigma_impl(*x_vec);
                    }
                }

                break;
            }
            default:
                throw std::runtime_error("[RCoT] Unreachable code");
        }
    }

    double pvalue(const std::string& x, const std::string& y) const override;
    template <typename ArrowType>
    double pvalue(const std::string& x, const std::string& y) const;

    double pvalue(const std::string& x, const std::string& y, const std::string& z) const override;
    template <typename ArrowType>
    double pvalue(const std::string& x, const std::string& y, const std::string& z) const;

    double pvalue(const std::string& x, const std::string& y, const std::vector<std::string>& z) const override;
    template <typename ArrowType>
    double pvalue(const std::string& x, const std::string& y, const std::vector<std::string>& z) const;

    int num_variables() const override { return m_df->num_columns(); }

    std::vector<std::string> variable_names() const override { return m_df.column_names(); }

    const std::string& name(int i) const override { return m_df.name(i); }

    bool has_variables(const std::string& name) const override { return m_df.has_columns(name); }

    bool has_variables(const std::vector<std::string>& cols) const override { return m_df.has_columns(cols); }

private:
    template <typename Scalar>
    Scalar rf_sigma(int index) const {
        if constexpr (std::is_same_v<Scalar, double>)
            return m_dsigma(index);
        else
            return m_fsigma(index);
    }

    template <typename Scalar>
    Matrix<Scalar, Dynamic, Dynamic>& fourier_x() const {
        if constexpr (std::is_same_v<Scalar, double>)
            return m_dfourier_x;
        else
            return m_ffourier_x;
    }

    template <typename Scalar>
    Matrix<Scalar, Dynamic, Dynamic>& fourier_y() const {
        if constexpr (std::is_same_v<Scalar, double>)
            return m_dfourier_y;
        else
            return m_ffourier_y;
    }

    template <typename Scalar>
    Matrix<Scalar, Dynamic, Dynamic>& fourier_z() const {
        if constexpr (std::is_same_v<Scalar, double>)
            return m_dfourier_z;
        else
            return m_ffourier_z;
    }

    template <typename Scalar>
    Matrix<Scalar, Dynamic, Dynamic>& tmp_cov() const {
        if constexpr (std::is_same_v<Scalar, double>)
            return m_tmp_dcov;
        else
            return m_tmp_fcov;
    }

    template <typename Mat>
    Matrix<typename Mat::Scalar, Dynamic, 1> eigenvalues_covariance(Mat& fourier_x, Mat& fourier_y) const;

    template <typename VectorType, typename FeatureType>
    double RIT_impl(
        VectorType& x, VectorType& y, FeatureType& feat_x, FeatureType& feat_y, double sigma_x, double sigma_y) const;
    template <bool contains_null, typename VectorType>
    double RIT(int x_index, int y_index, VectorType& x, VectorType& y) const;

    template <typename VectorType, typename MatType, typename FeatureType>
    double TestWithZ_impl(VectorType& x,
                          VectorType& y,
                          MatType& z,
                          FeatureType& feat_x,
                          FeatureType& feat_y,
                          FeatureType& feat_z,
                          double sigma_x,
                          double sigma_y,
                          double sigma_z) const;

    template <bool contains_null, typename VectorType>
    double RSingleZ(int x_index, int y_index, int z_index, VectorType& x, VectorType& y, VectorType& z) const;

    template <bool contains_null, typename VectorType, typename MatType>
    double RMultiZ(int x_index, int y_index, VectorType& x, VectorType& y, MatType& z) const;

    DataFrame m_df;
    int m_num_random_fourier_xy;
    int m_num_random_fourier_z;
    // Cache fourier matrices and sigmas (double or float).
    mutable MatrixXd m_dfourier_x;
    mutable MatrixXd m_dfourier_y;
    mutable MatrixXd m_dfourier_z;
    mutable MatrixXd m_tmp_dcov;
    VectorXd m_dsigma;
    mutable MatrixXf m_ffourier_x;
    mutable MatrixXf m_ffourier_y;
    mutable MatrixXf m_ffourier_z;
    mutable MatrixXf m_tmp_fcov;
    VectorXf m_fsigma;
};

template <typename InputMatrix, typename OutputMatrix>
void random_fourier_features(InputMatrix& m,
                             typename InputMatrix::Scalar sigma,
                             int num_features,
                             OutputMatrix& fourier_features) {
    static_assert(std::is_same_v<typename InputMatrix::Scalar, typename OutputMatrix::Scalar>,
                  "Input/Output matrices must have the same type");

    using Scalar = typename InputMatrix::Scalar;
    using MatrixType = Matrix<Scalar, Dynamic, Dynamic>;
    using VectorType = Matrix<Scalar, Dynamic, 1>;

    MatrixType W(m.cols(), num_features);
    VectorType b(num_features);

    std::mt19937 rng(std::random_device{}());
    std::normal_distribution<Scalar> normal;
    for (auto j = 0; j < W.cols(); ++j) {
        for (auto i = 0; i < W.rows(); ++i) {
            W(i, j) = normal(rng);
        }
    }
    W *= (1 / sigma);

    std::uniform_real_distribution<Scalar> unif;
    for (auto i = 0; i < num_features; ++i) {
        b(i) = unif(rng);
    }
    b *= 2 * util::pi<Scalar>;

    fourier_features.noalias() = (m * W).rowwise() + b.transpose();
    fourier_features = fourier_features.array().cos().matrix();
    fourier_features = fourier_features * util::root_two<Scalar>;
}

template <typename Mat, typename TmpMat>
Matrix<typename Mat::Scalar, Dynamic, 1> eigenvalues_covariance_impl(Mat& fourier_x, Mat& fourier_y, TmpMat& tmp_mat) {
    using Scalar = typename Mat::Scalar;
    using MatrixType = Matrix<Scalar, Dynamic, Dynamic>;

    for (int i = 0; i < fourier_x.cols(); ++i) {
        tmp_mat.block(0, i * fourier_y.cols(), tmp_mat.rows(), fourier_y.cols()) =
            fourier_y.array().colwise() * fourier_x.col(i).array();
    }
    auto cov = util::sse_mat(tmp_mat) * (1 / static_cast<Scalar>(fourier_x.rows()));
    auto eigen_solver = Eigen::SelfAdjointEigenSolver<MatrixType>(cov, Eigen::DecompositionOptions::EigenvaluesOnly);
    return eigen_solver.eigenvalues();
}

template <typename Mat>
Matrix<typename Mat::Scalar, Dynamic, 1> RCoT::eigenvalues_covariance(Mat& fourier_x, Mat& fourier_y) const {
    using Scalar = typename Mat::Scalar;
    auto& cc = tmp_cov<Scalar>();

    if (fourier_x.rows() != cc.rows()) {
        auto tmp = cc.topRows(fourier_x.rows());
        return eigenvalues_covariance_impl(fourier_x, fourier_y, tmp);
    } else {
        return eigenvalues_covariance_impl(fourier_x, fourier_y, cc);
    }
}

template <typename VectorType>
Matrix<typename VectorType::Scalar, Dynamic, 1> filter_positive_elements(const VectorType& v) {
    using Scalar = typename VectorType::Scalar;
    using NewVectorType = Matrix<Scalar, Dynamic, 1>;
    std::vector<Scalar> positive;
    for (int i = 0; i < v.rows(); ++i) {
        if (v(i) > 0) positive.push_back(v(i));
    }

    NewVectorType res(positive.size());
    for (size_t i = 0; i < positive.size(); ++i) {
        res(i) = positive[i];
    }

    return res;
}

template <typename VectorType, typename FeatureType>
double RCoT::RIT_impl(
    VectorType& x, VectorType& y, FeatureType& feat_x, FeatureType& feat_y, double sigma_x, double sigma_y) const {
    random_fourier_features(x, sigma_x, m_num_random_fourier_xy, feat_x);
    random_fourier_features(y, sigma_y, m_num_random_fourier_xy, feat_y);

    util::normalize_cols(feat_x);
    util::normalize_cols(feat_y);

    auto Cxy = util::cov(feat_x, feat_y);
    auto sta = x.rows() * Cxy.squaredNorm();
    auto eigs = eigenvalues_covariance(feat_x, feat_y);
    auto pos_eigs = filter_positive_elements(eigs);

    if (pos_eigs.rows() < 4) {
        auto pvalue = util::hbe_complement(pos_eigs, sta);
        if (pvalue < 0) return 0;
        return pvalue;
    }

    try {
        auto pvalue = util::lpb4_complement(pos_eigs, sta);
        if (pvalue < 0) return 0;
        return pvalue;
    } catch (std::exception&) {
        auto pvalue = util::hbe_complement(pos_eigs, sta);
        if (pvalue < 0) return 0;
        return pvalue;
    }
}

template <bool contains_null, typename VectorType>
double RCoT::RIT(int x_index, int y_index, VectorType& x, VectorType& y) const {
    using Scalar = typename VectorType::Scalar;

    if constexpr (contains_null) {
        Scalar sigma_x = rf_sigma_impl(x);
        Scalar sigma_y = rf_sigma_impl(y);

        auto feat_x = fourier_x<Scalar>().topRows(x.rows());
        auto feat_y = fourier_y<Scalar>().topRows(y.rows());

        return RIT_impl(x, y, feat_x, feat_y, sigma_x, sigma_y);
    } else {
        Scalar sigma_x = rf_sigma<Scalar>(x_index);
        Scalar sigma_y = rf_sigma<Scalar>(y_index);

        auto& feat_x = fourier_x<Scalar>();
        auto& feat_y = fourier_y<Scalar>();

        return RIT_impl(x, y, feat_x, feat_y, sigma_x, sigma_y);
    }
}

template <typename VectorType, typename MatType, typename FeatureType>
double RCoT::TestWithZ_impl(VectorType& x,
                            VectorType& y,
                            MatType& z,
                            FeatureType& feat_x,
                            FeatureType& feat_y,
                            FeatureType& feat_z,
                            double sigma_x,
                            double sigma_y,
                            double sigma_z) const {
    random_fourier_features(x, sigma_x, m_num_random_fourier_xy, feat_x);
    random_fourier_features(y, sigma_y, m_num_random_fourier_xy, feat_y);
    random_fourier_features(z, sigma_z, m_num_random_fourier_z, feat_z);

    util::normalize_cols(feat_x);
    util::normalize_cols(feat_y);
    util::normalize_cols(feat_z);

    auto Cxy = util::cov(feat_x, feat_y);

    auto Czz = util::cov(feat_z);
    Czz.diagonal().array() += 1e-10;

    auto i_Czz = Czz.inverse();

    auto Cxz = util::cov(feat_x, feat_z);
    auto Czy = util::cov(feat_z, feat_y);

    auto z_i_Czz = feat_z * i_Czz;
    feat_x = feat_x - z_i_Czz * Cxz.transpose();
    feat_y = feat_y - z_i_Czz * Czy;

    auto Cxy_z = Cxy - Cxz * i_Czz * Czy;

    auto sta = x.rows() * Cxy_z.squaredNorm();
    auto eigs = eigenvalues_covariance(feat_x, feat_y);
    auto pos_eigs = filter_positive_elements(eigs);

    if (m_num_random_fourier_z == 1 || pos_eigs.rows() < 4) {
        auto pvalue = util::hbe_complement(pos_eigs, sta);
        if (pvalue < 0) return 0;
        return pvalue;
    }

    try {
        auto pvalue = util::lpb4_complement(pos_eigs, sta);
        if (pvalue < 0) return 0;
        return pvalue;
    } catch (std::exception&) {
        auto pvalue = util::hbe_complement(pos_eigs, sta);
        if (pvalue < 0) return 0;
        return pvalue;
    }
}

template <bool contains_null, typename VectorType>
double RCoT::RSingleZ(int x_index, int y_index, int z_index, VectorType& x, VectorType& y, VectorType& z) const {
    using Scalar = typename VectorType::Scalar;

    if constexpr (contains_null) {
        Scalar sigma_x = rf_sigma_impl(x);
        Scalar sigma_y = rf_sigma_impl(y);
        Scalar sigma_z = rf_sigma_impl(z);

        auto feat_x = fourier_x<Scalar>().topRows(x.rows());
        auto feat_y = fourier_y<Scalar>().topRows(y.rows());
        auto feat_z = fourier_z<Scalar>().topRows(z.rows());

        return TestWithZ_impl(x, y, z, feat_x, feat_y, feat_z, sigma_x, sigma_y, sigma_z);
    } else {
        Scalar sigma_x = rf_sigma<Scalar>(x_index);
        Scalar sigma_y = rf_sigma<Scalar>(y_index);
        Scalar sigma_z = rf_sigma<Scalar>(z_index);

        auto& feat_x = fourier_x<Scalar>();
        auto& feat_y = fourier_y<Scalar>();
        auto& feat_z = fourier_z<Scalar>();

        return TestWithZ_impl(x, y, z, feat_x, feat_y, feat_z, sigma_x, sigma_y, sigma_z);
    }
}

template <bool contains_null, typename VectorType, typename MatType>
double RCoT::RMultiZ(int x_index, int y_index, VectorType& x, VectorType& y, MatType& z) const {
    using Scalar = typename VectorType::Scalar;

    if constexpr (contains_null) {
        Scalar sigma_x = rf_sigma_impl(x);
        Scalar sigma_y = rf_sigma_impl(y);
        Scalar sigma_z = rf_sigma_impl(z);

        auto feat_x = fourier_x<Scalar>().topRows(x.rows());
        auto feat_y = fourier_y<Scalar>().topRows(y.rows());
        auto feat_z = fourier_z<Scalar>().topRows(z.rows());

        return TestWithZ_impl(x, y, z, feat_x, feat_y, feat_z, sigma_x, sigma_y, sigma_z);
    } else {
        Scalar sigma_x = rf_sigma<Scalar>(x_index);
        Scalar sigma_y = rf_sigma<Scalar>(y_index);
        Scalar sigma_z = rf_sigma_impl(z);

        auto& feat_x = fourier_x<Scalar>();
        auto& feat_y = fourier_y<Scalar>();
        auto& feat_z = fourier_z<Scalar>();

        return TestWithZ_impl(x, y, z, feat_x, feat_y, feat_z, sigma_x, sigma_y, sigma_z);
    }
}

using DynamicRCoT = DynamicIndependenceTestAdaptator<RCoT>;

}  // namespace learning::independences::continuous

#endif  // PYBNESIAN_LEARNING_INDEPENDENCES_CONTINUOUS_RCOT_HPP