##
## Utility functions for this chapter.
##

exchange_two_rows = function(W, row_idx1, row_idx2){
    row_copy = W[row_idx1, ]
    W[row_idx1, ] = W[row_idx2, ]
    W[row_idx2, ] = row_copy
    W
}


extract_U_matrix = function(W, n, ps, qs){
    U = matrix(0, nrow=n, ncol=n)
    for(ii in 1:n){
        for(jj in ii:n){
            U[ii, jj] = W[ps[ii], qs[jj]]
        }
    }
    U
}


extract_L_matrix = function(W, n, ps, qs){
    L = matrix(0, nrow=n, ncol=n)
    for(ii in 1:n){
        for(jj in 1:ii){
            if( jj<ii ){
                L[ii, jj] = W[ps[ii], qs[jj]]
            }else{
                L[ii, jj] = 1.0 # the diagonal entry
            }
        }
    }
    L
}


GE_w_scaled_partial_pivoting = function(A, Bs, verbose=FALSE){
    ##
    ## Implements Algorithm 4.2 (Gaussian Elimination) with scaled partial pivoting in place.
    ## The 'in place' part means that we don't move the elements of A or Bs around so no copying of elements rows is done.
    ## The 'right-hand-side' Bs can be a single column or several columns (i.e. if we are solving Ax=b for several right-hand-sides each b is a column of the matrix Bs)
    ##
    n = dim(A)[1]

    ## Make sure that Bs is a matrix:
    if( !is.matrix(Bs) ){
        Bs = matrix(Bs, nrow=n, byrow=FALSE)
    }
    nRHS = dim(Bs)[2] ## the number of right-hand-sides we are solving for

    ## Compute the vector d[i] (absmax along each row):
    ##
    di = apply(abs(A), 1, max, na.rm=TRUE)

    ## ps(kk] is the pivot equation used in the kk-th step:
    ##
    ps = 1:n

    W = cbind(A, Bs)
    colnames(W) = NULL
    IFLAG = 0 ## counts the total number of row exchanges

    ## Do the 'upper-triangularization' part (this is done 'in-place')
    ##
    for(kk in 1:(n-1)){
        if( verbose ){
            print(sprintf('kk= %d', kk))
            print(W)
        }

        ## Select the pivot row for the kk-th step:
        ##
        unpivoted_rows = ps[kk:n] ## <- these rows have not been pivoted on yet
        if( verbose ){
            print('unpivoted_rows=')
            print(unpivoted_rows)
        }

        pivot_benefit = abs(W[unpivoted_rows, kk])/di[unpivoted_rows]
        unpivoted_row_indx = which.max(pivot_benefit)
        if( unpivoted_row_indx!=1 ){ ## we need to exchange two rows
            IFLAG = IFLAG + 1
        }
        if( verbose ){
            print('pivot_benefit=')
            print(pivot_benefit)
            print(sprintf('unpivoted_row_indx= %d; pivot row= %d', unpivoted_row_indx, unpivoted_rows[unpivoted_row_indx]))
        }

        ## Move this pivot row to the correct location in ps:
        ## The numbers in ps[1:kk] represent the rows and order we eliminate the variables x_1, x_2, ... x_kk variables (in the original matrix):
        ##
        tmp = ps[kk]
        ps[kk] = unpivoted_rows[unpivoted_row_indx]
        ps[(kk-1) + unpivoted_row_indx] = tmp
        if( verbose ){
            print('ps after pivot exchange=')
            print(ps)
        }

        ## Perform elimination of the variable x[kk] in the non-pivot rows using the row ps[kk]
        ##
        ps_kk = ps[kk]
        unpivoted_rows = ps[(kk+1):n] ## <- these rows need to have the variable x[kk] elliminated
        for(ri in unpivoted_rows){
            m = W[ri, kk] = W[ri, kk] / W[ps_kk, kk]
            W[ri, (kk+1):(n+nRHS)] = W[ri, (kk+1):(n+nRHS)] - m * W[ps_kk, (kk+1):(n+nRHS)]
        }

        if( verbose ){
            print('')
        }
    }

    ## Do the 'back-substitution' part:
    ##
    X = matrix(data=0, nrow=n, ncol=nRHS, byrow=TRUE)
    for(kk in n:1){
        rhs = W[ps[kk], (n+1):(n+nRHS)]
        if(kk==n){
            X[kk, ] = rhs / W[ps[kk], kk]
        }else{
            known_X = X[(kk+1):n, ]
            X_coefficients = W[ps[kk], (kk+1):n]
            X[kk, ] = ( rhs - (X_coefficients %*% known_X) ) / W[ps[kk], kk]
        }
    }

    ## Here we extract the L and U in the factorization A[ps, 1:n] = PA = LU (i.e. only row exchanges):
    ## W[ps, 1:n] - L %*% U == 0
    ##
    U = extract_U_matrix(W, n, ps, 1:n)
    L = extract_L_matrix(W, n, ps, 1:n)

    results = list('X'=X, 'ps'=ps, 'U'=U, 'L'=L, 'IFLAG'=IFLAG)
}



