#include <vector>
#include <iostream>
#include <iomanip> // 用于格式化输出

using namespace std;

int main() {
    // --------------------- 1. 初始化输入图像 ---------------------
    // 假设输入图像形状为 [3, 896, 896]（通道优先：C×H×W）
    const int C = 3;       // 通道数（RGB）
    const int H = 896;     // 图像高度
    const int W = 896;     // 图像宽度
    vector<float> inp_raw(C * H * W, 1.0f); // 示例：全1数据（实际需替换为真实图像数据）

    // --------------------- 2. 分块参数计算 ---------------------
    const int num_patches_h = 14; // Height方向分块数
    const int num_patches_w = 14; // Width方向分块数
    const int patch_h = H / num_patches_h; // 单个Patch的高度（64）
    const int patch_w = W / num_patches_w; // 单个Patch的宽度（64）

    // --------------------- 3. 提取每个Patch的统计特征（示例：均值） ---------------------
    // patch_means[ph][pw][c] 存储第ph行、第pw列Patch的第c通道均值
    vector<vector<vector<float>>> patch_means(
        num_patches_h, 
        vector<vector<float>>(num_patches_w, vector<float>(C))
    );

    for (int c = 0; c < C; ++c) {
        for (int ph = 0; ph < num_patches_h; ++ph) {
            for (int pw = 0; pw < num_patches_w; ++pw) {
                float sum = 0.0f;
                int count = 0;
                // 遍历Patch内所有像素
                for (int h = ph * patch_h; h < (ph + 1) * patch_h; ++h) {
                    for (int w = pw * patch_w; w < (pw + 1) * patch_w; ++w) {
                        int idx = c * H * W + h * W + w; // 计算inp_raw中像素的索引
                        sum += inp_raw[idx];
                        count++;
                    }
                }
                patch_means[ph][pw][c] = sum / count; // 计算该Patch的通道均值
            }
        }
    }

    // --------------------- 4. 构建目标维度：4096×2×3×14×14 ---------------------
    const int batch_size = 4096;   // Batch大小
    const int num_views = 2;       // 视角/分支数（如时间步、参考帧与当前帧）
    // 目标数组形状：[batch_size, num_views, C, num_patches_h, num_patches_w]
    vector<float> target(
        batch_size * num_views * C * num_patches_h * num_patches_w
    );

    // 填充目标数组：每个Batch元素、视角、通道、Patch位置均使用对应Patch的均值
    for (int n = 0; n < batch_size; ++n) {
        for (int m = 0; m < num_views; ++m) {
            for (int c = 0; c < C; ++c) {
                for (int ph = 0; ph < num_patches_h; ++ph) {
                    for (int pw = 0; pw < num_patches_w; ++pw) {
                        // 计算目标数组的一维索引
                        int target_idx = n * num_views * C * num_patches_h * num_patches_w 
                                        + m * C * num_patches_h * num_patches_w 
                                        + c * num_patches_h * num_patches_w 
                                        + ph * num_patches_w + pw;
                        target[target_idx] = patch_means[ph][pw][c];
                    }
                }
            }
        }
    }

    // --------------------- 5. 验证输出（可选） ---------------------
    cout << "目标数组总元素数: " << target.size() << "（预期: " 
         << batch_size * num_views * C * num_patches_h * num_patches_w << "）" << endl;
    
    // 打印第一个Batch、第一个视角、第一个通道的所有Patch值（前几个）
    cout << "\n第一个Batch、第一个视角、第一个通道的Patch均值（前20个）:" << endl;
    for (int ph = 0; ph < min(5, num_patches_h); ++ph) { // 只打印前5个Patch的行
        for (int pw = 0; pw < min(5, num_patches_w); ++pw) { // 只打印前5个Patch的列
            int idx = 0 * num_views * C * num_patches_h * num_patches_w 
                      + 0 * C * num_patches_h * num_patches_w 
                      + 0 * num_patches_h * num_patches_w 
                      + ph * num_patches_w + pw;
            cout << fixed << setprecision(2) << target[idx] << "\t";
        }
        cout << endl;
    }

    return 0;
}