// File: kernel.cpp
#include "conv2d_depthwise_contract.h"
#include "convolutiondepthwise_arm.h"
#include "mat.h"
#include "option.h"
#include <cstring>

int ncnn::conv2d_depthwise_kernel(
    const Mat& bottom_blob, Mat& top_blob,
    const Mat& weight_data, const Mat& bias_data,
    int kernel_h, int kernel_w,
    int stride_h, int stride_w,
    int dilation_h, int dilation_w,
    int pad,
    const Option& opt)
{
    const int C = top_blob.c;  // pre-set by binding.cpp (== input channels)

    // ncnn's pack1 convdw3x3s1_neon reads 12 floats per iteration (for 8 outputs),
    // potentially overreading past the last channel's allocation when the bordered blob
    // lands near a page boundary.  Use pack4 (4-channel interleaved) when possible to
    // take the safe convdw3x3s1_pack4_neon path instead.
    if (C % 4 == 0) {
        Option opt4 = opt;
        opt4.use_packing_layout = true;

        Mat bottom_pack4;
        convert_packing(bottom_blob, bottom_pack4, 4, opt4);

        ConvolutionDepthWise_arm conv;
        conv.num_output        = C;
        conv.kernel_h          = kernel_h;   conv.kernel_w    = kernel_w;
        conv.stride_h          = stride_h;   conv.stride_w    = stride_w;
        conv.dilation_h        = dilation_h; conv.dilation_w  = dilation_w;
        conv.pad_top           = pad;        conv.pad_bottom  = pad;
        conv.pad_left          = pad;        conv.pad_right   = pad;
        conv.pad_value         = 0.f;
        conv.bias_term         = (!bias_data.empty() && bias_data.total() > 0) ? 1 : 0;
        conv.weight_data_size  = static_cast<int>(weight_data.total());
        conv.group             = C;
        conv.int8_scale_term   = 0;
        conv.activation_type   = 0;
        conv.activation_params = Mat();
        conv.dynamic_weight    = 0;
        conv.weight_data       = const_cast<Mat&>(weight_data);
        if (conv.bias_term) conv.bias_data = const_cast<Mat&>(bias_data);

        if (conv.create_pipeline(opt4) != 0) return -1;
        Mat top_pack4;
        if (conv.forward(bottom_pack4, top_pack4, opt4) != 0) return -1;

        Mat local_top;
        convert_packing(top_pack4, local_top, 1, opt);

        for (int c = 0; c < C; ++c)
            std::memcpy((float*)top_blob.channel(c), (const float*)local_top.channel(c),
                        top_blob.h * top_blob.w * sizeof(float));
        return 0;
    }

    // Fallback pack1 path for C not divisible by 4.
    ConvolutionDepthWise_arm conv;
    conv.num_output        = C;
    conv.kernel_h          = kernel_h;   conv.kernel_w    = kernel_w;
    conv.stride_h          = stride_h;   conv.stride_w    = stride_w;
    conv.dilation_h        = dilation_h; conv.dilation_w  = dilation_w;
    conv.pad_top           = pad;        conv.pad_bottom  = pad;
    conv.pad_left          = pad;        conv.pad_right   = pad;
    conv.pad_value         = 0.f;
    conv.bias_term         = (!bias_data.empty() && bias_data.total() > 0) ? 1 : 0;
    conv.weight_data_size  = static_cast<int>(weight_data.total());
    conv.group             = C;
    conv.int8_scale_term   = 0;
    conv.activation_type   = 0;
    conv.activation_params = Mat();
    conv.dynamic_weight    = 0;
    conv.weight_data       = const_cast<Mat&>(weight_data);
    if (conv.bias_term) conv.bias_data = const_cast<Mat&>(bias_data);

    if (conv.create_pipeline(opt) != 0) return -1;
    Mat local_top;
    if (conv.forward(bottom_blob, local_top, opt) != 0) return -1;

    for (int c = 0; c < C; ++c)
        std::memcpy((float*)top_blob.channel(c), (const float*)local_top.channel(c),
                    top_blob.h * top_blob.w * sizeof(float));
    return 0;
}