# Variational calculation with truncated Hilbert space
# - numerical diagonalization of H in basis of N partocle-in-a-box states

 using LinearAlgebra
 using Plots

# Any potential within a box from x=x1 to x=x2 inside the box can be put here

 function potential(x,x1,x2,ll)
    v::Float64=20.0
    return v
 end 

# Particle-in-a-box wave function k

 function eigenfunc(x::Float64,k::Int,ll::Float64)::Float64
    f=sin(k*x*pi/ll)*(2/ll)^0.5
    return f
 end

# The function to be integrated for the potential-energy matrix elements

 function vpsikpsip(x,x1,x2,k,p,ll)::Float64
    psik=eigenfunc(x,k,ll)
    psip=eigenfunc(x,p,ll)
    return psik*psip*potential(x,x1,x2,ll)
 end

# Integrates the function func(x) from x=x2 to x=x2
# - uses n intervals dx (n+1 grid points)

 function simpson(func,x1::Float64,x2::Float64,n::Int)::Float64
    dx=(x2-x1)/n
    si::Float64=func(x1)+func(x2)
    for i=1:2:n-3
      x=x1+i*dx
      si=si+4*func(x)+2*func(x+dx)
    end 
    si=si+4*func(x2-dx)
    return si*dx/3
 end

# Makes the Hamiltonian matrix using Simpson's method

 function hamiltonian(n,x1,x2,ll)
    ham=Matrix{Float64}(undef,n,n)
    eps=1e-4
    for p=1:n
       for k=1:n
           ham[k,p] = simpson(x -> vpsikpsip(x,x1,x2,k,p,ll),x1,x2,1000)
       end
       ham[p,p] += (0.5*(pi*p/ll)^2)
    end
    return ham
 end

# Transforms from the eigenvector in the box basis to real space

 function psi(v,x)
    nx=length(x)
    ll=x[nx]
    p=zeros(Float64,nx)
    for j=1:nx
       for i=1:length(v)
          p[j] += v[i]*eigenfunc(x[j],i,ll)
       end
    end
    return p
 end

 n=5     # number of basis states
 ll=2.0  # box size (between x=0 and x=ll)
 x1=0.5  # lower edge of the potential inside the box
 x2=1.5  # upper edge of the potential inside the box

 h=hamiltonian(n,x1,x2,ll)

# Get the eigenvalues and eigenvectors
# - there is another function eigen() that does both, may be faster

 e=eigvals(h)
 v=eigvecs(h)
 
 nx=200                  # grid points (-1) for wave function
 dx=ll/nx                # grid spacing for wave function
 x=range(0.0,ll,length=nx+1)  # vector holding the grid points for the wave function 

 # Eigenvalue k is printed, egenfunction raphed (k in 1,...,n)             
 k=1         
 println(e[k],"  ",e[k+1])
 p=psi(v[:,k],x) 
 if p[2]<0 p=-p end    # overall sign choise for wave function
 
 display(plot(x,p))

 # now do it again with twice as many basis states
 n=n*2
 h=hamiltonian(n,x1,x2,ll)
 e=eigvals(h)
 v=eigvecs(h)
 println(e[k],"  ",e[k+1])
 p=psi(v[:,k],x) 
 if p[2]<0 p=-p end    # overall sign choise for wave function
 display(plot!(x,p))