/**
 * This program is free software, you can redistribute it and/or modify.
 * Copyright (c) 2026 Huawei Technologies Co., Ltd.
 * This file is a part of the CANN Open Software.
 * Licensed under CANN Open Software License Agreement Version 2.0 (the "License").
 * Please refer to the License for details. You may not use this file except in compliance with the License.
 * 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 FITNESS FOR A PARTICULAR PURPOSE. See LICENSE in the root of
 * the software repository for the full text of the License.
 */

#ifndef K_MAX_SHAPE_DIM
#define K_MAX_SHAPE_DIM 0
#endif

#include "catlass/gemm/kernel/a8w4_mx_matmul.hpp"

#include "catlass/arch/arch.hpp"
#include "catlass/catlass.hpp"
#include "catlass/gemm/block/block_mmad.hpp"
#include "catlass/gemm/block/block_swizzle.hpp"
#include "catlass/gemm/device/device_gemm.hpp"
#include "catlass/gemm/dispatch_policy.hpp"
#include "catlass/gemm/gemm_type.hpp"
#include "catlass/layout/layout.hpp"
#include "catlass/status.hpp"
#include "tla/layout.hpp"

#include "golden.hpp"
#include "helper.hpp"

using namespace Catlass;
using namespace tla;

using Options = GemmOptions;

