/*
 * Copyright (c) 2023 Arm Limited.
 *
 * SPDX-License-Identifier: MIT
 *
 * Permission is hereby granted, free of charge, to any person obtaining a copy
 * of this software and associated documentation files (the "Software"), to
 * deal in the Software without restriction, including without limitation the
 * rights to use, copy, modify, merge, publish, distribute, sublicense, and/or
 * sell copies of the Software, and to permit persons to whom the Software is
 * furnished to do so, subject to the following conditions:
 *
 * The above copyright notice and this permission notice shall be included in all
 * copies or substantial portions of the Software.
 *
 * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
 * IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
 * FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
 * AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
 * LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
 * OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
 * SOFTWARE.
 */

#ifndef CKW_PROTOTYPE_INCLUDE_CKW_KERNELWRITER_H
#define CKW_PROTOTYPE_INCLUDE_CKW_KERNELWRITER_H

#include "ckw/Kernel.h"
#include "ckw/TensorInfo.h"
#include "ckw/TensorOperand.h"
#include "ckw/TileInfo.h"
#include "ckw/TileOperand.h"
#include "ckw/types/ConvertPolicy.h"
#include "ckw/types/Functions.h"
#include "ckw/types/Operators.h"

#include <memory>

namespace ckw
{

namespace prototype
{
struct GpuKernelWriterAttribute;

class IGpuKernelWriter;
} // namespace prototype

/** Kernel writer. */
class KernelWriter
{
public:
    // =============================================================================================
    // Constructors and destructor
    // =============================================================================================

    /** Initialize a new instance of kernel writer.
     *
     * @param[in] kernel The kernel to be written to.
     */
    explicit KernelWriter(Kernel &kernel);

    /** Destructor */
    ~KernelWriter();

    /** No copy constructor. */
    KernelWriter(const KernelWriter &) = delete;

    /** No copy assignment. */
    KernelWriter &operator=(const KernelWriter &) = delete;

    // =============================================================================================
    // Scope management
    // =============================================================================================

    /** Get the current ID space. */
    int32_t id_space() const;

    /** Set the current ID space. */
    KernelWriter &id_space(int32_t id_space);

    /** Switch to and return a new ID space. */
    int32_t next_id_space();

    // =============================================================================================
    // Tensor and tile declaration
    // =============================================================================================

    /** Declare a tensor argument.
     *
     * @param[in] name         The name of the tensor.
     * @param[in] info         The tensor info.
     * @param[in] storage_type The tensor storage type.
     *
     * @return The @ref TensorOperand object.
     */
    TensorOperand &declare_tensor_argument(const std::string &name,
                                           const TensorInfo  &info,
                                           TensorStorageType  storage_type = TensorStorageType::BufferUint8Ptr);

    /** Declare a compile-time constant scalar argument.
     *
     * @param[in] name  The name of the tile.
     * @param[in] value The value of the tile.
     *
     * @return The @ref TileOperand object.
     */
    TileOperand &declare_tile_argument(const std::string &name, int32_t value);

    /** Declare a new tile.
     *
     * The name of the tile must be unique in the current ID space.
     *
     * @param[in] name The name of the tile.
     * @param[in] ...  The necessary arguments to create a new @ref TileOperand.
     *
     * @return The @ref TileOperand object.
     */
    template <typename... TArgs>
    TileOperand &declare_tile(const std::string &name, TArgs &&...args)
    {
        const auto var_name = generate_variable_name(name);
        auto       operand  = std::make_unique<TileOperand>(var_name, ::std::forward<TArgs>(args)...);

        return declare_tile_operand(std::move(operand));
    }

    // =============================================================================================
    // Load and store
    // =============================================================================================

