// SPDX-FileCopyrightText: 2017 - 2024 The Ginkgo authors
//
// SPDX-License-Identifier: BSD-3-Clause

#include "core/factorization/par_ilu_kernels.hpp"

#include <memory>

#include <ginkgo/core/base/math.hpp>
#include <ginkgo/core/matrix/coo.hpp>
#include <ginkgo/core/matrix/csr.hpp>

#include "omp/components/atomic.hpp"


namespace gko {
namespace kernels {
namespace omp {
/**
 * @brief The parallel ilu factorization namespace.
 *
 * @ingroup factor
 */
namespace par_ilu_factorization {


template <typename ValueType, typename IndexType>
void compute_l_u_factors(std::shared_ptr<const OmpExecutor> exec,
                         size_type iterations,
                         const matrix::Coo<ValueType, IndexType>* system_matrix,
                         matrix::Csr<ValueType, IndexType>* l_factor,
                         matrix::Csr<ValueType, IndexType>* u_factor)
{
    // If `iterations` is set to `Auto`, we do 3 fix-point sweeps as
    // experiments indicate this works well for many problems.
    iterations = (iterations == 0) ? 3 : iterations;
    const auto col_idxs = system_matrix->get_const_col_idxs();
    const auto row_idxs = system_matrix->get_const_row_idxs();
    const auto vals = system_matrix->get_const_values();
    const auto row_ptrs_l = l_factor->get_const_row_ptrs();
    const auto row_ptrs_u = u_factor->get_const_row_ptrs();
    const auto col_idxs_l = l_factor->get_const_col_idxs();
    const auto col_idxs_u = u_factor->get_const_col_idxs();
    auto vals_l = l_factor->get_values();
    auto vals_u = u_factor->get_values();
    for (size_type iter = 0; iter < iterations; ++iter) {
        // all elements in the incomplete factors are updated in parallel
#pragma omp parallel for
        for (size_type el = 0; el < system_matrix->get_num_stored_elements();
             ++el) {
            const auto row = row_idxs[el];
            const auto col = col_idxs[el];
            const auto val = vals[el];
            auto row_l = row_ptrs_l[row];
            auto row_u = row_ptrs_u[col];
            ValueType sum{val};
            ValueType last_operation{};
            while (row_l < row_ptrs_l[row + 1] && row_u < row_ptrs_u[col + 1]) {
                auto col_l = col_idxs_l[row_l];
                auto col_u = col_idxs_u[row_u];
                if (col_l == col_u) {
                    last_operation =
                        load(vals_l + row_l) * load(vals_u + row_u);
                    sum -= last_operation;
                } else {
                    last_operation = zero<ValueType>();
                }
                if (col_l <= col_u) {
                    ++row_l;
                }
                if (col_u <= col_l) {
                    ++row_u;
                }
            }
            // The loop above calculates: sum = system_matrix(row, col) -
            // dot(l_factor(row, :), u_factor(:, col))
            sum += last_operation;  // undo the last operation

            if (row > col) {  // modify entry in L
                auto to_write = sum / load(vals_u + row_ptrs_u[col + 1] - 1);
                if (is_finite(to_write)) {
                    store(vals_l + row_l - 1, to_write);
                }
            } else {  // modify entry in U
                auto to_write = sum;
                if (is_finite(to_write)) {
                    store(vals_u + row_u - 1, to_write);
                }
            }
        }
    }
}

GKO_INSTANTIATE_FOR_EACH_VALUE_AND_INDEX_TYPE(
    GKO_DECLARE_PAR_ILU_COMPUTE_L_U_FACTORS_KERNEL);


}  // namespace par_ilu_factorization
}  // namespace omp
}  // namespace kernels
}  // namespace gko