static void Run(const Options &options)
{
    aclrtStream stream{nullptr};

    ACL_CHECK(aclInit(nullptr));
    ACL_CHECK(aclrtSetDevice(options.deviceId));
    ACL_CHECK(aclrtCreateStream(&stream));

    // m、n、k
    uint32_t m = options.problemShape.m();
    uint32_t n = options.problemShape.n();
    uint32_t k = options.problemShape.k();
    uint32_t mxScaleK = CeilDiv<MX_SCALE_GROUP_NUM>(k);

    using ElementA = float8_e4m3_t;
    using ElementB = float8_e4m3_t;
    using ElementPrologueB = float4_e2m1x2_t;
    using ElementMxScale = float8_e8m0_t;
    using ElementC = float;
    using ElementBias = void;

    using ElementBiasType = std::conditional_t<std::is_void_v<ElementBias>, uint8_t, ElementBias>;

    // basic layout
    using LayoutA = layout::RowMajor;
    using LayoutB = layout::nZ;
    using LayoutPrologueB = layout::ColumnMajor;
    using LayoutC = layout::RowMajor;

    // makeLayout
    LayoutA tagA = LayoutA::MakeLayout<ElementA>(m, k);
    LayoutPrologueB tagPrologueB = LayoutPrologueB::MakeLayout<ElementPrologueB>(k, n);
    LayoutC tagC = LayoutC::MakeLayout<ElementC>(m, n);

    static constexpr uint32_t MX_k_ALIGN = 2;
    static constexpr uint32_t SIZE_MAGNIFICATION = 2;

    // data length
    size_t lenA = tagA.Capacity();
    size_t lenPrologueB = tagPrologueB.Capacity();
    uint32_t mxScaleAlignedK = RoundUp<MX_k_ALIGN>(mxScaleK);
    size_t lenMxScaleA = m * mxScaleAlignedK;
    size_t lenMxScaleB = mxScaleAlignedK * n;
    size_t lenC = tagC.Capacity();
    size_t lenBias = static_cast<size_t>(n);

    // data size(len * sizeof(element))
    size_t sizeA = lenA * sizeof(ElementA);
    size_t sizeB = lenPrologueB / SIZE_MAGNIFICATION;
    size_t sizeMxScaleA = lenMxScaleA * sizeof(ElementMxScale);
    size_t sizeMxScaleB = lenMxScaleB * sizeof(ElementMxScale);
    size_t sizeC = lenC * sizeof(ElementC);
    size_t sizeBias = lenBias * sizeof(ElementBiasType);
    size_t sizeWorkspace;

    // host
    std::vector<int8_t> hostA(sizeA);
    std::vector<int8_t> hostB(sizeB);
    std::vector<int8_t> hostMxScaleA(lenMxScaleA);
    std::vector<int8_t> hostMxScaleB(lenMxScaleB);
    std::vector<ElementBiasType> hostBias(lenBias);

    const auto releaseAclEarly = [&]() {
        ACL_CHECK(aclrtDestroyStream(stream));
        ACL_CHECK(aclrtResetDevice(options.deviceId));
        ACL_CHECK(aclFinalize());
    };
    //file read
    std::string inFileAName = "../../input/a_8.bin";
    std::ifstream inFileA(inFileAName, std::ios::binary);
    if (!inFileA.is_open()) {
        std::cerr << "Failed to open inFileA: " << inFileAName << std::endl;
        releaseAclEarly();
        return;
    } else {
        inFileA.read(reinterpret_cast<char *>(hostA.data()), sizeA);
        inFileA.close();
    }
    std::string inFileBName = "../../input/b_4.bin";
    std::ifstream inFileB(inFileBName, std::ios::binary);
    if (!inFileB.is_open()) {
        std::cerr << "Failed to open inFileB: " << inFileBName << std::endl;
        releaseAclEarly();
        return;
    } else {
        inFileB.read(reinterpret_cast<char *>(hostB.data()), sizeB);
        inFileB.close();
    }
    std::string inFileMxScaleAName = "../../input/a_scale.bin";
    std::ifstream inFileMxScaleA(inFileMxScaleAName, std::ios::binary);
    if (!inFileMxScaleA.is_open()) {
        std::cerr << "Failed to open inFileMxScaleA: " << inFileMxScaleAName << std::endl;
        releaseAclEarly();
        return;
    } else {
        inFileMxScaleA.read(reinterpret_cast<char *>(hostMxScaleA.data()), sizeMxScaleA);
        inFileMxScaleA.close();
    }
    std::string inFileMxScaleBName = "../../input/b_scale.bin";
    std::ifstream inFileMxScaleB(inFileMxScaleBName, std::ios::binary);
    if (!inFileMxScaleB.is_open()) {
        std::cerr << "Failed to open inFileMxScaleB: " << inFileMxScaleBName << std::endl;
        releaseAclEarly();
        return;
    } else {
        inFileMxScaleB.read(reinterpret_cast<char *>(hostMxScaleB.data()), sizeMxScaleB);
        inFileMxScaleB.close();
    }
    if constexpr (!std::is_void_v<ElementBias>) {
        std::string inFileBiasName = "../../input/bias.bin";
        std::ifstream inFileBias(inFileBiasName, std::ios::binary);
        if (!inFileBias.is_open()) {
            std::cerr << "Failed to open inFileBias: " << inFileBiasName << std::endl;
            releaseAclEarly();
            return;
        } else {
            inFileBias.read(reinterpret_cast<char *>(hostBias.data()), sizeBias);
            inFileBias.close();
        }
    }

    // device
    uint8_t *deviceA{nullptr};
    ACL_CHECK(aclrtMalloc(reinterpret_cast<void **>(&deviceA), sizeA, ACL_MEM_MALLOC_HUGE_FIRST));
    ACL_CHECK(aclrtMemcpy(deviceA, sizeA, hostA.data(), sizeA, ACL_MEMCPY_HOST_TO_DEVICE));

    uint8_t *deviceB{nullptr};
    ACL_CHECK(aclrtMalloc(reinterpret_cast<void **>(&deviceB), sizeB, ACL_MEM_MALLOC_HUGE_FIRST));
    ACL_CHECK(aclrtMemcpy(deviceB, sizeB, hostB.data(), sizeB, ACL_MEMCPY_HOST_TO_DEVICE));

    uint8_t *deviceMxScaleA{nullptr};
    ACL_CHECK(aclrtMalloc(reinterpret_cast<void **>(&deviceMxScaleA), sizeMxScaleA, ACL_MEM_MALLOC_HUGE_FIRST));
    ACL_CHECK(aclrtMemcpy(deviceMxScaleA, sizeMxScaleA, hostMxScaleA.data(), sizeMxScaleA, ACL_MEMCPY_HOST_TO_DEVICE));

    uint8_t *deviceMxScaleB{nullptr};
    ACL_CHECK(aclrtMalloc(reinterpret_cast<void **>(&deviceMxScaleB), sizeMxScaleB, ACL_MEM_MALLOC_HUGE_FIRST));
    ACL_CHECK(aclrtMemcpy(deviceMxScaleB, sizeMxScaleB, hostMxScaleB.data(), sizeMxScaleB, ACL_MEMCPY_HOST_TO_DEVICE));

    uint8_t *deviceC{nullptr};
    ACL_CHECK(aclrtMalloc(reinterpret_cast<void **>(&deviceC), sizeC, ACL_MEM_MALLOC_HUGE_FIRST));

    uint8_t *deviceBias{nullptr};
    if constexpr (!std::is_void_v<ElementBias>) {
        ACL_CHECK(aclrtMalloc(reinterpret_cast<void **>(&deviceBias), sizeBias, ACL_MEM_MALLOC_HUGE_FIRST));
        ACL_CHECK(aclrtMemcpy(deviceBias, sizeBias, hostBias.data(), sizeBias, ACL_MEMCPY_HOST_TO_DEVICE));
    }

    uint8_t *deviceWorkspace{nullptr};

    // Get the number of cube cores of the current hardware
    auto aicCoreNum = platform_ascendc::PlatformAscendCManager::GetInstance()->GetCoreNumAic();

    // archtag uniflag
    using ArchTag = Arch::Ascend950;
    constexpr bool enableUnitFlag = false;

    // shape & type
    using L1TileShape = Shape<Int<128>, Int<128>, Int<128>>;
    using L0TileShape = Shape<Int<128>, Int<128>, Int<128>>;
    using PrologueSrcType = Gemm::GemmType<ElementPrologueB, LayoutPrologueB>;
    using PrologueDstType = Gemm::GemmType<ElementB, LayoutB>;

    // DispatchPolicy
    using DispatchPolicyMmad = Gemm::MmadA8W4Mx<ArchTag, enableUnitFlag>;
    using DispatchPolicyPrologue = Gemm::MxA8W4Prologue<ArchTag>;

    // layout (tla)
    auto layoutA = tla::MakeLayout<ElementA, LayoutA>(m, k);
    auto layoutprologueB = tla::MakeLayout<ElementPrologueB, LayoutPrologueB>(k, n);
    auto layoutMxScaleA = tla::MakeMxScaleLayout<ElementMxScale, LayoutA, false>(m, mxScaleK);
    auto layoutMxScaleB = tla::MakeMxScaleLayout<ElementMxScale, LayoutPrologueB, true>(mxScaleK, n);
    auto layoutC = tla::MakeLayout<ElementC, LayoutC>(m, n);

    // tile
    using TileCopy = Gemm::Tile::PackedMxA8W4TileCopyTla<
        ArchTag, ElementA, LayoutA, ElementPrologueB, LayoutPrologueB, ElementB, LayoutB, ElementMxScale, decltype(layoutMxScaleA), ElementMxScale,
        decltype(layoutMxScaleB), ElementC, LayoutC, ElementBias, false, Gemm::Tile::ScaleGranularity::PER_TENSOR>;

    // BlockMmad
    using BlockMmad = Gemm::Block::BlockMmadA8W4Mx<
        DispatchPolicyMmad, L1TileShape, L0TileShape, ElementA, ElementB, ElementC, ElementPrologueB, ElementBias, TileCopy>;

    using BlockPrologue = Gemm::Block::BlockPrologue<
        DispatchPolicyPrologue, PrologueSrcType, PrologueDstType, L1TileShape, TileCopy>;

    // Epilogue
    using BlockEpilogue = void;

    //Swizzle offset is 3 and direction is 0.
    using BlockScheduler=  typename Gemm::Block::GemmIdentityBlockSwizzle<3, 0>;

    // kernel level
    using MatmulKernel = Gemm::Kernel::A8W4MxMatmul<BlockMmad, BlockPrologue, BlockEpilogue, BlockScheduler>;

    using MatmulAdapter = Gemm::Device::DeviceGemm<MatmulKernel>;

    MatmulKernel::Arguments arguments{
        options.problemShape, deviceA,        layoutA,        deviceB, layoutprologueB, deviceMxScaleA,
        layoutMxScaleA,       deviceMxScaleB, layoutMxScaleB, deviceC, layoutC, deviceBias
    };

    uint32_t taskNum = CeilDiv(options.problemShape.m(), tla::get<0>(L1TileShape{})) *
                       CeilDiv(options.problemShape.n(), tla::get<1>(L1TileShape{}));
    uint32_t aicCoreUsed = min(aicCoreNum, taskNum);

    MatmulAdapter matmulOp;
    matmulOp.CanImplement(arguments);
    sizeWorkspace = matmulOp.GetWorkspaceSize(arguments);
    if (sizeWorkspace > 0) {
        ACL_CHECK(aclrtMalloc(reinterpret_cast<void **>(&deviceWorkspace), sizeWorkspace, ACL_MEM_MALLOC_HUGE_FIRST));
    }
    matmulOp.Initialize(arguments, deviceWorkspace);
    matmulOp(stream, aicCoreUsed);

    ACL_CHECK(aclrtSynchronizeStream(stream));

    std::vector<ElementC> hostC(lenC);
    ACL_CHECK(aclrtMemcpy(hostC.data(), sizeC, deviceC, sizeC, ACL_MEMCPY_DEVICE_TO_HOST));

    std::vector<float> hostGolden(lenC);
    std::string outputFileName = "../../golden/expected_data.bin";
    if (!ReadFile(outputFileName, hostGolden.data(), sizeof(float) * hostGolden.size())) {
        ACL_CHECK(aclrtFree(deviceA));
        ACL_CHECK(aclrtFree(deviceB));
        ACL_CHECK(aclrtFree(deviceMxScaleA));
        ACL_CHECK(aclrtFree(deviceMxScaleB));
        ACL_CHECK(aclrtFree(deviceC));
        if constexpr (!std::is_void_v<ElementBias>) {
            ACL_CHECK(aclrtFree(deviceBias));
        }
        if (sizeWorkspace > 0) {
            ACL_CHECK(aclrtFree(deviceWorkspace));
        }
        ACL_CHECK(aclrtDestroyStream(stream));
        ACL_CHECK(aclrtResetDevice(options.deviceId));
        ACL_CHECK(aclFinalize());
        return;
    }

    std::vector<uint64_t> errorIndices = golden::CompareData(hostC, hostGolden, k);
    if (errorIndices.empty()) {
        std::cout << "Compare success." << std::endl;
    } else {
        std::cerr << "Compare failed. Error count: " << errorIndices.size() << std::endl;
    }

    ACL_CHECK(aclrtFree(deviceA));
    ACL_CHECK(aclrtFree(deviceB));
    ACL_CHECK(aclrtFree(deviceMxScaleA));
    ACL_CHECK(aclrtFree(deviceMxScaleB));
    ACL_CHECK(aclrtFree(deviceC));
    if constexpr (!std::is_void_v<ElementBias>) {
        ACL_CHECK(aclrtFree(deviceBias));
    }
    if (sizeWorkspace > 0) {
        ACL_CHECK(aclrtFree(deviceWorkspace));
    }

    ACL_CHECK(aclrtDestroyStream(stream));
    ACL_CHECK(aclrtResetDevice(options.deviceId));
    ACL_CHECK(aclFinalize());
}

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