#include "catlass/gemm/kernel/basic_matmul.hpp"
#include "catlass/gemm/device/device_gemm.hpp"
#include "catlass/status.hpp"
#include "golden.hpp" // 用於 CPU 結果校驗

using namespace Catlass;

// ==========================================
// 1. 問題建模：定義切分策略 (Tiling)
// ==========================================
// 針對 L1 Cache 大小進行優化: [M=128, N=256, K=256]
using L1TileShape = GemmShape<128, 256, 256>;
// 針對 Cube Unit 寄存器大小進行優化: [M=128, N=256, K=64]
using L0TileShape = GemmShape<128, 256, 64>;

// ==========================================
// 2. 架構適配：選擇調度策略 (Dispatch)
// ==========================================
using ArchTag = Arch::AtlasA2;
// 開啓 Atlas A2 專有的 Ping-Pong 流水線優化 (Double Buffer)
using DispatchPolicy = Gemm::MmadAtlasA2Pingpong<true>;

// ==========================================
// 3. 組件組裝：構建 Kernel (Composition)
// ==========================================
// 定義佈局與類型
using LayoutA = layout::RowMajor;
using LayoutB = layout::RowMajor;
using LayoutC = layout::RowMajor;

// 組裝核心計算塊：自動處理數據搬運與計算重疊
using BlockMmad = Gemm::Block::BlockMmad<
    DispatchPolicy,
    L1TileShape,
    L0TileShape,
    Gemm::GemmType<half, LayoutA>, // A 矩陣
    Gemm::GemmType<half, LayoutB>, // B 矩陣
    Gemm::GemmType<half, LayoutC>  // C 矩陣
>;

// 組裝完整 Kernel
using MatmulKernel = Gemm::Kernel::BasicMatmul<
    BlockMmad,
    void, // 無需 Epilogue
    Gemm::Block::GemmIdentityBlockSwizzle<3, 0> // 採用 Z 字形調度以優化 Cache 命中
>;

// Device Adapter: Host 端接口
using MatmulAdapter = Gemm::Device::DeviceGemm<MatmulKernel>;

// ==========================================
// 4. 運行時執行與驗證 (Execution & Verify)
// ==========================================
static void Run(const GemmOptions &options) {
    // 1. 初始化 ACL 資源
    aclrtStream stream{nullptr};
    ACL_CHECK(aclInit(nullptr));
    ACL_CHECK(aclrtSetDevice(options.deviceId));
    ACL_CHECK(aclrtCreateStream(&stream));

    // 2. 準備數據 (Data Preparation)
    // ==========================================
    size_t sizeA = options.problemShape.m() * options.problemShape.k() * sizeof(fp16_t);
    size_t sizeB = options.problemShape.k() * options.problemShape.n() * sizeof(fp16_t);
    size_t sizeC = options.problemShape.m() * options.problemShape.n() * sizeof(fp16_t);
    
    // 分配 Host 內存並初始化
    std::vector<fp16_t> hostA(options.problemShape.m() * options.problemShape.k());
    std::vector<fp16_t> hostB(options.problemShape.k() * options.problemShape.n());
    std::vector<fp16_t> hostC(options.problemShape.m() * options.problemShape.n(), 0);
    
    // 使用隨機數初始化 A、B 矩陣
    for (auto &v : hostA) v = static_cast<fp16_t>((rand() % 100) / 100.0f);
    for (auto &v : hostB) v = static_cast<fp16_t>((rand() % 100) / 100.0f);
    
    // 分配 Device 內存
    fp16_t *deviceA = nullptr;
    fp16_t *deviceB = nullptr;
    fp16_t *deviceC = nullptr;
    ACL_CHECK(aclrtMalloc((void**)&deviceA, sizeA, ACL_MEM_MALLOC_HUGE_FIRST));
    ACL_CHECK(aclrtMalloc((void**)&deviceB, sizeB, ACL_MEM_MALLOC_HUGE_FIRST));
    ACL_CHECK(aclrtMalloc((void**)&deviceC, sizeC, ACL_MEM_MALLOC_HUGE_FIRST));
    
    // 將數據從 Host 拷貝到 Device
    ACL_CHECK(aclrtMemcpy(deviceA, sizeA, hostA.data(), sizeA, ACL_MEMCPY_HOST_TO_DEVICE));
    ACL_CHECK(aclrtMemcpy(deviceB, sizeB, hostB.data(), sizeB, ACL_MEMCPY_HOST_TO_DEVICE));
    ACL_CHECK(aclrtMemset(deviceC, sizeC, 0, sizeC)); // 初始化 C 矩陣為 0

    // 3. 啓動 Kernel
    MatmulAdapter gemm;
    // 檢查 workspace 需求並分配
    MatmulKernel::Arguments args{options.problemShape, deviceA, deviceB, deviceC};
    size_t workspaceSize = gemm.GetWorkspaceSize(args);
    uint8_t* workspace = nullptr;
    if (workspaceSize > 0) aclrtMalloc((void**)&workspace, workspaceSize, ACL_MEM_MALLOC_HUGE_FIRST);

    // 執行計算
    gemm.Initialize(args, workspace);
    gemm(stream); // Launch!

    // 4. 同步與結果回傳
    ACL_CHECK(aclrtSynchronizeStream(stream));
    std::vector<fp16_t> hostC(lenC);
    ACL_CHECK(aclrtMemcpy(hostC.data(), sizeC, deviceC, sizeC, ACL_MEMCPY_DEVICE_TO_HOST));

    // 5. 結果校驗 (Golden Check)
    std::cout << "Verifying results against CPU reference..." << std::endl;
    auto errorIndices = golden::CompareData(hostC, hostGolden, options.problemShape.k());
    if (errorIndices.empty()) {
        std::cout << "Compare success. (Pass)" << std::endl;
    } else {
        std::cerr << "Compare failed. Errors: " << errorIndices.size() << std::endl;
    }

    // 資源釋放...
}

int main(int argc, char** argv) {
    GemmOptions options;
    if (options.Parse(argc, argv) != 0) return -1;
    Run(options);
    return 0;
}