//=======================================================================
// Chapter 10, code 5: Schrodinger equation as a matrix eigenvalue problem.
// Translation of ch10_schrodinger_matrix.f90.
// Dimensionless units: hbar = mass = 1.
// -0.5*psi'' + V(x)*psi = E*psi, with psi(a) = psi(b) = 0.
// Method: centered finite differences and symmetric Jacobi rotations.
// Jacobi follows Chapter 8, using a stable small-angle rotation formula.
//
// Example: harmonic oscillator, V(x) = 0.5*x*x.
// Exact infinite-line energies: E_j = j+0.5, j = 0,1,2,... .
// The finite interval approximates the infinite line for the lowest states.
// Check both grid spacing and interval size when testing accuracy.
// To solve another problem:
//   1. Edit the interval, grid, and number of states in main.
//   2. Edit potential(x) immediately below the driver.
//   3. Disable check_exact or replace the exact-energy comparison.
// All matrix eigenpairs are computed; only the lowest states are output.
// Jacobi uses a dense matrix: keep grids modest for this educational code.
// Compile: g++ -std=c++11 -O2 ch10_schrodinger_matrix.cpp -o ch10_schrodinger_matrix
//=======================================================================
#include <algorithm>
#include <cmath>
#include <fstream>
#include <iomanip>
#include <iostream>
#include <limits>
#include <vector>

using Matrix = std::vector<std::vector<double>>;
double potential(double x);
void sort_eigenpairs(int n, std::vector<double>& values, Matrix& vectors);
void jacobi(Matrix& a, Matrix& x, int n, double abserr, int max_sweeps,
            int& sweeps, int& rotations, double& offdiag_norm, int& info);

