/********************************************************
 *  ██████╗  ██████╗████████╗██╗
 * ██╔════╝ ██╔════╝╚══██╔══╝██║
 * ██║  ███╗██║        ██║   ██║
 * ██║   ██║██║        ██║   ██║
 * ╚██████╔╝╚██████╗   ██║   ███████╗
 *  ╚═════╝  ╚═════╝   ╚═╝   ╚══════╝
 * 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 "cholesky.h"

// Constructor

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

// Decomposition into triangular matrices

void gctl::cholesky::decompose()
{
    // Enumerate matrix columnwise
    for (int j = 0; j < decomposedMatrix.col_size(); j++)
    {
        for (int i = j; i < decomposedMatrix.row_size(); i++)
        {
            if (i == j)
            {
                double sum = 0.0;

                for (int k = 0; k < i; k++)
                {
                    sum += std::pow(decomposedMatrix[i][k], 2.0);
                }

                if (decomposedMatrix[i][j] - sum <= 0.0)
                {
                    // Not positive definite matrix
                    throw runtime_error("The input matrix is not positively defined. From gctl::cholesky::decompose()");
                    return;
                }

                decomposedMatrix[i][j] = std::sqrt(decomposedMatrix[i][j] - sum);
            }
            else
            {
                double sum = 0.0;

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

                decomposedMatrix[i][j] = (1 / decomposedMatrix[j][j]) * (decomposedMatrix[i][j] - sum);
                decomposedMatrix[j][i] = decomposedMatrix[i][j];
            }
        }
    }
        
    return;
}

// Solve for x in form Ax = b.  A is the original input matrix.

void gctl::cholesky::solve(const array<double>& b, array<double> &x)
{
    if (b.empty())
    {
        throw domain_error("Invalid target vector. From lu<double>::solve(...)");
    }

    x.resize(b.size());

    // First solve lower triangular * x = b with forward substitution

    for (int i = 0; i < b.size(); i++)
    {
        double sum = 0.0;

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

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

    // Now solve upper triangular (transpose of lower triangular) * x = x with back substitution.
    // Note that x can be solved in place using the existing x vector.  No need to allocate 
    // another vector.

    for (int i = static_cast<int>(b.size()) - 1; i >= 0; i--)
    {
        double sum = 0.0;

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

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

    return;
}