using PyPlot, LinearAlgebra

A1 = randn(500,500)
A2 = randn(1000,1000)
lu(A1); # do it once to make sure it is compiled
@elapsed(lu(A2)) / @elapsed(lu(A1))

qr(A1); # do it once to make sure it is compiled
@elapsed(qr(A2)) / @elapsed(qr(A1))

eigen(A1); # do it once to make sure it is compiled
@elapsed(eigen(A2)) / @elapsed(eigen(A1))

@time lu(A2)
@time eigen(A2);

A = [  2 -1  0  0  0  0
      -1  2 -1  0  0  0
       0 -1  2 -1  0  0
       0  0 -1  2 -1  0
       0  0  0 -1  2 -1
       0  0  0  0 -1  2]

using SparseArrays

sparse(A)

# compute the first-derivative finite-difference matrix
# for Dirichlet boundaries, given a grid x[:] of x points
# (including the endpoints where the function = 0!).
function sdiff1(x)
    N = length(x) - 2
    dx1 = [1/(x[i+1] - x[i]) for i = 1:N]
    dx2 = [-1/(x[i+1] - x[i]) for i = 2:N+1]
    spdiagm(N+1, N, 0=>dx1, -1=>dx2)
end

# compute the -∇⋅ c ∇ operator for a function c(x,y)
# and arrays x[:] and y[:] of the x and y points,
# including the endpoints where functions are zero
# (i.e. Dirichlet boundary conditions).
function Laplacian(x, y, c = (x,y) -> 1.0)
    Dx = sdiff1(x)
    Nx = size(Dx,2)
    Dy = sdiff1(y)
    Ny = size(Dy,2)
    
    # discrete Gradient operator:
    G = [kron(sparse(I,Ny,Ny), Dx); kron(Dy, sparse(I,Nx,Nx))]
    
    # grids for derivatives in x and y directions
    x′ = [0.5*(x[i]+x[i+1]) for i = 1:length(x)-1]
    y′ = [0.5*(y[i]+y[i+1]) for i = 1:length(y)-1]
    
    # evaluate c(x)
    C = spdiagm([ vec([c(X,Y) for X in x′, Y in y[2:end-1]]);
                  vec([c(X,Y) for X in x[2:end-1], Y in y′]) ])
    
    return G' * C * G # -∇⋅ c ∇
end

N = 400
x = range(-1,1,length=N+2)[2:end-1]
y = x'  # a row vector
r = hypot.(x,y)  # use broadcasting (.+) to make Nx x Ny matrix of radii
θ = atan.(y, x) # and angles
φ = @. exp(-(r - θ*0.5/π - 0.5)^2 / 0.3^2) - 0.5
imshow(φ .> 0, extent=[-1,1,-1,1], cmap="binary")
title("the domain")

x0 = range(-1,1,length=N+2) # includes boundary points, unlike x
Abox = Laplacian(x0, x0, (x,y) -> 1.0);
i = findall(>(0), vec(φ))
A = Abox[i,i]
size(A)

A

nnz(A) / length(A)

nnz(A) / size(A,1)

import KrylovKit

# to find the smallest |λ| eigenvectors, we need to multiply repeatedly by A⁻¹,
# not A, and we do this by calling F \ x with a "Cholesky" factorization (a type of LU) F:
F = cholesky(A)
@time λ, X, = KrylovKit.eigsolve(x -> F \ x, size(A,1), 20, ishermitian=true);

# the frequencies are the square roots of the eigenvalues of A, but
# since we applied eigsolve to A⁻¹ we also need to take 1/λ
ω = sqrt.(1 ./ λ)

@show size(X)

f = figure()
u = zeros(N,N)
for which_eig in 1:20
    display(
    withfig(f) do
        u[i] = X[which_eig]
        umax = maximum(abs, u)
        title("vibrating mode $which_eig, ω=$(ω[which_eig])")
        imshow(u, extent=[-1,1,-1,1], vmin=-umax,vmax=+umax, cmap="RdBu")
    end
    )
end

