// test/cpp/cpu/test_matmul_f16.cpp
#include <gtest/gtest.h>

#include "test_matmul.h"
#include "cpu/matmul_fp16.h"
#include "cpu/tensor.h"
#include "cpu/quantization_base.h"  // 提供 f16_to_f32 / f32_to_f16

namespace cpu_inference {

// 创建 FP16 Tensor 并用 FP32 数据初始化
Tensor make_fp16_tensor(int numa_id, const std::vector<int>& shape, const std::vector<float>& fp32_data) {
    size_t numel = 1;
    for (int s : shape) numel *= s;
    EXPECT_EQ(numel, fp32_data.size());

    std::vector<float16_t> fp16_data(numel);
    for (size_t i = 0; i < numel; ++i) {
        fp16_data[i] = fp32_to_fp16(fp32_data[i]);
    }

    Tensor tensor("test", numa_id, shape, "float16");
    tensor.write(numa_id, fp16_data.data(), fp16_data.size() * sizeof(float16_t));
    return tensor;
}

// 从 FP32 Tensor 提取数据
std::vector<float> extract_fp32_from_tensor(const Tensor& tensor, int numa_id) {
    size_t numel = tensor.compute_num_elements(tensor.get_shape());
    const float* ptr = static_cast<const float*>(tensor.at(numa_id));
    return std::vector<float>(ptr, ptr + numel);
}

class MatMulF16Test : public ::testing::Test {
protected:
    void SetUp() override {
        work.my_numa = 0;
        work.num_threads = 1;
        work.num_numas = 1;
        work.para = 1;
        work.global_tid = 0;
        work.tid = 0;
        work.threads_per_numa = 1;
        work.tid_in_numa = 0;
    }

    WorkDivider work;
};

TEST_F(MatMulF16Test, BasicCorrectness) {
    const int M = 64;
    const int K = 128;
    const int token_count = 2;
    const int numa_id = 0;

    // 构造确定性输入
    auto weight_fp32 = random_floats(M * K);
    auto input_fp32 = random_floats(token_count * K);

    // 创建张量
    auto weight = make_fp16_tensor(numa_id, {M, K}, weight_fp32);
    auto input  = make_fp16_tensor(numa_id, {token_count, K}, input_fp32);
    Tensor output("output", numa_id, {token_count, M}, "float32");

    // 执行 matmul
    MatMulBase* mm = new MatMulF16();
    mm->matmul(&work, output, weight, input);

    // 获取实际输出
    auto actual = extract_fp32_from_tensor(output, numa_id);

    // 计算参考输出
    auto expected = reference_matmul(weight_fp32, input_fp32, M, token_count, K);

    // 验证
    const float tolerance = 1e-6f;
    ASSERT_EQ(actual.size(), expected.size());
    for (size_t i = 0; i < actual.size(); ++i) {
        EXPECT_NEAR(expected[i], actual[i], std::max(0.01f, abs(expected[i] * tolerance)))
            << "Mismatch at index " << i;
    }
}

TEST_F(MatMulF16Test, HandlesNonAlignedK) {
    const int M = 33;
    const int K = 100;  // not multiple of 16 or 64
    const int token_count = 1;
    const int numa_id = 0;

    std::vector<float> weight_fp32(M * K, 1.0f);
    std::vector<float> input_fp32(token_count * K, 2.0f);

    auto weight = make_fp16_tensor(numa_id, {M, K}, weight_fp32);
    auto input  = make_fp16_tensor(numa_id, {token_count, K}, input_fp32);
    Tensor output("output", numa_id, {token_count, M}, "float32");

    MatMulF16 mm;
    mm.matmul(&work, output, weight, input);

    auto actual = extract_fp32_from_tensor(output, numa_id);
    std::vector<float> expected(token_count * M, K * 2.0f); // 1*2*K

    const float tolerance = 1e-3f;
    ASSERT_EQ(actual.size(), expected.size());
    for (size_t i = 0; i < actual.size(); ++i) {
        EXPECT_NEAR(expected[i], actual[i], expected[i] * tolerance);
    }
}

TEST_F(MatMulF16Test, SingleTokenSingleRow) {
    const int M = 1;
    const int K = 1;
    const int token_count = 1;
    const int numa_id = 0;

    std::vector<float> weight_fp32 = {3.0f};
    std::vector<float> input_fp32  = {2.0f};

    auto weight = make_fp16_tensor(numa_id, {M, K}, weight_fp32);
    auto input  = make_fp16_tensor(numa_id, {token_count, K}, input_fp32);
    Tensor output("output", numa_id, {token_count, M}, "float32");

    MatMulF16 mm;
    mm.matmul(&work, output, weight, input);

    auto actual = extract_fp32_from_tensor(output, numa_id);
    EXPECT_NEAR(6.0f, actual[0], 1e-6f);
}

} // namespace cpu_inference