GE_w_total_pivoting = function(A, Bs, verbose=FALSE){
    ##
    ## Implements Algorithm 4.2 (Gaussian Elimination) with total pivoting in place.
    ## The 'in place' part means that we don't move the elements of A or Bs around so no copying of elements rows is done
    ## The 'right-hand-side' Bs can be a single column or several columns (i.e. if we are solving Ax=b for several right-hand-sides each b is a column of the matrix Bs)
    ##
    n = dim(A)[1]

    ## Make sure that Bs is a matrix:
    if( !is.matrix(Bs) ){
        Bs = matrix(Bs, nrow=n, byrow=FALSE)
    }
    nRHS = dim(Bs)[2] ## the number of right-hand-sides we are solving for

    ## ps[kk] is the pivot equation used in the kk-th step:
    ##
    ps = 1:n

    ## qs[kk] is the variable we will eliminate from all non pivot equations in the kk-th step:
    ##
    qs = 1:n

    W = cbind(A, Bs)
    colnames(W) = NULL
    IFLAG = 0 ## counts the total number of row and column exchanges

    ## Do the 'upper-triangularization' part (this is done 'in-place')
    ##
    for(kk in 1:(n-1)){
        if( verbose ){
            print(sprintf('kk= %d', kk))
            print(W)
        }

        ## Select the pivot entry for the kk-th step:
        ##
        unpivoted_rows = ps[kk:n] ## <- these rows have not been pivoted on yet
        unpivoted_cols = qs[kk:n] ## <- these variables have not been eliminated yet
        if( verbose ){
            print('unpivoted_rows=')
            print(unpivoted_rows)
            print('unpivoted_cols=')
            print(unpivoted_cols)
        }

        pivot_benefit = abs(W[unpivoted_rows, unpivoted_cols])
        unpivoted_row_col_indx = which(pivot_benefit == max(pivot_benefit), arr.ind=TRUE)
        unpivoted_row_col_indx = unpivoted_row_col_indx[1, ] ## in case there are multiple equivalent maximums ... take only one
        unpivoted_row_indx = unpivoted_row_col_indx[1]
        unpivoted_col_indx = unpivoted_row_col_indx[2]
        if( unpivoted_row_indx!=1 ){ ## we need to exchange two rows
            IFLAG = IFLAG + 1
        }
        if( unpivoted_col_indx!=1 ){ ## we need to exchange two columns
            IFLAG = IFLAG + 1
        }
        if( verbose ){
            print('pivot_benefit=')
            print(pivot_benefit)
            print(sprintf('unpivoted_row_indx= %d; pivot_row= %d; unpivoted_col_indx= %d; pivot_col= %d', unpivoted_row_indx, unpivoted_rows[unpivoted_row_indx], unpivoted_col_indx, unpivoted_cols[unpivoted_col_indx]))
        }

        ## Move this pivot row to the correct location in ps and qs:
        ## The numbers in ps[1:kk] represent the rows (in the original matrix) and order we eliminate the
        ## variables qs[1:kk]
        ##
        tmp = ps[kk]
        ps[kk] = unpivoted_rows[unpivoted_row_indx]
        ps[(kk-1) + unpivoted_row_indx] = tmp
        tmp = qs[kk]
        qs[kk] = unpivoted_cols[unpivoted_col_indx]
        qs[(kk-1) + unpivoted_col_indx] = tmp
        if( verbose ){
            print('ps after pivot exchange=')
            print(ps)
            print('qs after pivot exchange=')
            print(qs)
        }

        ## Perform elimination of the variable qs[kk] in the non-pivot rows using the row ps[kk]
        ##
        ps_kk = ps[kk]
        qs_kk = qs[kk]
        unpivoted_rows = ps[(kk+1):n] ## <- these rows need to have the variable qs[kk] elliminated
        unpivoted_cols = qs[(kk+1):n] ## <- these cols need to have the variable qs[kk] elliminated
        for(ri in unpivoted_rows){
            m = W[ri, qs_kk] = W[ri, qs_kk] / W[ps_kk, qs_kk]
            W[ri, unpivoted_cols] = W[ri, unpivoted_cols] - m * W[ps_kk, unpivoted_cols] ## <- do the variables for this row
            W[ri, (n+1):(n+nRHS)] = W[ri, (n+1):(n+nRHS)] - m * W[ps_kk, (n+1):(n+nRHS)] ## <- do the right-hand-side for this row
        }

    }

    if( verbose ){
        print('')
    }

    ## Do the 'back-substitution' part:
    ##
    X = matrix(data=0, nrow=n, ncol=nRHS, byrow=TRUE)
    for(kk in n:1){
        rhs = W[ps[kk], (n+1):(n+nRHS)]
        if(kk==n){
            X[qs[kk], ] = rhs / W[ps[kk], qs[kk]]
        }else{
            known_X = X[qs[(kk+1):n], ]
            X_coefficients = W[ps[kk], qs[(kk+1):n]]
            X[qs[kk], ] = ( rhs - (X_coefficients %*% known_X) ) / W[ps[kk], qs[kk]]
        }
    }

    ## Here we extract the L and U in the factorization A[ps, qs] = PAQ = LU (i.e. both row and column exchanges):
    ## W[ps, qs] - L %*% U == 0
    ##
    U = extract_U_matrix(W, n, ps, qs)
    L = extract_L_matrix(W, n, ps, qs)

    results = list('X'=X, 'ps'=ps, 'qs'=qs, 'U'=U, 'L'=L, 'IFLAG'=IFLAG)
}


