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

#ifndef GKO_PUBLIC_CORE_PRECONDITIONER_JACOBI_HPP_
#define GKO_PUBLIC_CORE_PRECONDITIONER_JACOBI_HPP_


#include <ginkgo/core/base/array.hpp>
#include <ginkgo/core/base/lin_op.hpp>
#include <ginkgo/core/config/config.hpp>
#include <ginkgo/core/config/registry.hpp>
#include <ginkgo/core/matrix/csr.hpp>
#include <ginkgo/core/matrix/dense.hpp>
#include <ginkgo/core/matrix/diagonal.hpp>


namespace gko {
/**
 * @brief The Preconditioner namespace.
 *
 * @ingroup precond
 */
namespace preconditioner {


// TODO: replace this with a custom accessor
/**
 * Defines the parameters of the interleaved block storage scheme used by
 * block-Jacobi blocks.
 *
 * @tparam IndexType  type used for storing indices of the matrix
 *
 * @ingroup jacobi
 */
template <typename IndexType>
struct block_interleaved_storage_scheme {
    block_interleaved_storage_scheme() = default;

    block_interleaved_storage_scheme(IndexType block_offset,
                                     IndexType group_offset, uint32 group_power)
        : block_offset{block_offset},
          group_offset{group_offset},
          group_power{group_power}
    {}

    /**
     * The offset between consecutive blocks within the group.
     */
    IndexType block_offset;

    /**
     * The offset between two block groups.
     */
    IndexType group_offset;

    /**
     * Then base 2 power of the group.
     *
     * I.e. the group contains `1 << group_power` elements.
     */
    uint32 group_power;

    /**
     * Returns the number of elements in the group.
     *
     * @return the number of elements in the group
     */
    GKO_ATTRIBUTES IndexType get_group_size() const noexcept
    {
        return one<IndexType>() << group_power;
    }

    /**
     * Computes the storage space required for the requested number of blocks.
     *
     * @param num_blocks  the total number of blocks that needs to be stored
     *
     * @return the total memory (as the number of elements) that need to be
     *         allocated for the scheme
     *
     * @note  To simplify using the method in situations where the number of
     *        blocks is not known, for a special input `size_type{} - 1`
     *        the method returns `0` to avoid overallocation of memory.
     */
    GKO_ATTRIBUTES size_type
    compute_storage_space(size_type num_blocks) const noexcept
    {
        return (num_blocks + 1 == size_type{0})
                   ? size_type{0}
                   : ceildiv(num_blocks, this->get_group_size()) * group_offset;
    }

    /**
     * Returns the offset of the group belonging to the block with the given ID.
     *
     * @param block_id  the ID of the block
     *
     * @return the offset of the group belonging to block with ID `block_id`
     */
    GKO_ATTRIBUTES IndexType get_group_offset(IndexType block_id) const noexcept
    {
        return group_offset * (block_id >> group_power);
    }

    /**
     * Returns the offset of the block with the given ID within its group.
     *
     * @param block_id  the ID of the block
     *
     * @return the offset of the block with ID `block_id` within its group
     */
    GKO_ATTRIBUTES IndexType get_block_offset(IndexType block_id) const noexcept
    {
        return block_offset * (block_id & (this->get_group_size() - 1));
    }

    /**
     * Returns the offset of the block with the given ID.
     *
     * @param block_id  the ID of the block
     *
     * @return the offset of the block with ID `block_id`
     */
    GKO_ATTRIBUTES IndexType
    get_global_block_offset(IndexType block_id) const noexcept
    {
        return this->get_group_offset(block_id) +
               this->get_block_offset(block_id);
    }

