// SPDX-License-Identifier: MIT
// Copyright (c) 2025, Advanced Micro Devices, Inc. All rights reserved.

// Unit tests for tile_engine generated GEMM kernels
// Tests kernel correctness using tile_engine's verification methodology

#include <gtest/gtest.h>
#include <iostream>

#include "ck_tile/core.hpp"
#include "ck_tile/host.hpp"
#include "tile_engine/ops/gemm/gemm_common.hpp"

// The kernel header is included via compile command line with -include flag
// It defines SelectedKernel struct, KERNEL_NAME, and tensor data types

// Adaptive error threshold calculation matching tile_engine's implementation
template <typename ADataType, typename BDataType, typename AccDataType, typename CDataType>
auto calculate_rtol_atol(const ck_tile::index_t K,
                         const ck_tile::index_t kbatch,
                         const float max_accumulated_value)
{
    using ComputeType =
        std::conditional_t<sizeof(ADataType) < sizeof(BDataType), ADataType, BDataType>;
    // Calculate thresholds
    const auto rtol = ck_tile::get_relative_threshold<ComputeType, CDataType, AccDataType>(
        ck_tile::integer_divide_ceil(K, kbatch));
    const auto atol = ck_tile::get_absolute_threshold<ComputeType, CDataType, AccDataType>(
        max_accumulated_value / kbatch, ck_tile::integer_divide_ceil(K, kbatch));
    // Calculate error due to split_k accumulation
    const auto rtol_split_k =
        ck_tile::get_relative_threshold<CDataType, CDataType, CDataType>(kbatch);
    const auto atol_split_k = ck_tile::get_absolute_threshold<CDataType, CDataType, CDataType>(
        max_accumulated_value, kbatch);
    // Use higher threshold
    return ck_tile::make_tuple(std::max(rtol, rtol_split_k), std::max(atol, atol_split_k));
}

/// @brief Function to compare the results of the device and host computations (from tile_engine)
template <typename ADataType, typename BDataType, typename AccDataType, typename CDataType>
bool compare_results(std::string instanceName,
                     ck_tile::index_t K,
                     ck_tile::index_t kbatch,
                     ck_tile::HostTensor<CDataType>& c_m_n_dev_result,
                     ck_tile::HostTensor<CDataType>& c_m_n_host_result)
{
    const float max_accumulated_value =
        *std::max_element(c_m_n_host_result.mData.begin(), c_m_n_host_result.mData.end());
    const auto rtol_atol = calculate_rtol_atol<ADataType, BDataType, AccDataType, CDataType>(
        K, kbatch, max_accumulated_value);
    bool pass = ck_tile::check_err(c_m_n_dev_result,
                                   c_m_n_host_result,
                                   "Error: Incorrect results!",
                                   rtol_atol.at(ck_tile::number<0>{}),
                                   rtol_atol.at(ck_tile::number<1>{}));

    std::cout << "For " << instanceName << " Relative error threshold is "
              << rtol_atol.at(ck_tile::number<0>{}) << " Absolute error threshold is "
              << rtol_atol.at(ck_tile::number<1>{}) << std::endl;
    std::cout << "The verification result is:" << (pass ? "correct" : "fail") << std::endl;

    return pass;
}

// Test parameter structure for matrix dimensions and split_k values
struct GemmTestParams
{
    int m, n, k, split_k;
};

class GemmTileEngineTest : public ::testing::TestWithParam<GemmTestParams>
{
    protected:
    void SetUp() override
    {
        auto params = GetParam();
        m_          = params.m;
        n_          = params.n;
        k_          = params.k;
        split_k_    = params.split_k;

        // Calculate strides (following tile_engine pattern)
        if constexpr(std::is_same_v<ALayout, ck_tile::tensor_layout::gemm::RowMajor>)
        {
            stride_a_ = k_;
        }
        else
        {
            stride_a_ = m_;
        }

        if constexpr(std::is_same_v<BLayout, ck_tile::tensor_layout::gemm::RowMajor>)
        {
            stride_b_ = n_;
        }
        else
        {
            stride_b_ = k_;
        }

        if constexpr(std::is_same_v<CLayout, ck_tile::tensor_layout::gemm::RowMajor>)
        {
            stride_c_ = n_;
        }
        else
        {
            stride_c_ = m_;
        }
    }

    // Test dimensions
    int m_, n_, k_, split_k_;
    int stride_a_, stride_b_, stride_c_;
};