build_tridiagonal_matrix = function(n, a, b, c){
    ##
    ## a = element on the diagonal
    ## b = element on the superdiagonal
    ## c = element on the subdiagonal
    ##
    as = rep(a, n)
    A = diag(as)
    A[(row(A)-col(A)) == -1] = b
    A[(row(A)-col(A)) == +1] = c
    A
}


exact_tridiagonal_eigensystem = function(n, d, e){
    ##
    ## Returns the exact eigensystem (eigenvalues and eigenvectors) for the given tridiagonal system:
    ## d = element on the diagonal
    ## e = element on the subdiagonal and superdiagonal
    ##

    ## Extract the eigenvectors:
    ##
    U = matrix(0, nrow=n, ncol=n)
    for(jj in 1:n){
        U[, jj] = sin( ( (1:n) * jj * pi) / (n+1) )
    }

    ## Extract the eigenvalues:
    ##
    lambda = rep(0, n)
    for(jj in 1:n){
        lambda[jj] = d + 2*e*cos( ( jj * pi )/(n+1) )
    }

    ## Order the eigenvalues from smallest to largest:
    ##
    lambda_order = order(lambda)
    lambda = lambda[lambda_order]
    U = U[, lambda_order]

    results = list('U'=U, 'lambda'=lambda)
}


