/*
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
*/

#include "tensor.h"

#include <cstring>
#include <algorithm>
#include <cctype>
#include <stdexcept>
#include <utility>

#include "quantization/quantization_q4_0.h"
#include "quantization/quantization_q8_0.h"
#include "quantization/quantization_q8align.h"

namespace {

bool has_negative(const std::vector<int>& v)
{
    for (int x : v) {
        if (x < 0)
            return true;
    }
    return false;
}

const std::unordered_map<int, std::vector<int>> compute_numa_shape(
    const std::vector<int>& numa_list, const std::vector<int>& shape){
    std::unordered_map<int, std::vector<int>> numa_shape;
    for(int i = 0;i < numa_list.size();i++){
        int numa_id = numa_list[i];
        numa_shape[numa_id] = shape;
    }
    return numa_shape;
}

std::vector<int> compute_quant_shape(const std::vector<int>& shape, const std::string& quant_type){
    std::vector<int> shape_res = shape;
    int dim = shape_res.size();
    if(quant_type == "q4_0"){
        shape_res[dim - 1] =  shape_res[dim - 1] / QK4_0 * sizeof(cpu_inference::block_q4_0);
    }
    else if(quant_type == "q8_0"){
        shape_res[dim - 1] =  shape_res[dim - 1] / QK8_0 * sizeof(cpu_inference::block_q8_0);
    }
    return shape_res;
}

std::string compute_dtype(const std::string& dtype, const std::string& quant_type){
    // 如果quant_type是q4_0、q8_0、fp16，那么只根据quant_type来判断dtype_res。否则，返回dtype
    std::string dtype_res("");
    if(quant_type == "q4_0"){
        dtype_res = "int8";
    }
    else if(quant_type == "q8_0"){
        dtype_res = "int8";
    }
    else if (quant_type == "fp16") {
        dtype_res = "float16";
    }
    else if (quant_type == "q8align") {
        dtype_res = "q8align";
    }
    else {
        dtype_res = dtype;
    }
    return dtype_res;
}

}  // namespace