"""
Return `(A,x,y)` for the 2d Helmholtz problem -∇²-ω²ε.
"""
function Helmholtz2d(Lx, Ly, ε, ω; dpml=2, resolution=20, Rpml=1e-20)
    # PML σ = σ₀ x²/dpml², with σ₀ chosen so that the round-trip reflection is Rpml
    σ₀ = -log(Rpml) / (4dpml/3)
    
    M = round(Int, (Lx+2dpml) * resolution)
    N = round(Int, (Ly+2dpml) * resolution)
    dx = (Lx+2dpml) / (M+1)
    dy = (Ly+2dpml) / (N+1)
    x = (1:M) * dx # x grid
    y = (1:N) * dy # y grid
    x′ = @. ((0:M) + 0.5)*dx # 1st-derivative grid points
    y′ = @. ((0:N) + 0.5)*dy
    
    # 1st-derivative matrices
    ox = ones(M)/dx
    oy = ones(N)/dy
    Dx = spdiagm(M+1,M, -1 => -ox, 0 => ox)
    Dy = spdiagm(N+1,N, -1 => -oy, 0 => oy)
    
    # PML complex "stretch" factors 1/(1+iσ/ω) at both x and x' points:
    σx = [ξ < dpml ? σ₀*(dpml-ξ)^2 : ξ > Lx+dpml ? σ₀*(ξ-(Lx+dpml))^2 : 0.0 for ξ in x]
    Σx = spdiagm(@. inv(1 + (im/ω)*σx))
    σx′ = [ξ < dpml ? σ₀*(dpml-ξ)^2 : ξ > Lx+dpml ? σ₀*(ξ-(Lx+dpml))^2 : 0.0 for ξ in x′]
    Σx′ = spdiagm(@. inv(1 + (im/ω)*σx′))
    # similarly for y and y':
    σy = [ξ < dpml ? σ₀*(dpml-ξ)^2 : ξ > Ly+dpml ? σ₀*(ξ-(Ly+dpml))^2 : 0.0 for ξ in y]
    Σy = spdiagm(@. inv(1 + (im/ω)*σy))
    σy′ = [ξ < dpml ? σ₀*(dpml-ξ)^2 : ξ > Ly+dpml ? σ₀*(ξ-(Ly+dpml))^2 : 0.0 for ξ in y′]
    Σy′ = spdiagm(@. inv(1 + (im/ω)*σy′))
    
    # stretched 2nd-derivative matrices
    D2x = Σx * Dx' * Σx′ * Dx
    D2y = Σy * Dy' * Σy′ * Dy
    
    # combine x and y with Kronecker products
    Ix = spdiagm(ones(M))
    Iy = spdiagm(ones(N))
    x = x .- dpml
    y = y .- dpml
    return (kron(Ix, D2y) + kron(D2x, Iy) -
            spdiagm(vec([ω^2 * ε(ξ, ζ) for ζ in y, ξ in x])),
            x, y)
end

A, x, y = Helmholtz2d(20,20, (x,y) -> hypot(x,y) < 0.5 ? 12 : 1, 2π)
size(A), nnz(A) / size(A,1)

nnz(A)/size(A,1)

A

F = lu(A)
F.U

nnz(F.U) / size(A, 1)

Nx, Ny = length(x), length(y)
b = zeros(Nx,Ny)
b[Nx÷2, Ny÷4] = 1
@time u = reshape(A \ vec(b), Nx, Ny)
s = maximum(abs, real(u)) / 10
imshow(real(u), cmap="RdBu", vmin=-s, vmax=s,
       extent=(minimum(x),maximum(x),minimum(y),maximum(y)))

function SD(A, b, x=zero(b); tol=1e-8, maxiters=1000)
    bnorm = norm(b)
    r = b - A*x # initial residual
    rnorm = [norm(r)] # return the array of residual norms
    Ad = zero(r) # allocate space for Ad
    for i = 1:maxiters
        d = r # use the steepest-descent direction
        mul!(Ad, A, d) # store matvec A*r in-place in Ar
        α = dot(d, r) / dot(d, Ad)
        x .= x .+ α .* d  # in Julia, this "fuses" into a single in-place update
        r .= r .- α .* Ad # update the residual (without computing A*x again)
        push!(rnorm, norm(r))
        rnorm[end] ≤ tol*bnorm && break # converged
    end
    return x, rnorm
end

A = rand(100,100); A = A'*A # a random SPD matrix
b = rand(100)
x, rnorm = SD(A, b, maxiters=1000)
length(rnorm), rnorm[end]/norm(b)

semilogy(rnorm)
title("steepest-descent convergence")
ylabel("residual norm")
xlabel("iterations")

