/*
     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<vector<T>> data;
    Array2D() { }
    Array2D(int dim1, int dim2) { init(dim1, dim2); }
    void init(int dim1, int dim2) { data.resize(dim1); for (auto& v : data) v.resize(dim2); }
    inline T& get(int k1, int k2) { return data[k1][k2]; }
    // 根据关键字，在对应位置插入数据
    inline void put(const T& value, int k1, int k2) { data[k1][k2] = value; }
    inline vector<T>& operator[](size_t idx) { return data[idx]; }
    inline const vector<T>& operator[](size_t idx) const { return data[idx]; }
#define _Foreach2D(array, valName, codes) \
    for (auto& arr1 : array.data) {       \
        for (auto& valName : arr1) {      \
            codes;                        \
        }                                 \
    }
#define _Foreach2DK(array, valName, codes) \
    do {                                   \
        int k1 = 0;                        \
        for (auto& arr1 : array.data) {    \
            int k2 = 0;                    \
            for (auto& valName : arr1) {   \
                codes;                     \
                ++k2;                      \
            }                              \
            ++k1;                          \
        }                                  \
    } while (0)
};

template <class T>
struct Array3D {
    vector<vector<vector<T>>> data;
    Array3D() {}
    Array3D(int dim1, int dim2, int dim3) { init(dim1, dim2, dim3); }
    void init(int dim1, int dim2, int dim3) {
        data.resize(dim1);
        for (auto& v : data) v.resize(dim2);
        for (auto& v1 : data)
            for (auto& v2 : v1) v2.resize(dim3);
    }
    inline T& get(int k1, int k2, int k3) { return data[k1][k2][k3]; }
    // 根据关键字，在对应位置插入数据
    inline void put(const T& value, int k1, int k2, int k3) { data[k1][k2][k3] = value; }
    inline vector<vector<T>>& operator[](size_t idx) { return data[idx]; }
#define _Foreach3D(array, valName, codes) \
    for (auto& arr1 : array.data) {       \
        for (auto& arr2 : arr1) {         \
            for (auto& valName : arr2) {  \
                codes;                    \
            }                             \
        }                                 \
    }

#define _Foreach3DK(array, valName, codes)   \
    do {                                     \
        int k1 = 0;                          \
        for (auto& arr1 : array.data) {      \
            int k2 = 0;                      \
            for (auto& arr2 : arr1) {        \
                int k3 = 0;                  \
                for (auto& valName : arr2) { \
                    codes;                   \
                    ++k3;                    \
                }                            \
                ++k2;                        \
            }                                \
            ++k1;                            \
        }                                    \
    } while (0)
};

template <class T>
struct Array4D {
    vector<vector<vector<vector<T>>>> data;
    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) {
        data.resize(dim1);
        for (auto& v : data) v.resize(dim2);
        for (auto& v1 : data)
            for (auto& v2 : v1) v2.resize(dim3);
        for (auto& v1 : data)
            for (auto& v2 : v1)
                for (auto& v3 : v2) v3.resize(dim4);
    }
    inline T& get(int k1, int k2, int k3, int k4) { return data[k1][k2][k3][k4]; }
    // 根据关键字，在对应位置插入数据
    inline void put(const T& value, int k1, int k2, int k3, int k4) { data[k1][k2][k3][k4] = value; }
    inline vector<vector<vector<T>>>& operator[](size_t idx) { return data[idx]; }
#define _Foreach4D(array, valName, codes)    \
    for (auto& arr1 : array.data) {          \
        for (auto& arr2 : arr1) {            \
            for (auto& arr3 : arr2) {        \
                for (auto& valName : arr3) { \
                    codes;                   \
                }                            \
            }                                \
        }                                    \
    }

#define _Foreach4DK(array, valName, codes)       \
    do {                                         \
        int k1 = 0;                              \
        for (auto& arr1 : array.data) {          \
            int k2 = 0;                          \
            for (auto& arr2 : arr1) {            \
                int k3 = 0;                      \
                for (auto& arr3 : arr2) {        \
                    int k4 = 0;                  \
                    for (auto& valName : arr3) { \
                        codes;                   \
                        ++k4;                    \
                    }                            \
                    ++k3;                        \
                }                                \
                ++k2;                            \
            }                                    \
            ++k1;                                \
        }                                        \
    } 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;
    }
};