### Lecture 25 - efficient and inefficient gmres
using LinearAlgebra
"""
x = gmres(A,b,iters)
"""
function gmres(A,b,k)

n = length(b)
normb = norm(b)
q = b/norm(b)

Q = zeros(n,k+1)
H = zeros(k+1,k)
Q[:,1] = q
hist = zeros(k)
x = zeros(n)
for i = 1:k
    v = A*Q[:,i]
    for j = 1:i
        H[j,i]= v'*Q[:,j]
        v = v - H[j,i]*Q[:,j]
    end
    # at this point, v = H(i+1,i)*Q(:,i+1)
    H[i+1,i] = norm(v)
    Q[:,i+1] = v/H[i+1,i]

    ei = zeros(i+1)
    ei[1] = normb

    y = H[1:i+1,1:i]\ei
    x = Q[:,1:i]*y;
    hist[i] = norm(b-A*x)


    if H[i+1,i] < 100*eps()
        # this means we've finished the factorization early
        break
    end
end
return x,hist
end

n = 10
A = randn(n,n)
b = randn(n)

x1,hist1 = gmres(A,b,5)

([-0.352201745217469, 0.7081033750025103, -1.3393077353212062, -1.4790549734045693, 0.3542855813109064, -0.4687108554804207, 0.2869756722661793, -0.09066849919503164, -0.0773768290310974, -0.4103050804215759], [3.364936819267049, 2.9065236423847653, 2.755799449525736, 0.9934746477107153, 0.9100369517236918])

##

##
function planerot(x)
  a,b = -x[2],x[1]
  d = sqrt(b^2+a^2)
  a,b = a/d,b/d
  G = [b -a;a b]
  y = G*x
  return G,y
end

"""
An "efficient" implementation of GMRES that uses the Givens rotations
to update the solution of the least squares problem.
(Note that there are more efficiencies possible here, but this
  gets at the main ideas)
"""
function gmres_efficient(A,b,tol,maxit)
  n = size(b)
  beta = norm(b)

  Q = zeros(length(b),maxit+1)
  Q[:,1] = b/norm(b)

  g = zeros(maxit+1)
  g[1] = beta
  Js = zeros(2,maxit) # memory to store the Js
  hist = zeros(maxit)
  H = zeros(maxit+1,maxit)

  lasti = 1
  for i = 1:maxit
    lasti = i
    v = A*Q[:,i]
    H[i+1,i] = 0 # expand k
    for j = 1:i
      H[j,i]= v'*Q[:,j]
      v = v - H[j,i]*Q[:,j]
    end

    # at this point, v = H[i+1,i]*Q[:,i+1]
    H[i+1,i] = norm(v)
    #@show v
    #@show H[i+1,i]
    Q[:,i+1] = v/H[i+1,i]

    # apply J1, ... Ji-1 to z
    for j = 1:i-1
      Jj = [Js[1,j] Js[2,j]; -Js[2,j] Js[1,j]];
      H[j:j+1,i] = Jj*H[j:j+1,i];
    end
    # create Ji to zero out the final row
    Ji,d = planerot(H[i:i+1,i])
    H[i:i+1,i] = d
    Js[1,i] = Ji[1,1]
    Js[2,i] = Ji[1,2]

    # apply Ji to g
    g[i:i+1] = Ji*g[i:i+1]
    hist[i] = abs(g[i+1]);

    if (abs(g[i+1])/beta) < tol
      break
    end
  end

  i = lasti

  # produce the solution
  y = H[1:i,1:i]\g[1:i]
  x = Q[:,1:i]*y;

  if (norm(b-A*x)/beta) < tol
    flag = 0
  elseif (hist[i]/beta) < tol
    flag = -1
  else
    flag = 1
  end

  hist = hist[1:lasti]
  return x,hist,flag
end


x2,hist2,flag = gmres_efficient(A,b,1.0e-8,5)

@show hist1 - hist2

hist1 - hist2 = [-4.440892098500626e-16, 0.0, -4.440892098500626e-16, 0.0, 3.3306690738754696e-16]

5-element Vector{Float64}:
 -4.440892098500626e-16
  0.0
 -4.440892098500626e-16
  0.0
  3.3306690738754696e-16