//
// Program: quantum_wall.C
//
// Description: Risoluzione dell'eq. di Schroedinger in una buca di
//              potenziale.
//
// Author: A. Garfagnini
//
// Date: 17 December 2013
//
#include <iostream>
#include <cmath>
#include <fstream>
#include <unistd.h>
#include <libgen.h>
#include <stdlib.h>

using namespace std;

void usage(const char * pname)
{
    cerr << "Usage: " << pname << endl;
    cerr << "Risolve l'equazione di Schroedninger in una buca di";
    cerr << "potenziale a pareti infinite.\n";
    cerr << "\nOpzioni:\n";
    cerr << "  -l num : calcola fino all'autostato numero \"num\"\n";
}

// Calcola la funzione d'ona relativa all'energia E e ritorna il
// valore in x = 1
double fond(double erg, double psi[], double dx, int n)
{
    for (int j=0; j<n; j++) {
        int k = j+1;
        psi[k+1] = -psi[k-1] + 2*psi[k] - 2*dx*dx*erg*psi[k];
    }
    return psi[n+1];
}

//
// Funzione di bisezione ricorsiva
//
double bisezione(double a, double b, double wav[], double d, int t)
{
    double m = (a+b)/2;
    // cout << "m = " << m << endl;
    if ( fond(m, wav, d, t) == 0 ||
         fabs(fond(m,wav,d,t)) < 0.01)   return m;
    if ( fond(m, wav, d, t) *  fond(m, wav, d, t) < 0.0) {
        return bisezione(a, m, wav, d, t);
    } else {
        return bisezione(b, m, wav, d, t);
    }
}

// Normalizza la funzione d'onda calcolata
void normalizzazione(int n, double func[], double dx)
{
    double C = 0.0;

    for (int i=1; i<n; i++) {
        C += (func[i]*func[i] + func[i+1]*func[i+1])*dx/2.0;
    }

    C *= 2.0;
    C = sqrt(C);

    for (int i=0; i<n; i++) {
        func[i] /= C;
    }
}

int main(int argc, char * argv[])
{
   char filename[128] = "psi_";
   int opt;
   int energy_level = 1;

    // Handle command line options
    while((opt = getopt(argc, argv, "l:h?")) != -1) {

        switch (opt) {

        case 'l':
            energy_level = strtol(optarg, NULL, NULL);
            cout << "Energy level: " << energy_level << endl;
            if (energy_level < 1) {
                cerr << "Wrong energy level: " << energy_level;
                cerr << " (energy_level > 0)\n";
                usage( basename( argv[0] ) );
                return 1;
            }
            break;

        case 'h':
        case '?':
        default:
            usage( basename( argv[0] ) );
            return 1;
        }
    }

    // Handle user data input
    cout << "Numero di punti della griglia: ";
    int n;
    cin >> n;

    double * psi = new double[n + 2];

    // Definisci la griglia
    double dx = 1.0/(n+1);
    double e = 0.1, de = 0.1;

    double energia;

    for (int j=0; j<energy_level; j++) {

        for (int i=0; i<n+2; i++) {
            psi[i] = 0.0;
        }

        // Imposta le condizioni iniziali
        if (j%2 == 0) {

            // Funzioni d'onda pari
            psi[0] = 1.0;
            psi[1] = 1.0;

        } else {

            // Funzioni d'onda dispari
            psi[0] = - dx;
            psi[1] = 0;
        }

        double f1, f2;
        do {
            f1 = fond(e,psi,dx,n);
            f2 = fond(e+de,psi,dx,n);
            //cout << "E = " << e << " | f1 = " << f1 << " f2 = " << f2 << endl;
            e += de;
        } while(f1 * f2 > 0.0);
        e -= de;

        cout << "Energy = " << e << endl;

        // cout << "dx = " << dx << endl;
        energia = bisezione(e, e+de, psi, dx, n);
        normalizzazione(n+2, psi, dx);

        // Write only the requested energy level wave function
        if (energy_level == j+1) {
            // Write in a file
            sprintf(filename, "psi_%d.dat", energy_level);
            ofstream myfile(filename);
            if (j%2==0) {

                for (int i=n+1; i>1; i--)
                    myfile << (-i + 1.0)/n << " " << psi[i] << endl;
                for (int i=1; i<n+2; i++)
                    myfile << (i - 1.0)/n << " " << psi[i] << endl;

            } else {

                for (int i=n+1; i>1; i--)
                    myfile << (-i + 1.0)/n << " " << -psi[i] << endl;
                for (int i=1; i<n+2; i++)
                    myfile << (i - 1.0)/n << " " << psi[i] << endl;

            }

            myfile.close();
        }

        e += de;


    }

    return 0;
}