namespace cpu_inference {

Tensor::Tensor(const std::string& weight_name,
    const std::unordered_map<int, size_t>& numa_size_map, std::string dtype)
    : weight_name(weight_name), numa_size_map(numa_size_map), dtype(dtype)
{
    this->init_element_size();
    if (dtype == "q8align"){
        for (auto it = numa_size_map.begin(); it != numa_size_map.end(); ++it) {
            void **tmp = (void**)numa_alloc_onnode(2 * sizeof(void*), it->first);
            tmp[0] = (void*)numa_alloc_onnode(it->second * this->element_size, it->first);
            tmp[1] = (void*)numa_alloc_onnode(it->second * this->element_size / QK8_ALIGN * 2 * sizeof(float), it->first);
            memory_map[it->first] = (void*)tmp;
        }
    } else {
        for (auto it = numa_size_map.begin(); it != numa_size_map.end(); ++it) {
            memory_map[it->first] =
                static_cast<void*>(numa_alloc_onnode(it->second * this->element_size, it->first));
        }
    }
}

Tensor::Tensor(const std::string& weight_name,
    const std::unordered_map<int, size_t>& numa_size_map, std::string dtype, void* data)
    : weight_name(weight_name), numa_size_map(numa_size_map), dtype(dtype)
{
    this->init_element_size();
    const char* src = static_cast<const char*>(data);
    size_t offset_bytes = 0;
    if (dtype == "q8align"){
        throw std::runtime_error(std::string("Tensor not supported q8align with data"));
    }
    for (auto it = numa_size_map.begin(); it != numa_size_map.end(); ++it) {
        const int numa_node = it->first;
        const size_t elem_count = it->second;
        const size_t bytes = elem_count * this->element_size;
        void* dst = static_cast<void*>(numa_alloc_onnode(bytes, numa_node));
        memory_map[numa_node] = dst;
        if (src != nullptr) {
            memcpy(dst, src + offset_bytes, bytes);
            offset_bytes += bytes;
        }
    }
}

Tensor::Tensor(const std::string& weight_name,
    const std::unordered_map<int, std::vector<int>>& numa_shape_map,
    std::string dtype) :
    Tensor(weight_name, compute_numa_num_elements(numa_shape_map), dtype)
{
    for (auto it = numa_shape_map.begin(); it != numa_shape_map.end(); ++it) {
        const int numa_node = it->first;
        const std::vector<int>& shape = it->second;
        this->reshape(numa_node, shape);
    }
}

Tensor::Tensor(const std::string& weight_name,
    const std::unordered_map<int, std::vector<int>>& numa_shape_map,
    std::string dtype, void* data) :
    Tensor(weight_name, compute_numa_num_elements(numa_shape_map), dtype, data)
{
    for (auto it = numa_shape_map.begin(); it != numa_shape_map.end(); ++it) {
        const int numa_node = it->first;
        const std::vector<int>& shape = it->second;
        this->reshape(numa_node, shape);
    }
}

Tensor::Tensor(const std::string& weight_name, std::vector<int> numa_list,
    std::vector<int> shape, std::string dtype):
    Tensor(weight_name, compute_numa_shape(numa_list, shape), dtype)
{}

Tensor::Tensor(const std::string& weight_name, std::vector<int> numa_list,
    std::vector<int> shape, std::string dtype, std::string quant_type):
    Tensor(weight_name, numa_list, compute_quant_shape(shape, quant_type), 
    compute_dtype(dtype, quant_type))
{
    this->orignal_shape = shape;
}

Tensor::Tensor(Tensor&& other)
    : weight_name(std::move(other.weight_name)),
      dtype(std::move(other.dtype)),
      element_size(other.element_size),
      numa_size_map(std::move(other.numa_size_map)),
      memory_map(std::move(other.memory_map)),
      shape(std::move(other.shape)),
      orignal_shape(std::move(other.orignal_shape))
{
    // 确保移动后的对象可以安全销毁
    other.element_size = 0;
    other.numa_size_map.clear();
    other.memory_map.clear();
    other.shape.clear();
    other.orignal_shape.clear();
}

Tensor& Tensor::operator=(Tensor&& other)
{
    if (this == &other) {
        return *this;
    }

    // 释放当前NUMA分配
    if (dtype == "q8align"){
        for (auto it = memory_map.begin(); it != memory_map.end(); ++it) {
            const std::unordered_map<int, size_t>::iterator sz_it = numa_size_map.find(it->first);
            if (sz_it != numa_size_map.end()) {
                void** tmp = (void**)it->second;
                numa_free(tmp[0], sz_it->second * this->element_size);
                numa_free(tmp[1], sz_it->second * this->element_size / QK8_ALIGN * 2 * sizeof(float));
                numa_free(tmp, 2 * sizeof(void*));
            }
        }
    }else{
        for (auto it = memory_map.begin(); it != memory_map.end(); ++it) {
            const auto sz_it = numa_size_map.find(it->first);
            if (sz_it != numa_size_map.end() && it->second != nullptr) {
                numa_free(it->second, sz_it->second * this->element_size);
            }
        }
    }

    // 移动成员属性
    weight_name = std::move(other.weight_name);
    dtype = std::move(other.dtype);
    element_size = other.element_size;
    numa_size_map = std::move(other.numa_size_map);
    memory_map = std::move(other.memory_map);
    shape = std::move(other.shape);
    orignal_shape = std::move(other.orignal_shape);

    // 确保移动后的对象可以安全销毁
    other.element_size = 0;
    other.numa_size_map.clear();
    other.memory_map.clear();
    other.shape.clear();
    other.orignal_shape.clear();
    return *this;
}

Tensor::~Tensor()
{
    if (dtype == "q8align"){
        for (auto it = memory_map.begin(); it != memory_map.end(); ++it) {
            const std::unordered_map<int, size_t>::iterator sz_it = numa_size_map.find(it->first);
            if (sz_it != numa_size_map.end()) {
                void** tmp = (void**)it->second;
                numa_free(tmp[0], sz_it->second * this->element_size);
                numa_free(tmp[1], sz_it->second * this->element_size / QK8_ALIGN * 2 * sizeof(float));
                numa_free(tmp, 2 * sizeof(void*));
            }
        }
    }else{
        for (auto it = memory_map.begin(); it != memory_map.end(); ++it) {
            const std::unordered_map<int, size_t>::iterator sz_it = numa_size_map.find(it->first);
            if (sz_it != numa_size_map.end()) {
                numa_free(it->second, sz_it->second * this->element_size);
            }
        }
    }
}

bool Tensor::write(int numa_node, const void* data, size_t size)
{
    if (memory_map.find(numa_node) == memory_map.end()) {
        return false;
    }
    const std::unordered_map<int, size_t>::iterator it = numa_size_map.find(numa_node);
    if (it != numa_size_map.end() && size > it->second * this->element_size) {
        return false;
    }
    memcpy(memory_map[numa_node], data, size);
    return true;
}

bool Tensor::reshape(unsigned int numa_node, const std::vector<int> shape)
{
    if(shape.size() == 0) {     // shape不能为空
        return false;
    }
    if (has_negative(shape)) {  // 不能有负数
        return false;
    }
    // 检查numa_node是否存在
    if (this->numa_size_map.find(numa_node) == this->numa_size_map.end()) {
        return false;
    }
    // 检查shape的字节大小是否与numa_size_map[numa_node]一致
    size_t total = 1;
    for (const int& dim : shape) {
        total *= static_cast<size_t>(dim);
    }
    // numa_size_map stores element counts, not bytes.
    if (static_cast<size_t>(total) != numa_size_map[numa_node]) {
        return false;
    }
    this->shape = shape;
    return true;
}

size_t Tensor::get_element_size() const
{
    return this->element_size;
}

const std::vector<int>& Tensor::get_shape() const
{
    return this->shape;
}
const std::vector<int>& Tensor::get_orignal_shape() const{
    return this->orignal_shape;
}

const void* Tensor::at(int numa_node, const std::vector<int>& indices) const
{
    // numa检查
    const std::unordered_map<int, void*>::const_iterator mem_it = memory_map.find(numa_node);
    if (mem_it == memory_map.end()) {
        return nullptr;
    }
    // 空检查
    if (this->shape.empty()) {
        return nullptr;
    }
    // 负数检查
    if (has_negative(indices)) {
        return nullptr;
    }
    // 维度检查
    if (indices.size() > shape.size()) {
        return nullptr;
    }
    // 边界检查
    for (size_t i = 0; i < indices.size(); ++i) {
        if (indices[i] >= shape[i]) {
            return nullptr;
        }
    }

    // 行优先方式，检查偏移量
    std::vector<size_t> strides_(shape.size());
    if (!shape.empty()) {
        strides_.back() = 1;
        for (int i = shape.size() - 2; i >= 0; --i) {
            strides_[i] = strides_[i + 1] * static_cast<size_t>(shape[i + 1]);
        }
    }
    size_t offset_elems = 0;
    for (size_t i = 0; i < indices.size(); ++i) {
        offset_elems += static_cast<size_t>(indices[i]) * strides_[i];
    }
    const char* base = static_cast<const char*>(mem_it->second);
    const size_t byte_offset = offset_elems * this->element_size;
    return static_cast<const void*>(base + byte_offset);
}

void* Tensor::at(int numa_node, const std::vector<int>& indices)
{
    return const_cast<void*>(static_cast<const Tensor*>(this)->at(numa_node, indices));
}

void Tensor::init_element_size()
{
    this->element_size = dtype_nbytes(this->dtype);
    if (this->element_size == 0) {
        throw std::invalid_argument("Invalid dtype: " + this->dtype);
    }
}

std::string Tensor::normalize_dtype(std::string dtype) const
{
    if (dtype == "half") {
        return "float16";
    }
    dtype.erase(std::remove_if(dtype.begin(), dtype.end(), [](unsigned char c) { return std::isspace(c); }),
                dtype.end());
    std::transform(dtype.begin(), dtype.end(), dtype.begin(),
                    [](unsigned char c) { return static_cast<char>(std::tolower(c)); });
    return dtype;
}

size_t Tensor::dtype_nbytes(const std::string& dtype_raw) const
{
    const std::string dtype = normalize_dtype(dtype_raw);
    const auto it = dtype_nbytes_map.find(dtype);
    if (it != dtype_nbytes_map.end()) {
        return it->second;
    }
    return 0;
}

size_t Tensor::compute_num_elements(const std::vector<int>& shape)
{
    size_t num_elements = 1;
    for (const int& dim : shape) {
        num_elements *= static_cast<size_t>(dim);
    }
    return num_elements;
}

std::unordered_map<int, size_t> Tensor::compute_numa_num_elements(
    const std::unordered_map<int, std::vector<int>>& numa_shape_map)
{
    std::unordered_map<int, size_t> numa_size_map;
    for(auto it = numa_shape_map.begin(); it != numa_shape_map.end(); ++it) {
        int numa_id = it->first;
        size_t num_size = compute_num_elements(it->second);
        numa_size_map[numa_id] = num_size;
    }
    return numa_size_map;
}

}  // namespace cpu_inference