// Created by xiangzhou on 11/9/23.
//
#pragma once

#include "C_Tensor.h"

#ifdef __cplusplus
extern "C" {
#endif

//----------------------------------------------------------------------------
//      AveragePool
//------------------------------------------------------------------------
struct AveragePoolAttribute{
    // user
    const char* name_;
    uint32_t top_pad_;
    uint32_t bottom_pad_;
    uint32_t left_pad_;
    uint32_t right_pad_;
    uint32_t kernel_h_;
    uint32_t kernel_w_;
    uint32_t stride_h_;
    uint32_t stride_w_;
    int32_t count_include_pad_;  // flag to include padding

};
typedef struct AveragePoolAttribute AveragePoolAttribute;
//----------------------------------------------------------------------------
//      GlobalAveragePool
//------------------------------------------------------------------------
typedef struct Attribute GlobalAveragePoolAttribute;


//----------------------------------------------------------------------------
//      GlobalMaxPool
//------------------------------------------------------------------------
typedef struct Attribute GlobalMaxPoolAttribute;

//----------------------------------------------------------------------------
//      GlobalLpPoolAttribute
//------------------------------------------------------------------------
struct GlobalLpPoolAttribute{
    // user
    const char* name_;
    int p_;
};
typedef struct GlobalLpPoolAttribute GlobalLpPoolAttribute;
//----------------------------------------------------------------------------
//      LpPool
//------------------------------------------------------------------------
struct LpPoolAttribute {
    // user
    const char* name_;
    const char* auto_pad_;
    int ceil_mode_ ;
    C_Shape dilations_;
    C_Shape kernel_shape_;
    int p_ ;
    C_Shape pads_;
    C_Shape strides_;
};
typedef struct LpPoolAttribute LpPoolAttribute;


//----------------------------------------------------------------------------
//      Pooling
//------------------------------------------------------------------------
struct PoolingAttribute {
    // user
    const char* name_;
    C_Shape kernels_;
    C_Shape pads_;
    C_Shape strides_;

    uint32_t top_pad_;
    uint32_t bottom_pad_;
    uint32_t left_pad_;
    uint32_t right_pad_;
    uint32_t kernel_h_;
    uint32_t kernel_w_;
    uint32_t stride_h_;
    uint32_t stride_w_;

    // library
    int32_t in_c_, in_d_, in_h_, in_w_;
    int32_t out_d_, out_h_, out_w_;
    bool is_3d_;
};
typedef struct PoolingAttribute PoolingAttribute;
typedef struct PoolingAttribute MaxPoolAttribute;
typedef struct PoolingAttribute MaxUnpoolAttribute;

//----------------------------------------------------------------------------
//      MaxRoiPool
//------------------------------------------------------------------------
struct MaxRoiPoolAttribute {
    // user
    const char* name_;
    uint32_t pooled_shape_h_;
    uint32_t pooled_shape_w_;
    float spatial_scale_;
    int32_t roi_[5];  // this is just one roi, Don't support multiple roi yet
};
typedef struct MaxRoiPoolAttribute MaxRoiPoolAttribute;
//----------------------------------------------------------------------------
//      Pad
//------------------------------------------------------------------------
struct PadAttribute {
    // user
    const char* name_;
    C_Shape pads_;
};
typedef struct PadAttribute PadAttribute;
CREATE_DECL_SETUP_FUN(MaxPool)

CREATE_DECL_RUN_FUN(AveragePool)
CREATE_DECL_RUN_FUN(GlobalAveragePool)
CREATE_DECL_RUN_FUN(LpPool)
CREATE_DECL_RUN_FUN(GlobalLpPool)
CREATE_DECL_RUN_FUN(GlobalMaxPool)
CREATE_DECL_RUN_FUN(MaxPool)
CREATE_DECL_RUN_FUN(MaxRoiPool)
CREATE_DECL_RUN_FUN(MaxUnpool)
CREATE_DECL_RUN_FUN(Pad)

#ifdef __cplusplus
}
#endif