TEST_P(GemmTileEngineTest, BasicFunctionality)
{
    // Get tensor layouts from generated kernel
    const ALayout layout_a = ALayout{};
    const BLayout layout_b = BLayout{};
    const CLayout layout_c = CLayout{};

    // Use split_k from test parameters
    int split_k       = split_k_;
    int stride_a_calc = ck_tile::get_default_stride(m_, k_, 0, is_row_major(layout_a));
    int stride_b_calc = ck_tile::get_default_stride(k_, n_, 0, is_row_major(layout_b));
    int stride_c_calc = ck_tile::get_default_stride(m_, n_, 0, is_row_major(layout_c));

    // Create host tensors with proper descriptors
    ck_tile::HostTensor<ADataType> a_m_k(
        ck_tile::host_tensor_descriptor(m_, k_, stride_a_calc, is_row_major(layout_a)));
    ck_tile::HostTensor<BDataType> b_k_n(
        ck_tile::host_tensor_descriptor(k_, n_, stride_b_calc, is_row_major(layout_b)));
    ck_tile::HostTensor<CDataType> c_m_n_dev_result(
        ck_tile::host_tensor_descriptor(m_, n_, stride_c_calc, is_row_major(layout_c)));
    ck_tile::HostTensor<CDataType> c_m_n_host_result(
        ck_tile::host_tensor_descriptor(m_, n_, stride_c_calc, is_row_major(layout_c)));

    // Initialize input tensors with uniform random distribution [-1.0, 1.0] (matches tile_engine)
    ck_tile::FillUniformDistribution<ADataType>{-1.f, 1.f}(a_m_k);
    ck_tile::FillUniformDistribution<BDataType>{-1.f, 1.f}(b_k_n);

    // Allocate GPU device memory
    ck_tile::DeviceMem a_m_k_dev_buf(a_m_k.get_element_space_size_in_bytes());
    ck_tile::DeviceMem b_k_n_dev_buf(b_k_n.get_element_space_size_in_bytes());
    ck_tile::DeviceMem c_m_n_dev_buf(c_m_n_dev_result.get_element_space_size_in_bytes());

    // Copy data to device and zero output buffer
    a_m_k_dev_buf.ToDevice(a_m_k.data());
    b_k_n_dev_buf.ToDevice(b_k_n.data());
    c_m_n_dev_buf.SetZero();
    c_m_n_dev_result.SetZero();

    // Calculate reference result on host for verification
    ck_tile::reference_gemm<ADataType, BDataType, AccDataType, CDataType>(
        a_m_k, b_k_n, c_m_n_host_result);

    // Create GEMM kernel arguments
    ck_tile::GemmHostArgs gemm_args(a_m_k_dev_buf.GetDeviceBuffer(),
                                    b_k_n_dev_buf.GetDeviceBuffer(),
                                    c_m_n_dev_buf.GetDeviceBuffer(),
                                    split_k,
                                    m_,
                                    n_,
                                    k_,
                                    stride_a_calc,
                                    stride_b_calc,
                                    stride_c_calc);

    // Configure kernel execution for maximum speed (no timing, no debug output)
    ck_tile::stream_config stream_config{nullptr, // stream
                                         false,   // time_kernel (disable timing for speed)
                                         0,       // log_level (disable debug output)
                                         0,       // n_warmup
                                         1,       // n_repeat
                                         false,   // is_gpu_timer (unused when time_kernel=false)
                                         false,   // flush_cache
                                         1};      // rotating_count

    // Launch the generated kernel (no timing overhead for fastest execution)
    try
    {
        SelectedKernel::launch(gemm_args, stream_config);
        // Kernel launched successfully if no exception thrown
    }
    catch(const std::exception& e)
    {
        FAIL() << "Kernel launch failed: " << e.what();
    }

    // Copy result back from device
    c_m_n_dev_buf.FromDevice(c_m_n_dev_result.data());

    // Verify results using tile_engine's adaptive error thresholds
    bool verification_passed = compare_results<ADataType, BDataType, AccDataType, CDataType>(
        KERNEL_NAME, k_, split_k, c_m_n_dev_result, c_m_n_host_result);

    EXPECT_TRUE(verification_passed) << "GEMM result verification failed";
}

TEST_P(GemmTileEngineTest, KernelInfo)
{
    // Simple test to verify kernel information is available
    EXPECT_TRUE(strlen(KERNEL_NAME) > 0) << "Kernel name should not be empty";

    std::cout << "Testing kernel: " << KERNEL_NAME << std::endl;
    std::cout << "Problem size: " << m_ << "x" << n_ << "x" << k_ << " with split_k=" << split_k_
              << std::endl;
}

// Define test parameters for GEMM verification
INSTANTIATE_TEST_SUITE_P(GemmVerification,
                         GemmTileEngineTest,
                         ::testing::Values(GemmTestParams{256, 256, 128, 1},
                                           GemmTestParams{256, 256, 1024, 1},
                                           GemmTestParams{256, 512, 512, 1},
                                           GemmTestParams{512, 256, 512, 1}),
                         [](const ::testing::TestParamInfo<GemmTestParams>& param_info) {
                             return std::to_string(param_info.param.m) + "x" +
                                    std::to_string(param_info.param.n) + "x" +
                                    std::to_string(param_info.param.k) + "_splitk" +
                                    std::to_string(param_info.param.split_k);
                         });