θ = 0.9 # chosen to make a nice-looking plot
Q = [cos(θ) sin(θ); -sin(θ) cos(θ)] # 2x2 rotation by θ
A = Q * diagm([10,1]) * Q' # a 2x2 matrix with eigenvalues 10,1
b = A * [1,1] # right-hand side for solution (1,1)
x1 = range(-11,11,length=100)
contour(x1', x1, [dot([x1,x2], A*[x1,x2]) - 2*(x1*b[1]+x2*b[2]) for x1 in x1, x2 in x1], levels=range(1,2000,length=100))
plot(1,1, "r*")
x1s = Float64[]
x2s = Float64[]
for i = 0:20
    x, = SD(A, b, [-10.,-10.], maxiters=i)
    push!(x1s, x[1])
    push!(x2s, x[2])
end
plot(x2s, x1s, "k.-")

title("convergence of 2x2 steepest-descent")

function CG(A, b, x=zero(b); tol=1e-8, maxiters=1000)
    bnorm = norm(b)
    r = b - A*x # initial residual
    rnorm = [norm(r)] # return the array of residual norms
    d = copy(r) # initial direction is just steepest-descent
    Ad = zero(r) # allocate space for Ad
    for i = 1:maxiters
        mul!(Ad, A, d) # store matvec A*r in-place in Ar
        α = dot(d, r) / dot(d, Ad)
        x .= x .+ α .* d  # in Julia 0.6, this "fuses" into a single in-place update
        r .= r .- α .* Ad # update the residual (without computing A*x again)
        push!(rnorm, norm(r))
        d .= r .+ d .* (rnorm[end]/rnorm[end-1])^2 # conjugate direction update
        rnorm[end] ≤ tol*bnorm && break # converged
    end
    return x, rnorm
end

A = rand(100,100); A = A'*A # a random SPD matrix
b = rand(100)
x, rnorm = CG(A, b)
length(rnorm), rnorm[end]/norm(b)

semilogy(rnorm)
title("conjugate-gradient convergence")
ylabel("residual norm")
xlabel("iterations")

θ = 0.9 # chosen to make a nice-looking plot
Q = [cos(θ) sin(θ); -sin(θ) cos(θ)] # 2x2 rotation by θ
A = Q * diagm([10,1]) * Q' # a 2x2 matrix with eigenvalues 10,1
b = A * [1,1] # right-hand side for solution (1,1)
x1 = range(-11,11,length=100)
contour(x1', x1, [dot([x1,x2], A*[x1,x2]) - 2*(x1*b[1]+x2*b[2]) for x1 in x1, x2 in x1], levels=range(1,2000,length=100))
plot(1,1, "r*")
x1s = Float64[]
x2s = Float64[]
for i = 0:2
    x, = CG(A, b, [-10.,-10.], maxiters=i)
    push!(x1s, x[1])
    push!(x2s, x[2])
end
plot(x2s, x1s, "k.-")

title("convergence of 2x2 conjugate-gradient")

import IterativeSolvers

A = rand(100,100); A = A'*A # a random SPD matrix
b = rand(100)
x, ch = IterativeSolvers.cg(A, b, maxiter=300, log=true)
norm(A*x - b) / norm(b)

semilogy(ch[:resnorm])
xlabel("iteration")
ylabel(L"\Vert \mathrm{residual} \Vert")
title("convergence of conjugate gradient")

n = 300
L = SymTridiagonal(fill(2.0,n), fill(-1.0, n-1))
b = rand(n)

S = sprand(n,n,0.001) * 0.1  # a small random, sparse perturbation
S = S'*S # needs to be symmetric positive-definite

A = sparse(L + S)

x, ch = IterativeSolvers.cg(A, b, maxiter=300, log=true)
x′, ch′ = IterativeSolvers.cg(A, b, Pl=ldlt(L), maxiter=300, log=true)
semilogy(ch[:resnorm], "b-")
semilogy(ch′[:resnorm], "r-")
xlabel("iterations")
ylabel("residual norm")
legend(["no preconditioner", "preconditioner"], loc="lower right")

semilogy(sort!(eigvals(Matrix(A))), "b.")
semilogy(sort!(eigvals(Matrix(A),Matrix(L))), "r.")
xlabel(L"index $k$ of eigenvalue $\lambda_k$")
ylabel(L"eigenvalues $\lambda_k$")
legend([L"unpreconditioned $A$", L"preconditioned $L^{-1} A$"])