#include "nr3.h"
#include "eigen_sym.h"
#include "gamma.h"
#include "gauss_wgts.h"
#include "well.h"
#include "harmonic.h"
#define NBASE 50
#define NGAUSSLEG 300
double WIDTH; // width for square well potential
double mass=1; // mass, I use 1 everywhere

double V(double x) {
   return x*x*(x*x-2.0);
}

int main() {
   cin >> WIDTH;
   cout << "Tested with a square well base of width : " << WIDTH << endl;
   cout.precision(8);
   MatDoub hamiltonian(NBASE,NBASE);
   well w(mass,WIDTH); // mass 1
   harmonic h(mass,2.0); // string constant 2, mass 1
   // test full procedure by calculating the eigenfunctions from the
   // hamiltonian base in the well base.
   // integrate from -WIDTH to +WIDTH for the well, calculate harmonic oscillator wave function with C=2
   VecDoub x(NGAUSSLEG),weight(NGAUSSLEG);
   gauleg(-WIDTH,WIDTH,x,weight);

   for (int i=0; i<NBASE; i++) for (int j=0;j < NBASE; j++) {
       hamiltonian[i][j]=0;
       for (int k=0; k<NGAUSSLEG; k++) hamiltonian[i][j]+=weight[k]*w.wav1d(i,x[k])*(-0.5/mass*w.secdiv(j)+V(x[k]))*w.wav1d(j,x[k]);
   }
   Symmeig wellsolution(hamiltonian);
   gauleg(-12,12,x,weight);
   for (int i=0; i<NBASE; i++) for (int j=0;j < NBASE; j++) {
      hamiltonian[i][j]=0;
      for (int k=0; k<NGAUSSLEG; k++) hamiltonian[i][j]+=weight[k]*h.wav1d(i,x[k])*(-0.5/mass*h.secdiv(j,x[k])+V(x[k]))*h.wav1d(j,x[k]);
   }
   Symmeig harmonicsolution(hamiltonian);
   for (int i=1; i<11; i++) {
	cout << "state " << i-1 << " E well : " << wellsolution.d[NBASE-i] << " E harmonic " << harmonicsolution.d[NBASE-i] << " diff" << (wellsolution.d[NBASE-i]-harmonicsolution.d[NBASE-i]) << endl;
   }
   cout << endl << endl;
   cout << " Energy, x, psi_well, psi_harmonic " << endl;
// print all wave functions, well and harmonic
   for (int i=1; i<11; i++) {
   for (double xx=-10;xx<10; xx+=0.01) {
        cout << harmonicsolution.d[NBASE-i] << " " << xx << " ";
        double psiwell=0;
        double psihar=0;
	for (int j=0; j<NBASE; j++) {
             psiwell+=wellsolution.z[j][NBASE-i]*w.wav1d(j,xx);
             psihar+=harmonicsolution.z[j][NBASE-i]*h.wav1d(j,xx);
        }
        cout << psiwell << " " << psihar << endl;
    }
    }
    return 0;
}