#ifndef AI2D_H
#define AI2D_H

#include <vector>
#include <iostream>
#include <fstream>
#include <sstream>
#include <string>
#include <nncase/functional/ai2d/ai2d_builder.h>

// Use nncase namespace for easier access to AI2D related types and functions
using namespace nncase;
using namespace nncase::runtime;
using namespace nncase::runtime::k230;
using namespace nncase::F::k230;
using std::vector;

/**
 * @brief AI2D image preprocessing wrapper class
 * 
 * Provides cropping, resizing, padding, affine transformation, data type setting, 
 * and execution interfaces for image preprocessing before nncase inference.
 */
class AI2D
{
public:
    /**
     * @brief Constructor
     */
    explicit AI2D();

    /**
     * @brief Destructor
     */
    ~AI2D();

    // Disable copy constructor and copy assignment
    AI2D(const AI2D&) = default;
    AI2D& operator=(const AI2D&) = default;

    // Enable move constructor and move assignment
    AI2D(AI2D&&) noexcept = default;
    AI2D& operator=(AI2D&&) noexcept = default;

    /**
     * @brief Set AI2D input/output data types and formats
     * @param input_format Input data format
     * @param output_format Output data format
     * @param input_dtype Input data type
     * @param output_dtype Output data type
     * @return
     */
    void set_ai2d_dtype(ai2d_format input_format,
                        ai2d_format output_format,
                        typecode_t input_dtype,
                        typecode_t output_dtype);

    /**
     * @brief Set cropping parameters
     * @param x Crop start X coordinate
     * @param y Crop start Y coordinate
     * @param width Crop width
     * @param height Crop height
     * @return
     */
    void set_crop(size_t x, size_t y, size_t width, size_t height);

    /**
     * @brief Set resizing parameters
     * @param interp_method Interpolation method
     * @param interp_mode Interpolation mode
     * @return
     */
    void set_resize(ai2d_interp_method interp_method,
                    ai2d_interp_mode interp_mode);

    /**
     * @brief Set padding parameters
     * @param pad Padding amount per dimension
     * @param pad_mode Padding mode
     * @param pad_value Padding value
     * @return
     */
    void set_pad(const vector<int>& pad,
                 ai2d_pad_mode pad_mode,
                 const vector<int>& pad_value);

    /**
     * @brief Set affine transformation parameters
     * @param interp_method Interpolation method
     * @param cord_round Coordinate rounding mode
     * @param bound_ind Boundary index type
     * @param bound_val Boundary value
     * @param bound_smooth Boundary smoothing parameter
     * @param affine_matrix Affine transformation matrix (6 elements)
     * @return
     */
    void set_affine(ai2d_interp_method interp_method,
                    int cord_round,
                    int bound_ind,
                    int bound_val,
                    int bound_smooth,
                    const vector<float>& affine_matrix);

    /**
     * @brief Set shift operation parameters
     * @param shift_val Shift value
     * @return
     */
    void set_shift(int shift_val);

    /**
     * @brief Build the AI2D pipeline
     * @param input_shape Input tensor shape
     * @param output_shape Output tensor shape
     * @return
     */
    void build(const dims_t& input_shape,
               const dims_t& output_shape);

    /**
     * @brief Execute the AI2D pipeline
     * @param input_tensor Input tensor
     * @return
     */
    void run(runtime_tensor& input_tensor, runtime_tensor& output_tensor);

private:
    std::unique_ptr<ai2d_builder> ai2d_builder_;                                                                                        ///< AI2D builder

    ai2d_datatype_t ai2d_data_type_{};                                                                                                   ///< Data type configuration
    ai2d_crop_param_t ai2d_crop_param_{false, 0, 0, 0, 0};                                                                               ///< Crop parameters
    ai2d_resize_param_t ai2d_resize_param_{false, ai2d_interp_method::tf_bilinear, ai2d_interp_mode::half_pixel};                        ///< Resize parameters
    ai2d_pad_param_t ai2d_pad_param_{false, {{0, 0}, {0, 0}, {0, 0}, {0, 0}}, ai2d_pad_mode::constant, {0,0,0}};                         ///< Padding parameters
    ai2d_affine_param_t ai2d_affine_param_{false, ai2d_interp_method::cv2_bilinear, 0, 0, 127, 1, {0.5, 0.1, 0.0, 0.1, 0.5, 0.0}};       ///< Affine parameters
    ai2d_shift_param_t ai2d_shift_param_{false, 0};                                                                                      ///< Shift parameters

    dims_t input_shape_;                          ///< Input tensor shape
    dims_t output_shape_;                         ///< Output tensor shape
};

#endif // AI2D_H