int main()
{
    //-------------------- Problem and grid: edit here -------------------
    const int n = 120;     // Intervals: n+1 grid points, n-1 interior unknowns
    const int states = 3;
    double a = -6.0;
    double b = 6.0;
    double abserr = 1.0e-18; // Bound on SUM of squared off-diagonal entries
    int max_sweeps = 100;
    bool check_exact = true; // Harmonic-oscillator energy comparison
    //-------------------------------------------------------------------
    const int interior = n-1;
    if (n < 2 || states < 1 || states > interior || b <= a) {
        std::cerr << "Use n >= 2, 1 <= states <= n-1, and b > a.\n";
        return 1;
    }
    double h = (b-a)/n;
    std::vector<double> x(n+1), v(n+1), energies(interior);
    Matrix matrix(interior,std::vector<double>(interior,0.0));
    Matrix vectors(interior,std::vector<double>(interior));
    Matrix psi(n+1,std::vector<double>(states,0.0));
    for (int i = 0; i <= n; ++i) {
        x[i] = a+i*h;
        v[i] = potential(x[i]);
    }
    x[n] = b;
    v[n] = potential(b);

    // Only interior values are unknown. Zero endpoint values contribute
    // nothing to the first and last interior equations.
    // H(i,i) = 1/h^2 + V(x_i); H(i,i+1) = H(i+1,i) = -1/(2*h^2).
    // Matrix index i corresponds to physical grid point i+1.
    for (int i = 0; i < interior; ++i) matrix[i][i] = 1.0/(h*h)+v[i+1];
    for (int i = 0; i < interior-1; ++i) {
        matrix[i][i+1] = -0.5/(h*h);
        matrix[i+1][i] = matrix[i][i+1];
    }
    int sweeps = 0, rotations = 0, info;
    double offdiag_norm;
    std::cout << std::uppercase
              << "Schrodinger equation: finite difference matrix and Jacobi eigenvalues\n"
              << "Grid intervals: " << std::setw(8) << n << '\n'
              << "Matrix size:    " << std::setw(8) << interior << '\n'
              << "Grid spacing:   " << std::fixed << std::setprecision(6)
              << std::setw(12) << h << '\n';
    jacobi(matrix,vectors,interior,abserr,max_sweeps,sweeps,rotations,offdiag_norm,info);
    if (info != 0) {
        std::cerr << "Jacobi failed: check controls or increase the sweep limit.\n";
        return 1;
    }
    for (int i = 0; i < interior; ++i) energies[i] = matrix[i][i];
    sort_eigenpairs(interior,energies,vectors);

    // Matrix vectors have sum(vectors^2) = 1. Wavefunctions require
    // h*sum(psi^2) = 1. With zero endpoints, this is the trapezoidal norm.
    for (int j = 0; j < states; ++j) {
        double sum = 0.0;
        for (int i = 0; i < interior; ++i) sum += vectors[i][j]*vectors[i][j];
        double norm = std::sqrt(h*sum);
        for (int i = 1; i < n; ++i) psi[i][j] = vectors[i-1][j]/norm;
        // Fix the arbitrary overall sign: largest-magnitude component positive.
        int pivot = 0;
        for (int i = 1; i <= n; ++i)
            if (std::abs(psi[i][j]) > std::abs(psi[pivot][j])) pivot = i;
        if (psi[pivot][j] < 0.0)
            for (int i = 0; i <= n; ++i) psi[i][j] = -psi[i][j];
    }
    std::cout << '\n' << std::setw(7) << "state j"
              << ' ' << std::setw(16) << "numerical E";
    if (check_exact)
        std::cout << ' ' << std::setw(16) << "exact E"
                  << ' ' << std::setw(12) << "abs error";
    std::cout << '\n';
    for (int j = 0; j < states; ++j) {
        std::cout << std::setw(7) << j << std::fixed << std::setprecision(8)
                  << ' ' << std::setw(16) << energies[j];
        if (check_exact) {
            double exact = j+0.5;
            std::cout << ' ' << std::setw(16) << exact
                      << std::scientific << std::setprecision(4)
                      << ' ' << std::setw(12) << std::abs(energies[j]-exact);
        }
        std::cout << '\n';
    }
    // Check selected eigenpairs using the original tridiagonal operator,
    // rather than the matrix overwritten by Jacobi.
    double max_residual = 0.0;
    for (int j = 0; j < states; ++j) {
        for (int i = 1; i < n; ++i) {
            double residual = (1.0/(h*h)+v[i]-energies[j])*psi[i][j]
                              -0.5*(psi[i-1][j]+psi[i+1][j])/(h*h);
            max_residual = std::max(max_residual,std::abs(residual));
        }
    }
    double max_orthogonality = 0.0;
    for (int j = 0; j < states; ++j) {
        for (int k = 0; k < states; ++k) {
            double sum = 0.0;
            for (int i = 0; i <= n; ++i) sum += psi[i][j]*psi[i][k];
            double overlap = h*sum;
            double defect = (j == k) ? 1.0 : 0.0;
            max_orthogonality = std::max(max_orthogonality,std::abs(overlap-defect));
        }
    }
    std::cout << "Jacobi sweeps:            " << std::setw(8) << sweeps << '\n'
              << "Largest matrix residual: " << std::scientific << std::setprecision(4)
              << std::setw(12) << max_residual << '\n'
              << "Largest overlap defect:  " << std::setw(12) << max_orthogonality << '\n';

    // One file: x, potential, psi_0, psi_1, ... . Comments record energies.
    std::ofstream data("ch10_schrodinger_matrix.dat");
    if (!data) {
        std::cerr << "Cannot open ch10_schrodinger_matrix.dat.\n";
        return 1;
    }
    data << std::uppercase << std::scientific << std::setprecision(16)
         << "# x V(x) psi_0 psi_1 ...; integral(psi_j**2 dx) = 1\n";
    for (int j = 0; j < states; ++j)
        data << "# state " << std::setw(6) << j << " E = " << std::setw(24) << energies[j] << '\n';
    std::cout << "\nWavefunctions: columns are x, V(x), then the normalized states.\n";
    for (int i = 0; i <= n; ++i) {
        std::cout << std::fixed << std::setprecision(6) << std::setw(10) << x[i]
                  << std::scientific << std::setprecision(8)
                  << ' ' << std::setw(16) << v[i];
        data << ' ' << std::setw(24) << x[i] << ' ' << std::setw(24) << v[i];
        for (int j = 0; j < states; ++j) {
            std::cout << ' ' << std::setw(16) << psi[i][j];
            data << ' ' << std::setw(24) << psi[i][j];
        }
        std::cout << '\n';
        data << '\n';
    }
    data.close();
    std::cout << "Data file: ch10_schrodinger_matrix.dat\n";
    return 0;
}


