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

int ncnn::conv2d_depthwise_w8a8ch_kernel(
    const Mat& bottom_blob, Mat& top_blob,
    const Mat& weight_data, const Mat& bias_data,
    const Mat& input_scales, const Mat& weight_scales,
    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 == group)

    // Unlike plain Convolution_arm (where bottom/top int8 scales are a single
    // per-tensor scalar), ConvolutionDepthWise_arm expects ALL THREE scale Mats
    // sized `group` (== C here). ncnn's fields are *quantization* multipliers
    // (int8 = round(float*scale)); this baseline's weight_scales/input_scales
    // are *dequantization* multipliers (real = int8*scale) — invert them.
    Mat weight_int8_scales;
    weight_int8_scales.create(C, (size_t)4u);
    if (weight_int8_scales.empty()) return -1;
    const float* ws = (const float*)weight_scales.data;
    for (int g = 0; g < C; ++g)
        weight_int8_scales[g] = 1.0f / ws[g];

    Mat bottom_int8_scales;
    bottom_int8_scales.create(C, (size_t)4u);
    if (bottom_int8_scales.empty()) return -1;
    const float* is_ = (const float*)input_scales.data;
    for (int g = 0; g < C; ++g)
        bottom_int8_scales[g] = 1.0f / is_[g];

    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                = 1;
    conv.weight_data_size          = static_cast<int>(weight_data.total());
    conv.group                     = C;
    // Must be exactly 1 or 101 (not 2/102) to select the per-channel (size=group)
    // weight-scale branch — {2,102} would instead load a single scalar weight
    // scale broadcast to every channel, wrong for our per-channel weight_scales.
    // 1 (not 101) additionally keeps ncnn on the plain-dequantize branch
    // (out = acc*scale+bias, no rounding/clip) instead of requantizing to int8 —
    // w8a8ch only quantizes inputs, the task doesn't require a quantized output.
    conv.int8_scale_term           = 1;
    conv.activation_type           = 0;
    conv.activation_params         = Mat();
    conv.dynamic_weight             = 0;
    conv.weight_data                = const_cast<Mat&>(weight_data);
    conv.bias_data                  = const_cast<Mat&>(bias_data);
    conv.weight_data_int8_scales     = weight_int8_scales;
    conv.bottom_blob_int8_scales     = bottom_int8_scales;

    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;
}