    /** Load the data from the tensor memory to the tile using the sampling information.
     *
     * @param[out] tile       The tile to be loaded.
     * @param[in]  tensor     The tensor to be read.
     * @param[in]  sampler    The tensor sampling information.
     * @param[in]  dilation_y Dilation in the Y dimension.
     */
    void op_load(TileOperand             &tile,
                 const TensorOperand     &tensor,
                 const TensorTileSampler &sampler,
                 const TileOperand       &dilation_y = TileOperand("dil_y", 1));

    /** Load the data from the tensor memory to the tile using the indirect buffer approach and respective of the sampling information.
     *
     * @param[out] tile    The tile to be loaded.
     * @param[in]  tensor  The tensor to be read.
     * @param[in]  sampler The tensor sampling information.
     */
    void op_load_indirect(TileOperand &tile, const TensorOperand &tensor, const TensorTileSampler &sampler);

    /** Construct an indirection buffer in @p tile containing the precalculated addresses of elements in the source tensor.
     *
     * @param[out] tile    The tile to be loaded.
     * @param[in]  tensor  The tensor the be read.
     * @param[in]  sampler The tensor sampling information.
     * @param[in]  x       The X coordinate.
     * @param[in]  y       The Y coordinate.
     * @param[in]  x_off   Offset in the X dimension.
     * @param[in]  y_off   Offset in the Y dimension.
     */
    void util_get_indirect_buffer(TileOperand             &tile,
                                  const TensorOperand     &tensor,
                                  const TensorTileSampler &sampler,
                                  const TileOperand       &x,
                                  const TileOperand       &y,
                                  const TileOperand       &x_off,
                                  const TileOperand       &y_off);

    /** Store the tile to the tensor using the specified sampling information.
     *
     * @param[out] dst     The tensor that the tile is written to.
     * @param[in]  src     The tile to be stored.
     * @param[in]  sampler The tensor sampling information.
     */
    void op_store(TensorOperand &tensor, const TileOperand &tile, const TensorTileSampler &sampler);

    // =============================================================================================
    // Data processing
    // =============================================================================================

    /** Write assignment: `<dst> = <src>;`.
     *
     * @param[out] dst The destination tile.
     * @param[in]  src The source tile.
     */
    void op_assign(const TileOperand &dst, const TileOperand &src);

    /** Write the cast: `<dst> = convert_<dst.type><_sat>(<src>);`.
     *
     * @param[out] dst      The destination tile.
     * @param[in]  src      The source tile.
     * @param[in]  policy   The policy governing the behavior of the cast.
     */
    void op_cast_expression(const TileOperand &dst, const TileOperand &src, ConvertPolicy policy);

    /** Write the unary expression: `<dst> = <op> <src>`.
     *
     * @param[out]  dst The destination tile.
     * @param[in]   op  The unary operator.
     * @param[in]   src The source tile.
     */
    void op_unary_expression(const TileOperand &dst, UnaryOp op, const TileOperand &src);

    /** Write binary expression: `<dst> = <lhs> <op> <rhs>;`.
     *
     * @param[out] dst  The destination tile.
     * @param[in]  lhs  The LHS tile.
     * @param[in]  op   The binary operator.
     * @param[in]  rhs  The RHS tile.
     */
    void op_binary_expression(const TileOperand &dst, const TileOperand &lhs, BinaryOp op, const TileOperand &rhs);

    /** Write function applied to scalar value: `<dst> = <func>(<src>);`.
     *
     * @param[out] dst  The destination tile.
     * @param[in]  func The function to be applied to the source tile.
     * @param[in]  src  The source tile.
     */
    void op_unary_elementwise_function(const TileOperand &dst, UnaryFunction func, const TileOperand &src);

    /** Write function applied to scalar value: `<dst> = <func>(<first>, <second>);`.
     *
     * @param[out] dst      The destination tile.
     * @param[in]  func     The function to be applied to the source tiles.
     * @param[in]  first    The first argument tile.
     * @param[in]  second   The second argument tile.
     */
    void op_binary_elementwise_function(const TileOperand &dst,
                                        BinaryFunction     func,
                                        const TileOperand &first,
                                        const TileOperand &second);

