/*
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 model weight base class
*/

#include "model_weight_base.h"

#include <iostream>
#include <utility>
#include <sched.h>
#include <numa.h>
#include <cstdlib>
#include <set>
#include <algorithm>
#include <sstream>
#include <cctype>
#include <regex>
#include <pthread.h>
#include "utils.h"
#include "tp_method.h"
#include "cpu_utils.h"

namespace cpu_inference {

// Model-type -> (weight-key -> store-method)
const std::unordered_map<std::string, std::unordered_map<std::string, StoreType> > model_weight_store = {
    {"qwen2", qwen2_split_method},
    {"qwen3_moe", qwen3_moe_name_store},
};

ModelWeightBase::ModelWeightBase(const std::string& model_type, const std::string& quantization_type, 
    const std::map<int, std::vector<int>>& cpu_affinity){
    this->set_model_type(model_type);
    this->set_quantization_type(quantization_type);
    this->set_cpu_affinity(cpu_affinity);
    this->name_store_map = model_weight_store.at(model_type);
}

bool ModelWeightBase::add_weight(const std::string& weight_name, std::string dtype,
    std::vector<int> shape, void* data)
{
    std::regex pattern(R"((\d+)_(.+))");
    std::smatch match;
    std::string weight_type = "";
    if (std::regex_match(weight_name, match, pattern)) {
        weight_type = match[2].str();         // 第二个捕获组：后缀，如 "gate"
    }
    else {
        std::cerr << "weight type is't valid.\n";
    }
    StoreType weight_store_type = this->name_store_map[weight_type];
    // 计算weight_name对应权重在多numa中的形状
    std::vector<int> cur_numa_list;
    std::vector<int> cur_shape;
    if (weight_store_type == StoreType::FULL_IN_FIRST_NUMA) {
        cur_numa_list.push_back(numa_list[0]);
        cur_shape = shape;
    } else if (weight_store_type == StoreType::FULL_IN_ALL_NUMA) {
        cur_numa_list = numa_list;
        cur_shape = shape;
    }
    else if(weight_store_type == StoreType::ROW_SPLIT){
        cur_numa_list = numa_list;
        cur_shape = split_weight_row(shape);
    }
    return this->add_weight(weight_name, cur_numa_list, cur_shape, dtype, data);
}

bool ModelWeightBase::add_weight(const std::string& weight_name,
    const std::vector<int>& cur_numa_list,
    const std::vector<int>& cur_shape,
    std::string dtype,
    void* data)
{
    auto [weight_it, _] = this->weight_map.insert_or_assign(weight_name, Weight(weight_name, 
        cur_numa_list, cur_shape, dtype, this->quantization_type));
    Weight& weight = weight_it->second;         // 从迭代器中获得weight的引用
    const std::vector<int>& orignal_shape = weight.get_orignal_shape();
    size_t first_dim = orignal_shape[0];
    size_t elem_cnt = weight.compute_num_elements(orignal_shape) / first_dim;
    void* src_ptr = data;
    #pragma omp parallel
    {
        WorkDivider work;
        init_work_divider(&work, cur_numa_list.size());
        SingleNumaWorkRange pstSingleRange;
        divide_work_single_numa(&work, first_dim, &pstSingleRange);
        if(pstSingleRange.work_per_thread > 0){
            if (this->quantization_type == "q8align"){
                int block_cnt_in_row = orignal_shape[2] / QK8_ALIGN;
                void** ptr = (void**)weight.at(work.my_numa);
                void* cur_src_ptr = (void*)((char*)src_ptr + pstSingleRange.begin_thread * elem_cnt * sizeof(float16_t));
                void* dst_ptr_val = (void*)((char*)ptr[0] + pstSingleRange.begin_thread * elem_cnt * sizeof(int8_t));
                void* dst_ptr_d = (void*)((char*)ptr[1] + pstSingleRange.begin_thread * elem_cnt / QK8_ALIGN * 2 * sizeof(float));
                Q8Align::q8align_quantize((float16_t*)cur_src_ptr, (int8_t*)dst_ptr_val, (float*)dst_ptr_d, 0,
                    pstSingleRange.work_per_thread * elem_cnt / QK8_ALIGN, block_cnt_in_row, false);
            }else{
                void* cur_src_ptr = (void*)((char*)src_ptr + pstSingleRange.begin_thread * elem_cnt * sizeof(float16_t));
                void* dst_ptr = weight.at(work.my_numa, pstSingleRange.begin_thread);
                quantize_weight(cur_src_ptr, dst_ptr, pstSingleRange.work_per_thread * elem_cnt);
            }
        }
    }
    return true;
}

const Weight& ModelWeightBase::get_weight(const std::string& weight_name) const
{
    return this->weight_map.at(weight_name);
}


void ModelWeightBase::set_model_type(std::string model_type)
{
    this->model_type = model_type;
}

void ModelWeightBase::set_quantization_type(std::string quantization_type)
{
    this->quantization_type = quantization_type;
}

void ModelWeightBase::set_cpu_affinity(std::map<int, std::vector<int>> cpu_affinity)
{
    this->cpu_affinity = cpu_affinity;
    for(auto it = cpu_affinity.begin(); it != cpu_affinity.end(); ++it) {
        this->numa_list.push_back(it->first);
    }
}


void ModelWeightBase::quantize_weight(const void* const src_data, void* dst_data, size_t element_cnt)
{
    if(this->quantization_type == "fp32") {
        QuantizationBase<float16_t>::dequantize(
            static_cast<const float16_t*>(src_data), static_cast<float32_t*>(dst_data), element_cnt, "fp32");
    } else if(this->quantization_type == "fp16") {
        memcpy(dst_data, src_data, element_cnt * sizeof(float16_t));
    } else if(this->quantization_type == "q8_0" || this->quantization_type == "q4_0") {
        if (this->quantization_type == "q8_0") {
            this->template quantize_weight_impl<block_q8_0>(src_data, dst_data, element_cnt, QK8_0);
        } else if (this->quantization_type == "q4_0") {
            this->template quantize_weight_impl<block_q4_0>(src_data, dst_data, element_cnt, QK4_0);
        }
    } else {
        std::cerr << "Unsupported quantization type: " << this->quantization_type << std::endl;
    }
}


std::vector<int> ModelWeightBase::split_weight_row(const std::vector<int>& shape) const
{
    std::vector<int> shape_res = shape;
    const size_t numa_cnt = this->numa_list.size();
    int split_size = (shape[0] + numa_cnt - 1) / numa_cnt;
    shape_res[0] = split_size;
    return shape_res;
}


};