power_method = function(A, tol=1.e-3, max_iter=1000, init_choice=0){
    ##
    ## Use the power method to find the largest eigenvalue and associated eigenvector for a matrix A.
    ##

    ## Choose how to initialize our random vector z:
    ##
    if( init_choice==0 ){
        z = rep(0, n); z[1] = 1
    }else if( init_choice==1 ){
        z = A[, 1]
    }else if( init_choice==2 ){
        z = rep(1, n)
    }
    lambda_est = NaN

    u = z
    for(iter in 1:max_iter){
        u = A %*% u ## u = A^iter z
        u = u/norm(u, '2') ## normalize to prevent overflow
        u = sign(u[1])*u ## normalize the sign of the first element of u to be positive. This prevent oscillations between +v and -v.
        lambda_next = (t(u) %*% A %*% u)/(t(u) %*% u)
        if(iter==1){
            lambda_est = lambda_next
            eigenvector_est = u
        }else{
            lambda_error = lambda_est-lambda_next
            eigenvector_error = norm(u - eigenvector_est, '2')
            if( abs(lambda_error)>tol || eigenvector_error>tol ){
                lambda_est = lambda_next
                eigenvector_est = u
            }else{
                results = list('max_eigenvalue'=lambda_est[1], 'max_eigenvector'=u, 'n_iterations'=iter, 'max_iter'=max_iter)
                break
            }
        }
    }
    if( !exists('results') ){
        print(sprintf('ERROR: power method did not converge; iter= %d; max_iter= %d; lambda_error= %f; eigenvector_error= %f', iter, max_iter, lambda_error, eigenvector_error) )
        results = NULL
    }

    results
}


determinant_from_GE = function(A){
    ##
    ## Uses Gaussian elimination to compute the determinant of the matrix A
    ##
    if( FALSE ){
        det = det(A) ## for debugging
    }else{
        res = GE_w_scaled_partial_pivoting(A, rep(1, dim(A)[1]))
        ##res = GE_w_total_pivoting(A, rep(1, dim(A)[1]))
        det = (-1)^res$IFLAG * prod(diag(res$U))
    }
    det
}


householder_reduction = function(A, verbose=FALSE){
    ##
    ## Performs Householder reflections to reduce A to upper Hessenberg form.
    ##
    ## This code is designed to be easy to understand and not necessarily the most efficient.
    ##
    ## Returns upper Hessenberg matrix H and unitary P such that H = P * A * P' (with P'* P = I)
    ##
    n = dim(A)[1]

    H = A
    P = diag(n)
    for(kk in 1:(n-2)){
        ## Tranform the kkth column:
        x_col = H[, kk]
        y = rep(0, n)
        y[(kk+1):n] = x_col[(kk+1):n]
        beta = sign(x_col[kk+1])*norm(x_col[(kk+1):n], '2')
        y[kk+1] = y[kk+1] + beta
        if( abs(beta*y[kk+1])>1.e-6 ){ ## avoid "divide by zero"
            alpha = 1/(beta*y[kk+1])
        }else{
            alpha = 1/(sign(beta*y[kk+1])*1.e-6)
        }
        R = diag(n) - alpha * y %*% t(y)
        ##print(round(R %*% R, 2)) ## should equal I
        P = R %*% P
        H = R %*% H %*% R
        if( verbose ){
            print(sprintf('kk= %5d; y[kk+1]= %10.6e; beta= %10.6e; alpha= %10.6e', kk, y[kk+1], beta, alpha))
            print(round(H, 3))
        }
    }
    results = list('P'=P, 'H'=H)
    results
}


evaluate_symmetric_tridiagonal_characteristic_polynomial = function(lambda, as, bs, verbose=FALSE){
    ##
    ## Evaluates the characteristic polynomial p(lambda) of a symmetric tridiagonal matrix.
    ##
    ## as= diagonal elements
    ## bs= sup/super diagonal elements
    ##
    ## ps[n+1] = p_n(lambda)
    ## pprimes[n+1] = p_n'(lambda) 
    ## 
    n = length(as)
    stopifnot(length(bs)==(n-1))

    ps = rep(NA, n+1)
    pprimes = rep(NA, n+1)
    ps[1] = 1 ## p_0(lambda) 
    ps[2] = as[1]-lambda ## p_1(lambda)
    pprimes[1] = 0
    pprimes[2] = -1
    for(jj in 3:(n+1)){
        ps[jj] = (as[jj-1] - lambda)*ps[jj-1] - (bs[jj-2]^2)*ps[jj-2]
        pprimes[jj] = -ps[jj-1] + (as[jj-1] - lambda)*pprimes[jj-1] - (bs[jj-2]^2)*pprimes[jj-2] 
    }
    
    result = list('ps'=ps, 'pprimes'=pprimes)
    result
}