#!/usr/bin/env python3
"""Chapter 10, code 5: Schrodinger equation as a matrix eigenvalue problem.

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, with a stable small-angle rotation formula.

Example: harmonic oscillator, V(x) = 0.5*x^2.
Exact energies on the infinite line: E_j = j+0.5, j = 0,1,2,... .
A finite interval approximates the infinite line for the lowest states.
Check both grid spacing and interval size when testing accuracy.

Requires NumPy and Matplotlib. Run: python3 ch10_schrodinger_matrix.py
Alternatives: numpy.linalg.eigh solves a real symmetric matrix problem.
SciPy's eigh_tridiagonal uses just the diagonal and off-diagonal arrays;
eigsh can compute selected eigenpairs of a large sparse Hamiltonian.
Keep matching eigenvector columns when sorting energies.

To solve another problem:
  1. Edit the interval, grid size, and number of states in main.
  2. Edit potential immediately below the driver.
  3. Disable check_exact or replace the exact-energy comparison.
Jacobi computes all eigenpairs of a dense matrix. Keep grids modest.
"""

import numpy as np
import matplotlib.pyplot as plt


def main():
    # ------------------- Problem and grid: edit here ------------------
    a = -6.0
    b = 6.0
    n = 120                     # n intervals, n-1 interior unknowns
    states = 3
    abserr = 1.0e-18            # Bound on SUM of squared off-diagonal entries
    max_sweeps = 100
    check_exact = True          # Harmonic-oscillator energy comparison
    make_plot = True
    # -----------------------------------------------------------------
    interior = n-1
    if (n < 2 or n != int(n) or states < 1 or states > interior
            or states != int(states) or b <= a):
        raise ValueError('Use integer n >= 2, integer 1 <= states <= n-1, and b > a.')
    h = (b-a)/n
    x = np.linspace(a, b, n+1)
    v = np.zeros(n+1)
    for i in range(n+1):
        v[i] = potential(x[i])

    # Only interior values are unknown. Zero endpoints contribute nothing
    # to the first and last interior equations. Interior matrix index i
    # corresponds to full-grid point x[i+1].
    matrix = np.zeros((interior, interior))
    for i in range(interior):
        matrix[i,i] = 1/(h*h)+v[i+1]
    for i in range(interior-1):
        matrix[i,i+1] = -0.5/(h*h)
        matrix[i+1,i] = matrix[i,i+1]

    print('Schrodinger equation: finite difference matrix and Jacobi eigenvalues')
    print(f'Grid intervals: {n:8d}')
    print(f'Matrix size:    {interior:8d}')
    print(f'Grid spacing:   {h:12.6f}')
    matrix, vectors, sweeps, rotations, offdiag_norm, info = jacobi(
        matrix, abserr, max_sweeps)
    if info != 0:
        raise RuntimeError('Jacobi failed: check controls or increase sweep limit.')
    energies = np.diag(matrix).copy()
    energies, vectors = sort_eigenpairs(energies, vectors)

    # Matrix eigenvectors have sum(vectors**2) = 1. Wavefunctions need
    # h*sum(psi**2) = 1, also the trapezoidal integral with zero endpoints.
    psi = np.zeros((n+1, states))
    for j in range(states):
        norm_value = np.sqrt(h*np.sum(vectors[:,j]**2))
        psi[1:n,j] = vectors[:,j]/norm_value
        # Fix the arbitrary sign: largest-magnitude component positive.
        pivot = np.argmax(np.abs(psi[:,j]))
        if psi[pivot,j] < 0:
            psi[:,j] = -psi[:,j]

    print()
    if check_exact:
        print(f"{'state j':>7} {'numerical E':>16} {'exact E':>16} {'abs error':>12}")
    else:
        print(f"{'state j':>7} {'numerical E':>16}")
    for j in range(states):
        if check_exact:
            exact = j+0.5
            print(f'{j:7d} {energies[j]:16.8f} {exact:16.8f} {abs(energies[j]-exact):12.4e}')
        else:
            print(f'{j:7d} {energies[j]:16.8f}')

    # Check with the original tridiagonal operator, not the matrix
    # overwritten by Jacobi. These checks concern matrix eigenpairs;
    # they do not measure grid or finite-interval errors.
    max_residual = 0.0
    for j in range(states):
        for i in range(1, n):
            residual = ((1/(h*h)+v[i]-energies[j])*psi[i,j]
                        -0.5*(psi[i-1,j]+psi[i+1,j])/(h*h))
            max_residual = max(max_residual, abs(residual))
    max_orthogonality = 0.0
    for j in range(states):
        for k in range(states):
            overlap = h*np.sum(psi[:,j]*psi[:,k])
            target = 1.0 if j == k else 0.0
            max_orthogonality = max(max_orthogonality, abs(overlap-target))
    print(f'Jacobi sweeps:            {sweeps:8d}')
    print(f'Largest matrix residual:  {max_residual:12.4e}')
    print(f'Largest overlap defect:   {max_orthogonality:12.4e}')

    filename = 'ch10_schrodinger_matrix.dat'
    with open(filename, 'w', encoding='utf-8') as data_file:
        data_file.write('# x V(x) psi_0 psi_1 ...; integral(psi_j^2 dx) = 1\n')
        for j in range(states):
            data_file.write(f'# state {j:6d} E = {energies[j]:.16e}\n')
        print('\nWavefunctions: columns are x, V(x), then the normalized states.')
        for i in range(n+1):
            print(f'{x[i]:10.6f} {v[i]:16.8e}', end='')
            for j in range(states):
                print(f' {psi[i,j]:16.8e}', end='')
            print()
            row = [x[i], v[i], *psi[i]]
            data_file.write(' '.join(f'{value:.16e}' for value in row)+'\n')
    print(f'Data file: {filename}')

    if make_plot:
        fig, axes = plt.subplots(states, 2, figsize=(10, 2.5*states), squeeze=False)
        for j in range(states):
            ax = axes[j, 0]
            ax.plot(x, psi[:,j], 'k-', linewidth=1.2)
            ax.set_xlabel('x')
            ax.set_ylabel(r'$\psi_j(x)$')
            ax.set_title(f'State {j}: E = {energies[j]:.8f}')
            ax.grid(True)
            ax = axes[j, 1]
            ax.plot(x, psi[:,j]**2, 'k-', linewidth=1.2)
            ax.set_xlabel('x')
            ax.set_ylabel(r'$|\psi_j(x)|^2$')
            ax.set_title(f'Probability density, state {j}')
            ax.grid(True)
        fig.tight_layout()
        plt.show()


