#include <iostream>
#include <vector>
using namespace std;

vector<vector<int>> conv2d_naive(vector<vector<int>> img, vector<vector<int>> kernel, int stride)
{
    int img_h = img.size();
    int img_w = img[0].size();
    int k_h = kernel.size();
    int k_w = kernel[0].size();
    int output_h = ((img_h - k_h) / stride) + 1;
    int output_w = ((img_w - k_w) / stride) + 1;
    vector<vector<int>> output;
    for(int i = k_h; i < img_h - k_h; i+=stride)
    {
        vector<int> temp;
        for(int j = k_w; j < img_w - k_w; j+=stride)
        {
            int sum = 0;
            for(int i1 = 0; i1 < k_h; i1++)
            {
                for(int j1 = 0; j1 < k_w; j1++)
                {
                    sum += img[i - k_h + i1][j - k_w + j1] * kernel[i1][j1];
                }
            }
            temp.append(sum);
        }
        output.append(temp);
    }

    return output;
}

vector<vector<int>> conv2d_optimized(vector<int> img, vector<int> kernel, int img_h, int img_w,
                                    int k_h, int k_w, int stride)
{
    int output_h = ((img_h - k_h) / stride) + 1;
    int output_w = ((img_w - k_w) / stride) + 1;

    vector<vector<int>> output;
    vector<int> temp;
    int count = 0;
    for(int i = 0; i < img_h * img_w; i+=stride)
    {
        int sum = 0;
        int start = i;
        for(int j = 0; j < k_h * k_w; j++)
        {
            if(j != 0 && j % k_w == 0)
            {
                start += img_w - k_w;
            }

            sum += img[start] * kernel[j];
        }

        temp.append(sum)
        if(temp.size() == output_w)
        {
            output.append(temp);
            temp.clear();
        }
    }

    return output;
}