    /**
     * Returns the stride between columns of the block.
     *
     * @return stride between columns of the block
     */
    GKO_ATTRIBUTES IndexType get_stride() const noexcept
    {
        return block_offset << group_power;
    }
};


/**
 * A block-Jacobi preconditioner is a block-diagonal linear operator, obtained
 * by inverting the diagonal blocks of the source operator.
 *
 * The Jacobi class implements the inversion of the diagonal blocks using
 * Gauss-Jordan elimination with column pivoting, and stores the inverse
 * explicitly in a customized format.
 *
 * If the diagonal blocks of the matrix are not explicitly set by the user, the
 * implementation will try to automatically detect the blocks by first finding
 * the natural blocks of the matrix, and then applying the supervariable
 * agglomeration procedure on them. However, if problem-specific knowledge
 * regarding the block diagonal structure is available, it is usually beneficial
 * to explicitly pass the starting rows of the diagonal blocks, as the block
 * detection is merely a heuristic and cannot perfectly detect the diagonal
 * block structure. The current implementation supports blocks of up to 32 rows
 * / columns.
 *
 * The implementation also includes an improved, adaptive version of the
 * block-Jacobi preconditioner, which can store some of the blocks in lower
 * precision and thus improve the performance of preconditioner application by
 * reducing the amount of memory transfers. This variant can be enabled by
 * setting the Jacobi::Factory's `storage_optimization` parameter.  Refer to the
 * documentation of the parameter for more details.
 *
 * @tparam ValueType  precision of matrix elements
 * @tparam IndexType  integral type used to store pointers to the start of each
 *                    block
 *
 * @note The current implementation supports blocks of up to 32 rows / columns.
 * @note When using the adaptive variant, there may be a trade-off in terms of
 *       slightly longer preconditioner generation due to extra work required to
 *       detect the optimal precision of the blocks.
 * @note When the max_block_size is set to 1, specialized  kernels are used,
 *       both for generation (inverting the diagonals) and application (diagonal
 *       scaling) to reduce the overhead involved in the usual (adaptive) block
 *       case.
 *
 * @ingroup jacobi
 * @ingroup precond
 * @ingroup LinOp
 */
template <typename ValueType = default_precision, typename IndexType = int32>
class Jacobi : public EnableLinOp<Jacobi<ValueType, IndexType>>,
               public ConvertibleTo<matrix::Dense<ValueType>>,
               public WritableToMatrixData<ValueType, IndexType>,
               public Transposable {
    friend class EnableLinOp<Jacobi>;
    friend class EnablePolymorphicObject<Jacobi, LinOp>;
    GKO_ASSERT_SUPPORTED_VALUE_AND_INDEX_TYPE;

public:
    using EnableLinOp<Jacobi>::convert_to;
    using EnableLinOp<Jacobi>::move_to;
    using ConvertibleTo<matrix::Dense<ValueType>>::convert_to;
    using ConvertibleTo<matrix::Dense<ValueType>>::move_to;
    using value_type = ValueType;
    using index_type = IndexType;
    using mat_data = matrix_data<ValueType, IndexType>;
    using transposed_type = Jacobi<ValueType, IndexType>;

    /**
     * Returns the number of blocks of the operator.
     *
     * @return the number of blocks of the operator
     *
     * @internal
     * TODO: replace with ranges
     */
    size_type get_num_blocks() const noexcept { return num_blocks_; }

    /**
     * Returns the storage scheme used for storing Jacobi blocks.
     *
     * @return the storage scheme used for storing Jacobi blocks
     *
     * @internal
     * TODO: replace with ranges
     */
    const block_interleaved_storage_scheme<index_type>& get_storage_scheme()
        const noexcept
    {
        return storage_scheme_;
    }

    /**
     * Returns the pointer to the memory used for storing the block data.
     *
     * Element (`i`, `j`) of block `b` is stored in position
     * `(get_block_pointers()[b] + i) * stride + j` of the array.
     *
     * @return the pointer to the memory used for storing the block data
     *
     * @internal
     * TODO: replace with ranges
     */
    const value_type* get_blocks() const noexcept
    {
        return blocks_.get_const_data();
    }

    /**
     * Returns an array of 1-norm condition numbers of the blocks.
     *
     * @return an array of 1-norm condition numbers of the blocks
     *
     * @note This value is valid only if adaptive precision variant is used, and
     *       implementations of the standard non-adaptive variant are allowed to
     *       omit the calculation of condition numbers.
     */
    const remove_complex<value_type>* get_conditioning() const noexcept
    {
        return conditioning_.get_const_data();
    }

    /**
     * Returns the number of elements explicitly stored in the matrix.
     *
     * @return the number of elements explicitly stored in the matrix
     */
    size_type get_num_stored_elements() const noexcept
    {
        return blocks_.get_size();
    }

    void convert_to(matrix::Dense<value_type>* result) const override;

    void move_to(matrix::Dense<value_type>* result) override;

    void write(mat_data& data) const override;

    std::unique_ptr<LinOp> transpose() const override;

    std::unique_ptr<LinOp> conj_transpose() const override;

    /**
     * Copy-assigns a Jacobi preconditioner. Preserves executor, copies all
     * data and parameters.
     */
    Jacobi& operator=(const Jacobi& other);

    /**
     * Move-assigns a Jacobi preconditioner. Preserves executor, moves all data
     * and parameters. The moved-from object will be empty (0x0 and default
     * parameters).
     */
    Jacobi& operator=(Jacobi&& other);

    /**
     * Copy-constructs a Jacobi preconditioner. Inherits executor, copies all
     * data and parameters.
     */
    Jacobi(const Jacobi& other);

    /**
     * Move-assigns a Jacobi preconditioner. Inherits executor, moves all data
     * and parameters. The moved-from object will be empty (0x0 and default
     * parameters).
     */
    Jacobi(Jacobi&& other);

    GKO_CREATE_FACTORY_PARAMETERS(parameters, Factory)
    {
        /**
         * Maximal size of diagonal blocks.
         *
         * @note This value has to be between 1 and 32 (NVIDIA)/64 (AMD). For
         * efficiency, when the max_block_size is set to 1, specialized kernels
         * are used and the additional objects (block_ptrs etc) are set to null
         * values.
         */
        uint32 GKO_FACTORY_PARAMETER_SCALAR(max_block_size, 32u);

        /**
         * Stride between two columns of a block (as number of elements).
         *
         * Should be a multiple of cache line size for best performance.
         *
         * @note If this value is 0, it uses 64 in hip AMD but 32 in NVIDIA or
         *       reference executor. The allowed value: 0, 64 for AMD and 0, 32
         *       for NVIDIA
         */
        uint32 GKO_FACTORY_PARAMETER_SCALAR(max_block_stride, 0u);

        /**
         * @brief `true` means it is known that the matrix given to this
         *        factory will be sorted first by row, then by column index,
         *        `false` means it is unknown or not sorted, so an additional
         *        sorting step will be performed during the preconditioner
         *        generation (it will not change the matrix given).
         *        The matrix must be sorted for this preconditioner to work.
         *
         * The `system_matrix`, which will be given to this factory, must be
         * sorted (first by row, then by column) in order for the algorithm
         * to work. If it is known that the matrix will be sorted, this
         * parameter can be set to `true` to skip the sorting (therefore,
         * shortening the runtime).
         * However, if it is unknown or if the matrix is known to be not sorted,
         * it must remain `false`, otherwise, this preconditioner might be
         * incorrect.
         */
        bool GKO_FACTORY_PARAMETER_SCALAR(skip_sorting, false);

        /**
         * Starting (row / column) indexes of individual blocks.
         *
         * An index past the last block has to be supplied as the last value.
         * I.e. the size of the array has to be the number of blocks plus 1,
         * where the first value is 0, and the last value is the number of
         * rows / columns of the matrix.
         *
         * @note Even if not set explicitly, this parameter will be set to
         *       automatically detected values once the preconditioner is
         *       generated.
         * @note If the parameter is set automatically, the size of the array
         *       does not correlate to the number of blocks, and is
         *       implementation defined. To obtain the number of blocks `n` use
         *       Jacobi::get_num_blocks(). The starting indexes of the blocks
         *       are stored in the first `n+1` values of this array.
         * @note If the block-diagonal structure can be determined from the
         *       problem characteristics, it may be beneficial to pass this
         *       information specifically via this parameter, as the
         *       autodetection procedure is only a rough approximation of the
         *       true block structure.
         * @note The maximum block size set by the max_block_size parameter
         *       has to be respected when setting this parameter. Failure to do
         *       so will lead to undefined behavior.
         */
        gko::array<index_type> GKO_FACTORY_PARAMETER_VECTOR(block_pointers,
                                                            nullptr);

        /**
         * Use L1 Jacboi, which is introduced in the paper A. H. Baker et al.
         * "Multigrid smoothers for ultraparallel computing." This paper
         * discusses this type of smoother with the following matrix property.
         * $ A_{ii} \geq \theta \sum_{j \in \text{off diagonal block}} |A_{ij}|,
         * \theta \geq 0 $
         * If it is true, it generates the preconditioner on A +
         * Diag(sum_{k in off-diagonal block of i} |A_ik|) instead of A. We
         * aggregate the absolute value of the entries with the same row in the
         * off-diagonal block into the diagonal value.
         */
        bool GKO_FACTORY_PARAMETER_SCALAR(aggregate_l1, false);

    private:
        // See documentation of storage_optimization parameter for details about
        // this class
        struct storage_optimization_type {
            storage_optimization_type(precision_reduction p)
                : is_block_wise{false}, of_all_blocks{p}
            {}

            storage_optimization_type(
                const array<precision_reduction>& block_wise_opt)
                : is_block_wise{block_wise_opt.get_size() > 0},
                  block_wise{block_wise_opt}
            {}

            storage_optimization_type(
                array<precision_reduction>&& block_wise_opt)
                : is_block_wise{block_wise_opt.get_size() > 0},
                  block_wise{std::move(block_wise_opt)}
            {}

            operator precision_reduction() { return of_all_blocks; }

            bool is_block_wise;
            precision_reduction of_all_blocks;
            gko::array<precision_reduction> block_wise;
        };

    public:
        /**
         * The precisions to use for the blocks of the matrix.
         *
         * This parameter can either be a single instance of precision_reduction
         * or an array of precision_reduction values. If set to
         * `precision_reduction(0, 0)` (this is the default), a regular
         * full-precision block-Jacobi will be used. Any other value (or an
         * array of values) will map to the adaptive variant.
         *
         * The best starting point when evaluating the potential of the adaptive
         * version is to set this parameter to
         * `precision_reduction::autodetect()`. This option will cause the
         * preconditioner to reduce the memory transfer volume as much as
         * possible, while trying to maintain the quality of the preconditioner
         * similar to that of the full precision block-Jacobi.
         *
         * For finer control, specific instances of precision_reduction can be
         * used. Supported values are `precision_reduction(0, 0)`,
         * `precision_reduction(0, 1)` and `precision_reduction(0, 2)`. Any
         * other value will have the same effect as `precision_reduction(0, 0)`.
         *
         * If the ValueType template parameter is set to `double` (or the
         * complex variant `std::complex<double>`), `precision_reduction(0, 0)`
         * will use IEEE double precision for preconditioner storage,
         * `precision_reduction(0, 1)` will use IEEE single precision, and
         * `precision_reduction(0, 2)` will use IEEE half precision.
         *
         * It ValueType is set to `float` (or `std::complex<float>`),
         * `precision_reduction(0, 0)` will use IEEE single precision for
         * preconditioner storage, and both `precision_reduction(0, 1)` and
         * `precision_reduction(0, 2)` will use IEEE half precision.
         *
         * Instead of specifying the same precision for all blocks, the
         * precision of the elements can be specified on per-block basis by
         * passing an array of precision_reduction objects. All values discussed
         * above are supported, with the same meaning. It is worth mentioning
         * that a value of `precision_reduction::autodetect()` will cause
         * autodetection on the per-block basis, so blocks whose precisions are
         * autodetected can end up having different precisions once the
         * preconditioner is generated. The detected precision generally depends
         * on the conditioning of the block.
         *
         * If the number of diagonal blocks is larger than the number of
         * elements in the passed array, the entire array will be replicated
         * until enough values are available. For example, if the original array
         * contained two precisions `(x, y)` and the preconditioner contains 5
         * blocks, the array will be transformed into `(x, y, x, y, x)` before
         * generating the preconditioner. As a consequence, specifying a single
         * value for this property is exactly equivalent to specifying an array
         * with a single element set to that value.
         *
         * Once an instance of the Jacobi linear operator is generated, the
         * precisions used for the blocks can be obtained by reading this
         * property. Whether the parameter was set to a single value or to an
         * array of values can be queried by reading the
         * `storage_optimization.is_block_wise` boolean sub-property. If it is
         * set to `false`, the precision used for all blocks can be obtained
         * using `storage_optimization.of_all_blocks` or by casting
         * `storage_optimization` to `precision_reduction`. Independently of the
         * value of `storage_optimization.is_block_wise`, the
         * `storage_optimization.block_wise` property will return an array of
         * precisions used for each block. All values set to
         * `precision_reduction::autodetect()` will be replaced with the value
         * representing the precision used for the corresponding block.
         * If the non-adaptive version of Jacobi is used, the
         * `storage_optimization.block_wise` array will be empty.
         */
        storage_optimization_type GKO_FACTORY_PARAMETER_VECTOR(
            storage_optimization, precision_reduction(0, 0));

        /**
         * The relative accuracy of the adaptive Jacobi variant.
         *
         * This parameter is only used if the adaptive version of the algorithm
         * is selected (see storage_optimization parameter for more details).
         * The parameter is used when detecting the optimal precisions of blocks
         * whose precision has been set to precision_reduction::autodetect().
         *
         * The parameter represents the number of correct digits in the result
         * of Jacobi::apply() operation of the adaptive variant, compared to the
         * non-adaptive variant. In other words, the total preconditioning error
         * will be:
         *
         * ```
         * || inv(A)x - inv(M)x|| / || inv(A)x || <= c * (dropout + accuracy)
         * ```
         *
         * where `c` is some constant depending on the problem size and roundoff
         * error and `dropout` the error introduced by disregarding off-diagonal
         * elements.
         *
         * Larger values reduce the volume of memory transfer, but increase
         * the error compared to using full precision storage. Thus, tuning the
         * accuracy to a value as close as possible to `dropout` will result in
         * optimal memory savings, while not degrading the quality of solution.
         */
        remove_complex<value_type> GKO_FACTORY_PARAMETER_SCALAR(
            accuracy, static_cast<remove_complex<value_type>>(1e-1));
    };
    GKO_ENABLE_LIN_OP_FACTORY(Jacobi, parameters, Factory);
    GKO_ENABLE_BUILD_METHOD(Factory);

    /**
     * Create the parameters from the property_tree.
     * Because this is directly tied to the specific type, the value/index type
     * settings within config are ignored and type_descriptor is only used
     * for children configs.
     *
     * @param config  the property tree for setting
     * @param context  the registry
     * @param td_for_child  the type descriptor for children configs. The
     *                      default uses the value/index type of this class.
     *
     * @return parameters
     *
     * @note Jacobi does not support block_pointers and storage_optimization
     *       array.
     */
    static parameters_type parse(
        const config::pnode& config, const config::registry& context,
        const config::type_descriptor& td_for_child =
            config::make_type_descriptor<ValueType, IndexType>());

protected:
    /**
     * Creates an empty Jacobi preconditioner.
     *
     * @param exec  the executor this object is assigned to
     */
    explicit Jacobi(std::shared_ptr<const Executor> exec)
        : EnableLinOp<Jacobi>(exec),
          num_blocks_{},
          blocks_(exec),
          conditioning_(exec)
    {
        parameters_.block_pointers.set_executor(exec);
        parameters_.storage_optimization.block_wise.set_executor(exec);
    }

    /**
     * Creates a Jacobi preconditioner from a matrix using a Jacobi::Factory.
     *
     * @param factory  the factory to use to create the preconditioner
     * @param system_matrix  the matrix this preconditioner should be created
     *                       from
     */
    explicit Jacobi(const Factory* factory,
                    std::shared_ptr<const LinOp> system_matrix)
        : EnableLinOp<Jacobi>(factory->get_executor(),
                              gko::transpose(system_matrix->get_size())),
          parameters_{factory->get_parameters()},
          storage_scheme_{this->compute_storage_scheme(
              parameters_.max_block_size, parameters_.max_block_stride)},
          num_blocks_{parameters_.block_pointers.get_size() - 1},
          blocks_(factory->get_executor(),
                  storage_scheme_.compute_storage_space(
                      parameters_.block_pointers.get_size() - 1)),
          conditioning_(factory->get_executor())
    {
        parameters_.block_pointers.set_executor(this->get_executor());
        parameters_.storage_optimization.block_wise.set_executor(
            this->get_executor());
        this->generate(system_matrix.get(), parameters_.skip_sorting);
    }

    /**
     * Computes the storage scheme suitable for storing blocks of a given
     * maximum size.
     *
     * @param max_block_size  the maximum size of the blocks
     *
     * @return a suitable storage scheme
     */
    block_interleaved_storage_scheme<index_type> compute_storage_scheme(
        uint32 max_block_size, uint32 param_max_block_stride)
    {
        uint32 default_block_stride = 32;
        // If the executor is hip, the warp size is 32 or 64
        if (auto hip_exec = std::dynamic_pointer_cast<const gko::HipExecutor>(
                this->get_executor())) {
            default_block_stride = hip_exec->get_warp_size();
        }
        uint32 max_block_stride = default_block_stride;
        if (param_max_block_stride != 0) {
            // if parameter max_block_stride is not zero, set max_block_stride =
            // param_max_block_stride
            max_block_stride = param_max_block_stride;
            if (this->get_executor() != this->get_executor()->get_master() &&
                max_block_stride != default_block_stride) {
                // only support the default value on the gpu device
                GKO_NOT_SUPPORTED(this);
            }
        }
        if (parameters_.max_block_size > max_block_stride ||
            parameters_.max_block_size < 1) {
            GKO_NOT_SUPPORTED(this);
        }
        const auto group_size = static_cast<uint32>(
            max_block_stride / get_superior_power(uint32{2}, max_block_size));
        const auto block_offset = max_block_size;
        const auto block_stride = group_size * block_offset;
        const auto group_offset = max_block_size * block_stride;
        return {static_cast<index_type>(block_offset),
                static_cast<index_type>(group_offset),
                get_significant_bit(group_size)};
    }

    /**
     * Generates the preconditioner.
     *
     * @param system_matrix  the source matrix used to generate the
     *                       preconditioner
     * @param skip_sorting  determines if the sorting of system_matrix can be
     *                      skipped (therefore, marking that it is already
     *                      sorted)
     */
    void generate(const LinOp* system_matrix, bool skip_sorting);

    /**
     * Detects the diagonal blocks and allocates the memory needed to store the
     * preconditioner.
     *
     * @param system_matrix  the source matrix whose diagonal block pattern is
     *                       to be detected
     */
    void detect_blocks(const matrix::Csr<ValueType, IndexType>* system_matrix);

    void apply_impl(const LinOp* b, LinOp* x) const override;

    void apply_impl(const LinOp* alpha, const LinOp* b, const LinOp* beta,
                    LinOp* x) const override;

private:
    block_interleaved_storage_scheme<index_type> storage_scheme_{};
    size_type num_blocks_;
    array<value_type> blocks_;
    array<remove_complex<value_type>> conditioning_;
};


}  // namespace preconditioner
}  // namespace gko


#endif  // GKO_PUBLIC_CORE_PRECONDITIONER_JACOBI_HPP_
