/*
Copyright (c) 2025-2025 Huawei Technologies Co., Ltd.

sysHAX-adapter is licensed under Mulan PSL v2.
You can use this software according to the terms and conditions of the Mulan PSL v2.
You may obtain a copy of Mulan PSL v2 at:
    http://license.coscl.org.cn/MulanPSL2
THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR
IMPLIED, INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY OR FIT FOR A PARTICULAR
PURPOSE.
See the Mulan PSL v2 for more details.
Created: 2026-1-31
Desc: CPU inference tensor
*/

#ifndef TENSOR_BASE_H
#define TENSOR_BASE_H

#include <string>
#include <unordered_map>
#include <vector>
#include <cstdint>
#include <numa.h>

typedef float f32;

namespace cpu_inference {

inline const std::unordered_map<std::string, size_t> dtype_nbytes_map = {
    {"half", 2},
    {"float16", 2},
    {"bfloat16", 2},
    {"float32", 4},
    {"float64", 8},
    {"int8", 1},
    {"uint8", 1},
    {"int16", 2},
    {"uint16", 2},
    {"q8align", 1},

};

class Tensor
{
public:
    Tensor() = default;
    Tensor(const std::string& weight_name, const std::unordered_map<int, size_t>& numa_size_map,
        std::string dtype);
    Tensor(const std::string& weight_name, const std::unordered_map<int, size_t>& numa_size_map,
        std::string dtype, void* data);
    Tensor(const std::string& weight_name, const std::unordered_map<int, std::vector<int>>& numa_shape_map,
        std::string dtype);
    Tensor(const std::string& weight_name, const std::unordered_map<int, std::vector<int>>& numa_shape_map,
        std::string dtype, void* data);
    Tensor(const std::string& weight_name, std::vector<int> numa_list,
        std::vector<int> shape, std::string dtype);
    Tensor(const std::string& weight_name, std::vector<int> numa_list,
        std::vector<int> shape, std::string dtype, std::string quant_type);
    virtual ~Tensor();
    

    // 允许移动构造
    Tensor(Tensor&& other);
    Tensor& operator=(Tensor&& other);

    // 禁用拷贝构造
    Tensor(const Tensor&) = delete;
    Tensor& operator=(const Tensor&) = delete;

    bool write(int numa_node, const void* data, size_t size);
    bool reshape(unsigned int numa_node, const std::vector<int> shape);
    size_t get_element_size() const;
    const std::vector<int>& get_shape() const;
    const std::vector<int>& get_orignal_shape() const;


    // 返回指向子张量的指针
    // 索引按行优先顺序，基于shape进行解释
    // 例如: shape = [3, 5]
    // - at(1, {0})    -> 指向NUMA 1, 第0行的起始位置
    // - at(0, {2, 4}) -> 指向NUMA 0, [2,4]的位置
    const void* at(int numa_node, const std::vector<int>& indices) const;

    template <typename... Index>
    const void* at(int numa_node, Index... idx) const
    {
        return at(numa_node, {static_cast<int>(idx)...});
    }

    void* at(int numa_node, const std::vector<int>& indices);

    template <typename... Index>
    void* at(int numa_node, Index... idx)
    {
        return at(numa_node, {static_cast<int>(idx)...});
    }

    static size_t compute_num_elements(const std::vector<int>& shape);
    static std::unordered_map<int, size_t> compute_numa_num_elements(
        const std::unordered_map<int, std::vector<int>>& numa_shape_map);
    std::unordered_map<int, void*> memory_map;

protected:
    std::string weight_name;
    std::string dtype;
    size_t element_size = 0;
    std::unordered_map<int, size_t> numa_size_map;
    std::vector<int> shape;
    std::vector<int> orignal_shape;
protected:
    void init_element_size();
    std::string normalize_dtype(std::string dtype) const;
    size_t dtype_nbytes(const std::string& dtype_raw) const;
};

};  // namespace cpu_inference

#endif