def potential(x):
    """Supply V in -0.5*psi'' + V*psi = E*psi.

    Default: dimensionless harmonic oscillator.
    Change this expression for another potential and disable check_exact.
    """
    return 0.5*x*x


def sort_eigenpairs(values, vectors):
    """Sort energies in ascending order, moving eigenvector columns too."""
    values = values.copy()
    vectors = vectors.copy()
    n = len(values)
    for i in range(n-1):
        smallest = i
        for j in range(i+1, n):
            if values[j] < values[smallest]:
                smallest = j
        if smallest == i:
            continue
        values[i], values[smallest] = values[smallest], values[i]
        work = vectors[:,i].copy()
        vectors[:,i] = vectors[:,smallest]
        vectors[:,smallest] = work
    return values, vectors


def jacobi(a, abserr, max_sweeps):
    """Thresholded Jacobi rotations for a real symmetric matrix.

    A copy of a is overwritten by a nearly diagonal similar matrix;
    its diagonal contains eigenvalues. vectors contains eigenvector columns.
    info: 0 converged, 1 invalid input, 2 sweep limit.
    abserr bounds the SUM of squared off-diagonal entries.

    Rotate rows and columns by the same angle and rotate eigenvector
    columns. Recompute the actual off-diagonal sum after each sweep.
    The stable tangent formula avoids cancellation for unequal diagonals.
    Temporary arrays preserve old values during each row or column rotation.
    """
    a = a.copy()
    n = a.shape[0]
    vectors = np.eye(n)
    sweeps = 0
    rotations = 0
    offdiag_norm = np.inf
    info = 1
    if (n < 1 or a.shape != (n, n) or abserr <= 0
            or max_sweeps < 1 or max_sweeps != int(max_sweeps)):
        return a, vectors, sweeps, rotations, offdiag_norm, info
    if not np.array_equal(a, a.T) or not np.all(np.isfinite(a)):
        return a, vectors, sweeps, rotations, offdiag_norm, info

    while True:
        b2 = 0.0
        for i in range(n):
            for j in range(n):
                if i != j:
                    b2 += a[i,j]**2
        offdiag_norm = np.sqrt(b2)
        if b2 <= abserr:
            info = 0
            return a, vectors, sweeps, rotations, offdiag_norm, info
        if sweeps >= max_sweeps:
            info = 2
            return a, vectors, sweeps, rotations, offdiag_norm, info
        bar = 0.5*b2/(n*n)
        for i in range(n-1):
            for j in range(i+1, n):
                if a[j,i]**2 <= bar:
                    continue
                b2 = max(0.0, b2-2*(a[j,i]**2))
                bar = 0.5*b2/(n*n)
                beta = (a[j,j]-a[i,i])/(2*a[j,i])
                # copysign also handles beta = 0 as Fortran sign(1,beta).
                tangent = -np.copysign(1.0, beta)/(abs(beta)+np.sqrt(1+beta**2))
                c = 1/np.sqrt(1+tangent**2)
                s = tangent*c

                # Rotate rows i and j.
                row_i = c*a[i,:]+s*a[j,:]
                row_j = -s*a[i,:]+c*a[j,:]
                a[i,:] = row_i
                a[j,:] = row_j

                # Rotate columns of the matrix and eigenvectors.
                column_i = c*a[:,i]+s*a[:,j]
                column_j = -s*a[:,i]+c*a[:,j]
                a[:,i] = column_i
                a[:,j] = column_j
                column_i = c*vectors[:,i]+s*vectors[:,j]
                column_j = -s*vectors[:,i]+c*vectors[:,j]
                vectors[:,i] = column_i
                vectors[:,j] = column_j
                rotations += 1
        sweeps += 1


if __name__ == '__main__':
    main()