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

#include "test_matmul.h"
#include "cpu/tensor.h"
#include "cpu/matmul_q8.h"
#include "cpu/quantization_q8_0.h"

namespace cpu_inference{

class MatMulQ8Test : 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;

    std::vector<float> run_matmul(
        int M, int K, int token_count,
        const std::vector<float>& weight_f32,
        const std::vector<float>& input_f32)
    {
        EXPECT_EQ(weight_f32.size(), M * K);
        EXPECT_EQ(input_f32.size(), token_count * K);
        EXPECT_EQ(K % QK8_0, 0);

        std::vector<block_q8_0> weight_q8(M * (K / QK8_0));
        std::vector<block_q8_0> input_q8(token_count * (K / QK8_0));

        QuantizationBase<block_q8_0>::quantize(
            weight_f32.data(), weight_q8.data(), M * K, "fp32");
        QuantizationBase<block_q8_0>::quantize(
            input_f32.data(), input_q8.data(), token_count * K, "fp32");
        const int node = 0;
        Tensor weight("weight", node, {M, K / QK8_0 * (int)sizeof(block_q8_0)}, "int8");
        weight.write(node, weight_q8.data(), weight_q8.size() * (int)sizeof(block_q8_0));

        Tensor input("input", node, {token_count, K / QK8_0 * (int)sizeof(block_q8_0)}, "int8");
        input.write(node, input_q8.data(), input_q8.size() * (int)sizeof(block_q8_0));

        Tensor output("output", node, {token_count, M}, "float32");

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

        // 提取结果
        const float* out_ptr = static_cast<const float*>(output.at(node));
        return std::vector<float>(out_ptr, out_ptr + token_count * M);
    }
};

// === 测试用例 ===

TEST_F(MatMulQ8Test, BasicShape) {
    const int M = 64, K = 128, token_count = 2;
    auto weight_f32 = random_floats(M * K);
    auto input_f32 = random_floats(token_count * K);

    auto result = run_matmul(M, K, token_count, weight_f32, input_f32);
    auto ref = reference_matmul(weight_f32, input_f32, M, token_count, K);

    const float tol = 0.1f;
    for (size_t i = 0; i < result.size(); ++i) {
        EXPECT_NEAR(result[i], ref[i], std::max(0.5f, abs(tol * ref[i])))
            << "Mismatch at index " << i
            << " (token=" << i / M << ", out_dim=" << i % M << ")";
    }
}

TEST_F(MatMulQ8Test, LargerMatrix) {
    const int M = 512, K = 256, token_count = 8;
    auto weight_f32 = random_floats(M * K);
    auto input_f32 = random_floats(token_count * K);

    auto result = run_matmul(M, K, token_count, weight_f32, input_f32);
    auto ref = reference_matmul(weight_f32, input_f32, M, token_count, K);

    const float tol = 0.1f;
    for (size_t i = 0; i < result.size(); ++i) {
        EXPECT_NEAR(result[i], ref[i], std::max(0.5f, abs(tol * ref[i])));
    }
}

TEST_F(MatMulQ8Test, SingleToken) {
    const int M = 128, K = 64, token_count = 1;
    auto weight_f32 = random_floats(M * K);
    auto input_f32 = random_floats(token_count * K);

    auto result = run_matmul(M, K, token_count, weight_f32, input_f32);
    auto ref = reference_matmul(weight_f32, input_f32, M, token_count, K);

    const float tol = 0.1f;
    for (size_t i = 0; i < result.size(); ++i) {
        EXPECT_NEAR(result[i], ref[i], std::max(0.5f, abs(tol * ref[i])));
    }
}

TEST_F(MatMulQ8Test, SmallK) {
    const int M = 32, K = QK8_0, token_count = 4; // K = 32
    auto weight_f32 = random_floats(M * K);
    auto input_f32 = random_floats(token_count * K);

    auto result = run_matmul(M, K, token_count, weight_f32, input_f32);
    auto ref = reference_matmul(weight_f32, input_f32, M, token_count, K);

    const float tol = 0.1f;
    for (size_t i = 0; i < result.size(); ++i) {
        EXPECT_NEAR(result[i], ref[i], std::max(0.5f, abs(tol * ref[i])));
    }
}

TEST_F(MatMulQ8Test, OddK) {
    const int M = 32, K = 3 * QK8_0, token_count = 4; // K = 96
    auto weight_f32 = random_floats(M * K);
    auto input_f32 = random_floats(token_count * K);

    auto result = run_matmul(M, K, token_count, weight_f32, input_f32);
    auto ref = reference_matmul(weight_f32, input_f32, M, token_count, K);

    const float tol = 0.1f;
    for (size_t i = 0; i < result.size(); ++i) {
        EXPECT_NEAR(result[i], ref[i], std::max(0.5f, abs(tol * ref[i])));
    }
}

TEST_F(MatMulQ8Test, OddM) {
    const int M = 35, K = 3 * QK8_0, token_count = 4; // K = 96
    auto weight_f32 = random_floats(M * K);
    auto input_f32 = random_floats(token_count * K);

    auto result = run_matmul(M, K, token_count, weight_f32, input_f32);
    auto ref = reference_matmul(weight_f32, input_f32, M, token_count, K);

    const float tol = 0.1f;
    for (size_t i = 0; i < result.size(); ++i) {
        EXPECT_NEAR(result[i], ref[i], std::max(0.5f, abs(tol * ref[i])));
    }
}

} // namespace cpu_inference