/*
     Description: 多维度嵌套的数组

     Copyright : All right reserved by ICT

     Author : Zhang Zhonghai
     Date : 2025/12/24
*/
#pragma once

#include <vector>
#include <assert.h>
#include "spdlog/spdlog.h"

using std::vector;

template <class T>
struct Array2D {
    vector<T> data;
    int d1 = 0, d2 = 0;
    int s1 = 0, cap = 0;  // strides for each dimension
    Array2D() {}
    Array2D(int dim1, int dim2) { init(dim1, dim2); }
    void init(int dim1, int dim2) { d1 = dim1; d2 = dim2; s1 = dim2; cap = d1 * d2; data.resize(cap); }
    inline T& operator()(int k1, int k2) { return data[k1 * s1 + k2]; }
#define _Foreach2D(array, valName, codes) \
    for (auto& valName : array.data) {    \
        codes;                            \
    }

#define _Foreach2DK(array, valName, codes) \
    do {                                   \
        int k1 = 0, k2 = 0, kg = 0;        \
        for (auto& valName : array.data) { \
            codes;                         \
            ++kg;                          \
            k1 = kg / (array.s1);          \
            k2 = kg % (array.s1);          \
        }                                  \
    } while (0)
};

// 不能是bool
template <class T>
struct Array3D {
    vector<T> data;
    int d1 = 0, d2 = 0, d3 = 0;
    int s1 = 0, s2 = 0, cap = 0;  // strides for each dimension
    Array3D() {}
    Array3D(int dim1, int dim2, int dim3) { init(dim1, dim2, dim3); }
    void init(int dim1, int dim2, int dim3) {
        d1 = dim1; d2 = dim2; d3 = dim3;
        s1 = dim2 * dim3; s2 = dim3; cap = d1 * d2 * d3;
        data.resize(cap);
    }
    inline T& operator()(int k1, int k2, int k3) {
        return data[k1 * s1 + k2 * s2 + k3];
    }

#define _Foreach3D(array, valName, codes) \
    for (auto& valName : array.data) {    \
        codes;                            \
    }

#define _Foreach3DK(array, valName, codes)    \
    do {                                      \
        int k1 = 0, k2 = 0, k3 = 0, kg = 0;   \
        for (auto& valName : array.data) {    \
            codes;                            \
            ++kg;                             \
            const int mod1 = kg % (array.s1); \
            k1 = kg / (array.s1);             \
            k2 = mod1 / (array.s2);           \
            k3 = mod1 % (array.s2);           \
        }                                     \
    } while (0)
};

template <class T>
struct Array4D {
    vector<T> data;
    int d1 = 0, d2 = 0, d3 = 0, d4;
    int s1 = 0, s2 = 0, s3 = 0, cap = 0;  // strides for each dimension
    Array4D() {}
    Array4D(int dim1, int dim2, int dim3, int dim4) { init(dim1, dim2, dim3, dim4); }
    void init(int dim1, int dim2, int dim3, int dim4) {
        d1 = dim1; d2 = dim2; d3 = dim3; d4 = dim4;
        s1 = d2 * d3 * d4; s2 = d3 * d4; s3 = d4; cap = d1 * d2 * d3 * d4;
        data.resize(cap);
    }
    inline T& operator()(int k1, int k2, int k3, int k4) { return data[k1 * s1 + k2 * s2 + k3 * s3 + k4]; }

#define _Foreach4D(array, valName, codes) \
    for (auto& valName : array.data) {    \
        codes;                            \
    }

#define _Foreach4DK(array, valName, codes)          \
    do {                                            \
        int k1 = 0, k2 = 0, k3 = 0, k4 = 0, kg = 0; \
        for (auto& valName : array.data) {          \
            codes;                                  \
            ++kg;                                   \
            const int mod1 = kg % (array.s1);       \
            const int mod2 = mod1 % (array.s2);     \
            k1 = kg / (array.s1);                   \
            k2 = mod1 / (array.s2);                 \
            k3 = mod2 / (array.s3);                 \
            k4 = mod2 % (array.s3);                 \
        }                                           \
    } while (0)
};

// 类似一个tensor，性能太低，不用了
template <class T>
struct NestedArray {
    vector<T> data;
    vector<int> dimensions;
    vector<int> dim_offset;

    NestedArray() { }

    template <typename... Args>
    NestedArray(Args... dims) {
        init(dims...);
    }

    template <typename... Args>
    void init(Args... dims) {
        (dimensions.emplace_back(std::forward<Args>(dims)), ...);
        // spdlog::info("dimensions: {}", dimensions.size());
        //  for (auto& val : dimensions) spdlog::info("dim: {}", val);
        dim_offset.resize(dimensions.size(), 1);
        for (int i = dimensions.size() - 2; i >= 0; --i) {
            dim_offset[i] = dim_offset[i + 1] * dimensions[i + 1];
        }
        data.resize(dimensions[0] * dim_offset[0]);
        // for (int i = 0; i < data.size(); ++i) data[i] = i;
        // for (auto& val : dim_offset) spdlog::info("dim offset: {}", val);
    }

    // 根据索引位置获取数据
    template <typename... Args>
    T& get(Args... keys) {
        vector<int> keyArr;
        (keyArr.emplace_back(std::forward<Args>(keys)), ...);
        assert(keyArr.size() == dimensions.size());
        int idx = 0;
        for (int i = 0; i < keyArr.size(); ++i) {
            idx += keyArr[i] * dim_offset[i];
        }
        return data[idx];
    }

    // 根据关键字，在对应位置插入数据
    template <typename... Args>
    void put(T value, Args... keys) {
        vector<int> keyArr;
        (keyArr.emplace_back(std::forward<Args>(keys)), ...);
        assert(keyArr.size() == dimensions.size());
        int idx = 0;
        for (int i = 0; i < keyArr.size(); ++i) {
            idx += keyArr[i] * dim_offset[i];
        }
        data[idx] = value;
    }
};