double potential(double x)
{
    // Supply V(x) in -0.5*psi'' + V(x)*psi = E*psi.
    // Default: dimensionless harmonic oscillator.
    // For another potential, change this expression and disable check_exact.
    // Zero endpoint conditions are imposed by the matrix construction.
    return 0.5*x*x;
}


void sort_eigenpairs(int n, std::vector<double>& values, Matrix& vectors)
{
    // Sort energies in ascending order, moving their eigenvector columns too.
    for (int i = 0; i < n-1; ++i) {
        int smallest = i;
        for (int j = i+1; j < n; ++j)
            if (values[j] < values[smallest]) smallest = j;
        if (smallest == i) continue;
        std::swap(values[i],values[smallest]);
        for (int k = 0; k < n; ++k) std::swap(vectors[k][i],vectors[k][smallest]);
    }
}


void jacobi(Matrix& a, Matrix& x, int n, double abserr, int max_sweeps,
            int& sweeps, int& rotations, double& offdiag_norm, int& info)
{
    // Thresholded Jacobi rotations for a real symmetric matrix.
    // a is overwritten; its final diagonal contains the eigenvalues.
    // Columns of x are the associated orthonormal eigenvectors.
    // abserr bounds the SUM of squared off-diagonal entries.
    // info = 0: converged; 1: invalid input; 2: sweep limit reached.
    // The stable tangent formula avoids cancellation when diagonal entries
    // differ greatly. The input matrix must be exactly symmetric.
    for (int i = 0; i < n; ++i) {
        std::fill(x[i].begin(),x[i].end(),0.0);
        x[i][i] = 1.0;
    }
    sweeps = 0;
    rotations = 0;
    offdiag_norm = std::numeric_limits<double>::max();
    info = 1;
    if (n < 1 || abserr <= 0.0 || max_sweeps < 1) return;
    for (int i = 0; i < n; ++i)
        for (int j = i+1; j < n; ++j)
            if (a[i][j] != a[j][i]) return;

    while (true) {
        // Recompute the actual off-diagonal sum at each sweep.
        double b2 = 0.0;
        for (int i = 0; i < n; ++i)
            for (int j = 0; j < n; ++j)
                if (i != j) b2 += a[i][j]*a[i][j];
        offdiag_norm = std::sqrt(b2);
        if (b2 <= abserr) {
            info = 0;
            return;
        }
        if (sweeps >= max_sweeps) {
            info = 2;
            return;
        }
        double bar = 0.5*b2/(double(n)*n);
        for (int i = 0; i < n-1; ++i) {
            for (int j = i+1; j < n; ++j) {
                if (a[j][i]*a[j][i] <= bar) continue;
                b2 = std::max(0.0,b2-2.0*(a[j][i]*a[j][i]));
                bar = 0.5*b2/(double(n)*n);
                // Small rotation eliminating a[j][i]. The minus sign follows
                // the row rotation [c s; -s c] below.
                double beta = (a[j][j]-a[i][i])/(2.0*a[j][i]);
                double tangent = -std::copysign(1.0,beta)
                                 /(std::abs(beta)+std::sqrt(1.0+beta*beta));
                double c = 1.0/std::sqrt(1.0+tangent*tangent);
                double s = tangent*c;
                // Rotate rows i and j.
                for (int k = 0; k < n; ++k) {
                    double cs = c*a[i][k]+s*a[j][k];
                    double sc = -s*a[i][k]+c*a[j][k];
                    a[i][k] = cs;
                    a[j][k] = sc;
                }
                // Rotate columns i and j. Apply the same column rotation to
                // x so that its columns become the eigenvectors.
                for (int k = 0; k < n; ++k) {
                    double cs = c*a[k][i]+s*a[k][j];
                    double sc = -s*a[k][i]+c*a[k][j];
                    a[k][i] = cs;
                    a[k][j] = sc;
                    cs = c*x[k][i]+s*x[k][j];
                    sc = -s*x[k][i]+c*x[k][j];
                    x[k][i] = cs;
                    x[k][j] = sc;
                }
                ++rotations;
            }
        }
        ++sweeps;
    }
}