/********************************************************
 *  ██████╗  ██████╗████████╗██╗
 * ██╔════╝ ██╔════╝╚══██╔══╝██║
 * ██║  ███╗██║        ██║   ██║
 * ██║   ██║██║        ██║   ██║
 * ╚██████╔╝╚██████╗   ██║   ███████╗
 *  ╚═════╝  ╚═════╝   ╚═╝   ╚══════╝
 * Geophysical Computational Tools & Library (GCTL)
 *
 * Copyright (c) 2022  Yi Zhang (yizhang-geo@zju.edu.cn)
 *
 * GCTL is distributed under a dual licensing scheme. You can redistribute 
 * it and/or modify it under the terms of the GNU Lesser General Public 
 * License as published by the Free Software Foundation, either version 2 
 * of the License, or (at your option) any later version. You should have 
 * received a copy of the GNU Lesser General Public License along with this 
 * program. If not, see <http://www.gnu.org/licenses/>.
 * 
 * If the terms and conditions of the LGPL v.2. would prevent you from using 
 * the GCTL, please consider the option to obtain a commercial license for a 
 * fee. These licenses are offered by the GCTL's original author. As a rule, 
 * licenses are provided "as-is", unlimited in time for a one time fee. Please 
 * send corresponding requests to: yizhang-geo@zju.edu.cn. Please do not forget 
 * to include some description of your company and the realm of its activities. 
 * Also add information on how to contact you by electronic and paper mail.
 ******************************************************/

#include "lu.h"

gctl::lu::lu(matrix<double> &sourceMatrix) : decomposedMatrix(sourceMatrix)
{
    if (sourceMatrix.empty() || sourceMatrix.row_size() != sourceMatrix.col_size())
    {
        throw domain_error("Invalid input matrix. From lu::lu(...)");
    }
}

// Decomposition into triangular matrices
void gctl::lu::decompose()
{
    // Initialize the permutation vector
    int n = decomposedMatrix.row_size();
    rowPermutation.resize(n);
    for (int i = 0; i < n; i++)
    {
        rowPermutation[i] = i;
    }

    // LU factorization
    double tmp, det = 1.0;       
    for (int p = 1; p <= n - 1; p++) 
    {
        // Find pivot element.
        for (int i = p + 1; i <= n; i++) 
        {
            if (std::fabs(decomposedMatrix[rowPermutation[i - 1]][p - 1]) > std::fabs(decomposedMatrix[rowPermutation[p - 1]][p - 1]))
            {
                // Switch the index for the p-1 pivot row if necessary.
                tmp = rowPermutation[p - 1]; rowPermutation[p - 1] = rowPermutation[i - 1]; rowPermutation[i - 1] = tmp;
                det = -det;
            }
        }

        if (decomposedMatrix[rowPermutation[p - 1]][p - 1] == 0.0)
        {
            // The matrix is singular, at least to precision of algorithm
            throw runtime_error("The input matrix is singular. From gctl::lu::decompose()");
            return;
        }

        // Multiply the diagonal elements.
        det = det * decomposedMatrix[rowPermutation[p - 1]][p - 1];

        // Form multiplier.
        for (int i = p + 1; i <= n; i++)
        {
            decomposedMatrix[rowPermutation[i - 1]][p - 1] /= decomposedMatrix[rowPermutation[p - 1]][p - 1];

            // Eliminate [p-1].
            for (int j = p + 1; j <= n; j++)
            {
                decomposedMatrix[rowPermutation[i - 1]][j - 1] -= decomposedMatrix[rowPermutation[i - 1]][p - 1] * decomposedMatrix[rowPermutation[p - 1]][j - 1];
            }
        }
    }

    det = det * decomposedMatrix[rowPermutation[n - 1]][n - 1];
    if (det == 0.0)
    {
            throw runtime_error("Determinant of the input matrix is zero. From gctl::lu::decompose()");
    }

    return;
}

// solve for x in form Ax = b.  A is the original input matrix.
// Note: b is modified in-place for row permutations
void gctl::lu::solve(const array<double>& b, array<double> &x)
{
    // Our decomposed matrix is comprised of both the lower and upper diagonal matrices.

    // The rows of this matrix have been permutated during the decomposition process.  The
    // rowPermutation indicates the proper row order.
    
    // The lower diagonal matrix only include elements below the diagonal with diagonal 
    // elements set to 1.

    // The upper diagonal matrix is fully specified.

    // First solve Ly = Pb for x using forward substitution. P is a permutated identity matrix.

    if (b.empty())
    {
        throw domain_error("Invalid target vector. From gctl::lu::solve(...)");
    }

    x.resize(b.size());

    for (int i = 0; i < x.size(); i++)
    {
        int currentRow = rowPermutation[i];

        double sum = 0.0;

        for (int j = 0; j < i; j++)
        {
            sum += (decomposedMatrix[currentRow][j] * x[j]);
        }

        x[i] = (b[currentRow] - sum);
    }

    // Now solve Uy = x for y using back substitution.  Note that 
    // y can be solved in place using the existing y vector.  No need
    // to allocate another vector.

    for (int i = b.size()-1; i >= 0; i--)
    {
        int currentRow = rowPermutation[i];

        double sum = 0.0;

        for (int j = b.size()-1; j > i; j--)
        {
            sum += (decomposedMatrix[currentRow][j] * x[j]);
        }

        x[i] = (x[i] - sum) / decomposedMatrix[currentRow][i];
    }

    return;
}