#include<iostream>
#include<iomanip>
#include<cmath>
#include"ROL_StdVector.hpp"
#include"Teuchos_LAPACK.hpp"
#include"LinearAlgebra.hpp"
 

#ifndef __ORTHOGONAL_POLYNOMIALS__
#define __ORTHOGONAL_POLYNOMIALS__
           

/** \brief Generate the Jacobi polynomial recursion coeffcients \f$a_k,b_k\f$
    
    The Jacobi polynomials satisfy the recurrence relation
    \f[ P^{(\alpha,\beta)}_{k+1}(x) = (x-a_k)P_{k}^{(\alpha,\beta)}(x) - b_k P_{k-1}^{(\alpha,\beta)}(x) \f]
    and form an orthogonal basis on \f$[-1,1]\f$ with respect to the weight function 
    \f$w(x)=(1-x)^\alpha(1+x)^\beta\f$.

    @param[in]    lapack   is a pointer to the Teuchos::LAPACK interface
    @param[in]    alpha    is a parameter that defines the weight function
    @param[in]    beta     is a parameter that defines the weight function
    @param[out]   ap       is a vector of recursion coefficients
    @param[out]   bp       is a vector of recursion coefficients
  
    Adapted from the MATLAB code by Dirk Laurie and Walter Gautschi
    http://www.cs.purdue.edu/archives/2002/wxg/codes/r_jacobi.m  */
template<class Real>
void rec_jacobi( ROL::Ptr<Teuchos::LAPACK<int,Real> > lapack,
                 const double alpha, 
                 const double beta,
                 std::vector<Real> &a,
                 std::vector<Real> &b ) {

    int N = a.size();
    Real nu = (beta-alpha)/double(alpha+beta+2.0);
    Real mu  = pow(2.0,alpha+beta+1.0)*tgamma(alpha+1.0)*tgamma(beta+1.0)/tgamma(alpha+beta+2.0);
    Real nab;
    Real sqdif = pow(beta,2)-pow(alpha,2);
                              
    a[0] = nu;
    b[0] = mu;
    
    if(N>1){
     
        for(int n=1;n<N;++n) {
            nab = 2*n+alpha+beta;
            a[n] = sqdif/(nab*(nab+2));              
        }

        b[1] = 4.0*(alpha+1.0)*(beta+1.0)/(pow(alpha+beta+2.0,2)*(alpha+beta+3.0));

        if(N>2) {
            for(int n=2;n<N;++n) {
                nab = 2*n+alpha+beta;
                b[n] = 4.0*(n+alpha)*(n+beta)*n*(n+alpha+beta)/(nab*nab*(nab+1.0)*(nab-1.0));
            }
        }
    }    
}

/** \brief Construct the generalized Vandermonde matrix (in column stacked form)
           based upon the recurrence coefficients (a,b) on the grid x 

    @param[in]   a   vector of recursion coefficients
    @param[in]   b   vector of recursion coefficients
    @param[in]   x   vector of quadrature nodes
    @param[in]   V   column-stacked Vandermonde matrix
*/   

template<class Real>
void vandermonde( const std::vector<Real> &a,
                  const std::vector<Real> &b,
                  const std::vector<Real> &x,
                        std::vector<Real> &V) {
    const int n = a.size();

    for(int i=0;i<n;++i) {
        // Zeroth polynomial is always 1
        V[i]   = 1.0;
        // First polynomial is (x-a) 
        V[i+n] = x[i] - a[0]; 
    }
    for(int j=1;j<n-1;++j) {
        for(int i=0;i<n;++i) { 
            V[i+(j+1)*n] = (x[i] - a[j])*V[i+j*n] - b[j]*V[i+(j-1)*n];  
        }      
    }        
}


/** \brief Compute the Gauss quadrature nodes and weights for the polynomials generated by the 
           recurrence coefficients.
    
    @param[in]   lapack   pointer to the Teuchos::LAPACK interface
    @param[in]   a        vector of recursion coefficients
    @param[in]   b        vector of recursion coefficients
    @param[out]  x        vector of quadrature nodes
    @param[out]  w        vector of quadrature weights

    Adapted from the MATLAB code by Walter Gautschi
    http://www.cs.purdue.edu/archives/2002/wxg/codes/gauss.m  */
template<class Real>
void gauss( ROL::Ptr<Teuchos::LAPACK<int,Real> > lapack,
            const std::vector<Real> &a,
            const std::vector<Real> &b,
            std::vector<Real> &x,
            std::vector<Real> &w ) {
    int INFO;
    const int N = a.size();  
    const int LDZ = N;
    const char COMPZ = 'I';

    std::vector<Real> D(N,0.0);
    std::vector<Real> E(N,0.0);
    std::vector<Real> WORK(4*N,0.0);

    // Column-stacked matrix of eigenvectors needed for weights
    std::vector<Real> Z(N*N,0);

    D = a;
     
    for(int i=0;i<N-1;++i) {
        E[i] = sqrt(b[i+1]);  
    }

    lapack->STEQR(COMPZ,N,&D[0],&E[0],&Z[0],LDZ,&WORK[0],&INFO);

    for(int i=0;i<N;++i){
        x[i] = D[i];
        w[i] = b[0]*pow(Z[N*i],2);
    } 

}



/** \brief Modify the given recurrence coefficients so that the set of zeros of the maximal
           order polynomial include the two prescribed points
    
    @param[in]      lapack   pointer to the Teuchos::LAPACK interface
    @param[in]      xl1      location of one pre-assigned node
    @param[in]      xl2      location of another pre-assigned node
    @param[in/out]  ap       pointer to vector of recursion coefficients
    @param[in/out]  bp       pointer to vector of recursion coefficients

    Based on the section 7 of the paper 
    "Some modified matrix eigenvalue problems" 
    by Gene Golub, SIAM Review Vol 15, No. 2, April 1973, pp.318--334 */
template<class Real>
void rec_lobatto( ROL::Ptr<Teuchos::LAPACK<int,Real> > const lapack,
                  const double xl1, 
                  const double xl2,
                  std::vector<Real> &a,
                  std::vector<Real> &b ) {
    const int N = a.size()-1;

    std::vector<Real> amod(N,0.0);
    std::vector<Real> bmod(N-1,0.0);
    std::vector<Real> en(N,0.0);
    std::vector<Real> g(N,0.0);

    // Nth canonical vector
    en[N-1] = 1;

    for(int i=0;i<N-1;++i) {
        bmod[i] = sqrt(b[i+1]);
    }

    for(int i=0;i<N;++i) {
        amod[i] = a[i]-xl1;
    }
    
    trisolve(lapack,bmod,amod,bmod,en,g);         
    Real g1 = g[N-1];

    for(int i=0;i<N;++i) {
        amod[i] = a[i]-xl2;
    }
    
    trisolve(lapack,bmod,amod,bmod,en,g);         
    Real g2 = g[N-1];

    a[N] = (g1*xl2-g2*xl1)/(g1-g2);
    b[N] = (xl2-xl1)/(g1-g2);
    
}

 
#endif

             
