#include <iostream>
#include <iosfwd>
#include <fstream>
#include <math.h>
#include <dmtk/vector.h>
#include <dmtk/matrix.h>
#include <LAPACK/dsygv.h>

int main()
{
	dmtk::Matrix<double> H;
	dmtk::Matrix<double> S;
	dmtk::Vector<double> c;
	
	int dim;
	std::cout << "Input basis size: " << std::endl;
	std::cin >> dim; 
	
	H.resize(dim,dim);
	S.resize(dim,dim);
	c.resize(dim);
	
	// calculate matrix elements
	
	for(int i = 0; i < dim; i++){
		for(int j = i; j < dim; j++){
			if((i+j)%2 == 0){ // think of another *smarter* way
				H(i,j) = -8.*double(1-i-j-2*i*j)/double((i+j+3)*(i+j+1)*(i+j-1));
				S(i,j) = 2./double(i+j+5)-4./double(i+j+3)+2./double(i+j+1);
				H(j,i) = H(i,j);
				S(j,i) = S(i,j);
			}
		}
	}
	
	// solve the generalized eigenvalue problem

	dmtk::Vector<double> w(dim); //eigenvalues
	dmtk::Vector<double> aux(3*dim); // auxiliary "work" array for lapack
	int info;
	dsygv_(1, 'V', 'U', 
		   H.rows(), H.array(), H.rows(), 
		   S.array(), S.rows(), 
		   w.array(), aux.array(), 3*dim, 
		   info);
	
	
    // output the results

    // energies
    
    for(int i = 0; i < dim; i++){
    	std::cout << i << " " << w(i) << std::endl;
    }
    
    // eigenfunctions
    
    for(int i = 0; i < dim; i++){
        char file_name[255];
        std::snprintf(file_name, 255, "state_%d.dat", i);
        std::ofstream output_file(file_name,std::ios::out);
        if(!output_file) {
          std::cout << "*** ERROR: Could not open file " << file_name << std::endl;
          exit(-1);
        }
    	

        for(double x=-1.; x < 1; x+=0.01){
        	double psi = 0.;
        	for(int j = 0; j < dim; j++){
                        c(j) = H(i,j); // Fortran notation H(col,row)
        		psi += c(j)*pow(x,j)*(x-1)*(x+1);
        	}
        	output_file << x << " " << psi << std::endl;
        }
        
        output_file.close();
    }
    
    return 0;
}