    /** Write function applied to scalar value: `<dst> = <func>(<first>, <second>, <third>);`.
     *
     * @param[out] dst      The destination tile.
     * @param[in]  func     The function to be applied to the source tiles.
     * @param[in]  first    The first argument tile.
     * @param[in]  second   The second argument tile.
     * @param[in]  third    The third argument tile.
     */
    void op_ternary_elementwise_function(const TileOperand &dst,
                                         TernaryFunction    func,
                                         const TileOperand &first,
                                         const TileOperand &second,
                                         const TileOperand &third);

    /** Write if-statement: `if(<lhs> <op> <rhs>) { <body> }`.
     *
     * @param[in] lhs   The LHS tile of the condition.
     * @param[in] op    The relational binary operator.
     * @param[in] rhs   The RHS tile of the condition.
     * @param[in] body  The body of the if-statement.
     */
    void op_if(const TileOperand &lhs, BinaryOp op, const TileOperand &rhs, const std::function<void()> &body);

    /** Write else-if-statement: `else if(<lhs> <op> <rhs>) { <body> }`.
     *
     * @param[in] lhs   The LHS tile of the condition.
     * @param[in] op    The relational binary operator.
     * @param[in] rhs   The RHS tile of the condition.
     * @param[in] body  The body of the else-if-statement.
     */
    void op_else_if(const TileOperand &lhs, BinaryOp op, const TileOperand &rhs, const std::function<void()> &body);

    /** Write an else-statement: `else { <body> }`.
     *
     * @param[in] body The body of the else-statement.
     */
    void op_else(const std::function<void()> &body);

    /** Write for-loops: `for(; <var> <cond_op> <cond_value>; <var> <update_op> <update_value>) { body }`.
     *
     * @param[in]       var_name          The name of the variable used in condition.
     * @param[in]       cond_op           The relational binary operator used in condition.
     * @param[in]       cond_value_name   The value which the variable is compared against.
     * @param[in]       update_var_name   The name of the variable which is updated.
     * @param[in]       update_op         The assignment operator used for updating the update value.
     * @param[in, out]  update_value      The value which is updated at every iteration.
     * @param[in]       body              The body of the for-loop.
     */
    void op_for_loop(const TileOperand           &var_name,
                     BinaryOp                     cond_op,
                     const TileOperand           &cond_value_name,
                     const TileOperand           &update_var_name,
                     AssignmentOp                 update_op,
                     const TileOperand           &update_value_name,
                     const std::function<void()> &body);

    /** Write the return statement: `return;`
     */
    void op_return();

    // =============================================================================================
    // Misc
    // =============================================================================================

    /** Set `dst` the global ID of dimension `dim`.
     *
     * @param[out] dst The tile to be written to.
     * @param[in]  dim The global ID dimension.
     */
    void op_get_global_id(const TileOperand &dst, int32_t dim);

    // =============================================================================================
    // Code generation
    // =============================================================================================

    /** Generate the source code of the kernel. */
    ::std::string generate_code();

private:
    /** Generate the full variable name based on the original name and the ID space.
     *
     * @param[in] name The name of the variable.
     *
     * @return The full variable name.
     */
    ::std::string generate_variable_name(const std::string &name) const;

    /** Declare the tile operand.
     *
     * @param[in] operand   The tile operand to be declared.
     */
    TileOperand &declare_tile_operand(std::unique_ptr<TileOperand> operand);

private:
    Kernel                                                *_kernel;
    ::std::unique_ptr<prototype::GpuKernelWriterAttribute> _impl_attr;
    ::std::unique_ptr<prototype::IGpuKernelWriter>         _impl;

    int32_t _id_space{0};
    int32_t _max_id_space{0};
};

} // namespace ckw

#endif // CKW_PROTOTYPE_INCLUDE_CKW_KERNELWRITER_H