using Pkg
Pkg.add("BenchmarkTools")
Pkg.add("LinearAlgebra")
Pkg.add("Plots")

using LinearAlgebra
using BenchmarkTools
# R has a tictoc library or system time.

function solve_QHO(n, xL, xR)
    h = (xR-xL)/(n-2);  # step size
    x = range(xL, xR, step=h); # grid of x values

    # with for loops
    H = zeros(Float64, n-1, n-1); # intialize
    for i in range(1,n-1, step=1);
      @inbounds H[i,i] = 1/h^2 + .5*x[i]^2;  # diagonal of H
    end

    for i in range(1,n-2)
      @inbounds H[i+1,i] = -1/(2*h^2); # upper diag, one row above
      @inbounds H[i,i+1] = -1/(2*h^2); # lower diag
    end

    vals, vecs = eigen(H);
    
    return x, vals, vecs
end

x, vals, vecs = solve_QHO(1000, -5, 5);
vals[1]

# data for timing comparison
ttable = m = Array{Float64}(undef, 70, 4)
itr = 1;
for n in range(500,3500, step=500);
  for x in range(1,10);
    r = @belapsed solve_QHO($n, -$x, $x)
    ttable[itr,1] = -x
    ttable[itr,2] = x
    ttable[itr,3] = n
    ttable[itr,4] = r
    itr+=1
  end
end

r = @belapsed solve_QHO(1000, -5,5)
@btime solve_QHO(2000, -5,5)


#Pkg.add("Plots")
using Plots;
plot(x, [vecs[:,1], vecs[:,2], vecs[:,3]],
    title="wave functions",
    label=["E0", "E1", "E2"], linewidth=3, gridlinewidth=3)

function Laplace_time(;alpha, L, T, h, u_tblr)
  # ; lets you ref arguments by name
  # Laplace_time(alpha=20, L=25, T=)
  # u_tblr = [100, 0, 0, 0] # boundary conditions
  dt = h^2/(4*alpha+1);  # ensures numerical stability
  n = floor(Int64, L/h); # coerce to integer, spatial loops
  m = floor(Int64, T/dt); # number of time steps
  u = zeros(n,n,m) # intialize solution 3d array

  function impose_bc(u_tblr, u, n, k)
      u_top    = u_tblr[1];   u[1,:,k] .= u_top; # top bc at time k
      u_bottom = u_tblr[2];   u[n,:,k] .= u_bottom; # top bc at time k
      u_left   = u_tblr[3];   u[:,1,k] .= u_left; # top bc at time k
      u_right  = u_tblr[4];   u[:,n,k] .= u_right; # top bc at time k
      return u
  end
  u = impose_bc(u_tblr, u, n, 1)  # intial u at time k=1
  for k in range(2,m-1)         # time
    for i in range(2,n-1)     # x 
        for j in range(2,n-1) # y
          u[i,j,k+1] = (1-(4*alpha*dt)/(h^2))*u[i,j,k] + ((alpha*dt)/(h^2))*(u[i+1,j,k]+u[i-1,j,k]+u[i,j+1,k]+u[i,j-1,k]) # method of lines
        end
    end
    u = impose_bc(u_tblr, u, n, k) # re-impose the bc
  end
  return u
end

my_u = Laplace_time(alpha=20, L=25, T=20, h=1, u_tblr=[100, 0, 100, 0]);
my_u_b = Laplace_time(alpha=20, L=50, T=40, h=2, u_tblr=[100, 0, 100, 0]);
my_u_c = Laplace_time(alpha=20, L=100, T=60, h=5, u_tblr=[100, 0, 100, 0]);


my_dims = size(my_u)
using Plots;
heatmap(1:my_dims[1], 1:my_dims[2], my_u[:,:,1])

@gif for k = 1:100
  heatmap(1:my_dims[1], 1:my_dims[2], my_u[:,:,k],
      xlabel = join(["time idx = ", k])
  )
end