/**
 * Copyright (c) 2025 Huawei Technologies Co., Ltd.
 * This program is free software, you can redistribute it and/or modify it under the terms and conditions of
 * 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.
 */

/*!
 * \file moe_custom_row_idx_gather_droppad.h
 * \brief
 */
#ifndef MOE_CUSTOM_ROW_IDX_GATHER_DROPPAD_H
#define MOE_CUSTOM_ROW_IDX_GATHER_DROPPAD_H

#include "moe_custom_common.h"

namespace MoeInitRoutingCustom {
using namespace AscendC;

template <typename T, typename TilingData>
class MoeCustomSrcToDstWithCapacity {
public:
    __aicore__ inline MoeCustomSrcToDstWithCapacity(){};
    __aicore__ inline void Init(GM_ADDR expandedRowIdx, GM_ADDR expandedX, GM_ADDR expandedScale, GM_ADDR workspace,
                                const TilingData *tilingData, TPipe *tPipe);
    __aicore__ inline void Process();

private:
    __aicore__ inline void CopyIn(int64_t progress);
    __aicore__ inline void CopyOut(int64_t progress);
    __aicore__ inline void CopyOutRemain();
    __aicore__ inline void SyncAll();
    __aicore__ inline void AssistInit();

private:
    TPipe *pipe;
    TQue<QuePosition::VECIN, 1> copyInQueue;
    TQue<QuePosition::VECOUT, 1> copyOutQueue;
    TQue<QuePosition::VECOUT, 1> copyOutZeroQueue;
    TQue<QuePosition::VECOUT, 1> scaleOutZeroQueue;

    GlobalTensor<int32_t> expandDstToSrcRowGm;
    GlobalTensor<int32_t> expandedRowIdxGm;
    GlobalTensor<int32_t> expertIdxValueGm;
    GlobalTensor<int32_t> expandedExpertIdxGm;
    GlobalTensor<T> expandedXGm;
    GlobalTensor<float> expandedScaleGm;

    LocalTensor<T> outTmpLocal;
    LocalTensor<float> scaleLocal;

    const MoeCustomSrcToDstCapacityComputeTilingData *srcToDstTilingData;
    int64_t coreNum;
    int64_t blockIdx;
    int64_t totalLength;
    int64_t currentLoopRows;
    int64_t coreRows;
    int64_t perLoopRows;
    int64_t lastLoopRows;
    int64_t rowLoops;
    int64_t expertCapacity;
    int64_t expertNum;
    int64_t cols;
    int64_t perLoopCols;
    int64_t lastLoopCols;
    int64_t colLoops;
    int64_t isInputScale_;
    int64_t quantMode_;

    int64_t tokenCount = 0;
    int32_t lastExpertId = -1;
    int32_t lastCoreExpertId = 0;
    int32_t lastCoreExpertIdNum = 0;
    bool needScaleCopy = false;
};

template <typename T, typename TilingData>
__aicore__ inline void MoeCustomSrcToDstWithCapacity<T, TilingData>::AssistInit()
{
    if constexpr (IsSameType<T, int8_t>::value) {
        LocalTensor<int16_t> outLocal = copyOutZeroQueue.AllocTensor<int16_t>();
        Duplicate<int16_t>(outLocal, static_cast<int16_t>(0), this->perLoopCols);
        copyOutZeroQueue.EnQue<int16_t>(outLocal);
    } else {
        LocalTensor<T> outLocal = copyOutZeroQueue.AllocTensor<T>();
        Duplicate<T>(outLocal, static_cast<T>(0), this->perLoopCols);
        copyOutZeroQueue.EnQue<T>(outLocal);
    }
    if (this->needScaleCopy) {
        LocalTensor<float> scaleOutLocal = scaleOutZeroQueue.AllocTensor<float>();
        Duplicate<float>(scaleOutLocal, 0.0f, FP32_ONE_BLOCK_NUM);
        scaleOutZeroQueue.EnQue<float>(scaleOutLocal);
    }

    if (this->blockIdx != 0) {
        this->lastCoreExpertId = expertIdxValueGm.GetValue((this->blockIdx - 1) * 2);
        this->lastCoreExpertIdNum = expertIdxValueGm.GetValue((this->blockIdx - 1) * 2 + 1);
        for (int64_t i = this->blockIdx - 2; i >= 0; i--) {
            int32_t lastExpertIdx = expertIdxValueGm.GetValue(i * 2);
            if (lastExpertIdx < this->lastCoreExpertId) {
                break;
            }
            int32_t lastExpertNum = expertIdxValueGm.GetValue(i * 2 + 1);
            this->lastCoreExpertIdNum += lastExpertNum;
        }
    }
}

template <typename T, typename TilingData>
__aicore__ inline void MoeCustomSrcToDstWithCapacity<T, TilingData>::CopyIn(int64_t progress)
{
    LocalTensor<int32_t> inLocal = copyInQueue.AllocTensor<int32_t>();
    int64_t length = Align(currentLoopRows, sizeof(int32_t));
    DataCopy(inLocal, expandDstToSrcRowGm[progress * perLoopRows], length);
    DataCopy(inLocal[length], expandedExpertIdxGm[progress * perLoopRows], length);
    copyInQueue.EnQue<int32_t>(inLocal);
}

template <typename T, typename TilingData>
__aicore__ inline void MoeCustomSrcToDstWithCapacity<T, TilingData>::CopyOut(int64_t progress)
{
    LocalTensor<int32_t> inLocal = copyInQueue.DeQue<int32_t>();
    LocalTensor<int32_t> outLocal = copyOutQueue.AllocTensor<int32_t>();
    int64_t length = Align(currentLoopRows, sizeof(int32_t));
    DataCopyExtParams copyParams{static_cast<uint16_t>(1), static_cast<uint32_t>(sizeof(int32_t)), 0, 0, 0};
    DataCopyExtParams ScaleParams{1, static_cast<uint32_t>(sizeof(float)), 0, 0, 0};

    SetWaitFlag<HardEvent::MTE2_S>(HardEvent::MTE2_S);
    if (this->lastExpertId == -1) {
        this->lastExpertId = this->lastCoreExpertId;
        this->tokenCount = this->lastCoreExpertIdNum;
    }
    for (int64_t idx = 0; idx < currentLoopRows; idx++) {
        int32_t expertIdx = inLocal[length].GetValue(idx);
        int32_t index = 0;
        while (this->lastExpertId < expertIdx) {
            while (this->tokenCount < this->expertCapacity) {
                index = this->lastExpertId * this->expertCapacity + this->tokenCount;
                if (this->needScaleCopy) {
                    DataCopyPad(expandedScaleGm[index], this->scaleLocal, ScaleParams);
                }
                int64_t col = this->perLoopCols;
                for (int64_t i = 0; i < this->colLoops; i++) {
                    if (i == this->colLoops - 1) {
                        col = this->lastLoopCols;
                    }
                    DataCopyExtParams copyParams1{static_cast<uint16_t>(1), static_cast<uint32_t>(col * sizeof(T)), 0,
                                                  0, 0};
                    DataCopyPad(expandedXGm[index * this->cols + i * this->perLoopCols], this->outTmpLocal,
                                copyParams1);
                }
                this->tokenCount++;
            }
            this->tokenCount = 0;
            this->lastExpertId++;
        }

        if (this->tokenCount < this->expertCapacity) {
            int32_t outOffset = inLocal.GetValue(idx);
            index = expertIdx * this->expertCapacity + this->tokenCount;
            outLocal.SetValue(0, index);
            SetWaitFlag<HardEvent::S_MTE3>(HardEvent::S_MTE3);
            DataCopyPad(expandedRowIdxGm[outOffset], outLocal, copyParams);
            this->tokenCount++;
        }
    }
    copyInQueue.FreeTensor(inLocal);
    copyOutQueue.FreeTensor(outLocal);
}

template <typename T, typename TilingData>
__aicore__ inline void MoeCustomSrcToDstWithCapacity<T, TilingData>::CopyOutRemain()
{
    if (this->blockIdx != this->srcToDstTilingData->needCoreNum - 1) {
        copyOutZeroQueue.FreeTensor(this->outTmpLocal);
        if (this->needScaleCopy) {
            scaleOutZeroQueue.FreeTensor(this->scaleLocal);
        }
        return;
    }
    DataCopyExtParams ScaleParams{1, static_cast<uint32_t>(sizeof(float)), 0, 0, 0};
    while (this->lastExpertId < this->expertNum) {
        while (this->tokenCount < this->expertCapacity) {
            int32_t index = this->lastExpertId * this->expertCapacity + this->tokenCount;
            if (this->needScaleCopy) {
                DataCopyPad(expandedScaleGm[index], this->scaleLocal, ScaleParams);
            }
            int64_t col = this->perLoopCols;
            for (int64_t i = 0; i < this->colLoops; i++) {
                if (i == this->colLoops - 1) {
                    col = this->lastLoopCols;
                }
                DataCopyExtParams copyParams{static_cast<uint16_t>(1), static_cast<uint32_t>(col * sizeof(T)), 0, 0, 0};
                DataCopyPad(expandedXGm[index * this->cols + i * this->perLoopCols], this->outTmpLocal, copyParams);
                SetWaitFlag<HardEvent::MTE3_S>(HardEvent::MTE3_S);
            }
            this->tokenCount++;
        }
        this->tokenCount = 0;
        this->lastExpertId++;
    }
    copyOutZeroQueue.FreeTensor(this->outTmpLocal);
    if (this->needScaleCopy) {
        scaleOutZeroQueue.FreeTensor(this->scaleLocal);
    }
}

template <typename T, typename TilingData>
__aicore__ inline void MoeCustomSrcToDstWithCapacity<T, TilingData>::SyncAll()
{
    if (coreNum == 1) {
        return;
    }
#ifndef __CCE_KT_TEST__
    AscendC::SyncAll();
#endif
}

template <typename T, typename TilingData>
__aicore__ inline void MoeCustomSrcToDstWithCapacity<T, TilingData>::Init(GM_ADDR expandedRowIdx, GM_ADDR expandedX,
                                                                      GM_ADDR expandedScale, GM_ADDR workspace,
                                                                      const TilingData *tilingData,
                                                                      TPipe *tPipe)
{
    int64_t blockNum = GetBlockNum();
    this->pipe = tPipe;
    this->blockIdx = GetBlockIdx();

    this->coreNum = tilingData->coreNum;
    this->totalLength = tilingData->n * tilingData->k;
    this->srcToDstTilingData = &(tilingData->srcToDstDropPadParamsOp);
    this->expertNum = tilingData->expertNum;
    this->expertCapacity = tilingData->expertCapacity;
    this->cols = tilingData->cols;
    this->isInputScale_ = tilingData->isInputScale;
    this->quantMode_ = tilingData->quantMode;

    if (this->blockIdx == this->srcToDstTilingData->needCoreNum - 1) {
        this->coreRows = this->srcToDstTilingData->lastCoreRows;
        this->perLoopRows = this->srcToDstTilingData->lastCorePerLoopRows;
        this->lastLoopRows = this->srcToDstTilingData->lastCoreLastLoopRows;
        this->rowLoops = this->srcToDstTilingData->lastCoreLoops;
    } else {
        this->coreRows = this->srcToDstTilingData->perCoreRows;
        this->perLoopRows = this->srcToDstTilingData->perCorePerLoopRows;
        this->lastLoopRows = this->srcToDstTilingData->perCoreLastLoopRows;
        this->rowLoops = this->srcToDstTilingData->perCoreLoops;
    }
    this->perLoopCols = this->srcToDstTilingData->perLoopCols;
    this->lastLoopCols = this->srcToDstTilingData->lastLoopCols;
    this->colLoops = this->srcToDstTilingData->colLoops;
    this->needScaleCopy = (this->isInputScale_ != 0 && this->quantMode_ == -1);

    expandedScaleGm.SetGlobalBuffer((__gm__ float *)expandedScale);

    int64_t length = Align(this->totalLength, sizeof(int32_t));
    expandedRowIdxGm.SetGlobalBuffer((__gm__ int32_t *)expandedRowIdx, length);
    expandedXGm.SetGlobalBuffer((__gm__ T *)expandedX, this->expertNum * this->expertCapacity * this->cols);

    expandedExpertIdxGm.SetGlobalBuffer((__gm__ int32_t *)workspace +
                                            this->blockIdx * this->srcToDstTilingData->perCoreRows,
                                        Align(this->coreRows, sizeof(int32_t)));
    expandDstToSrcRowGm.SetGlobalBuffer((__gm__ int32_t *)workspace + length +
                                            this->blockIdx * this->srcToDstTilingData->perCoreRows,
                                        Align(this->coreRows, sizeof(int32_t)));
    expertIdxValueGm.SetGlobalBuffer(
        (__gm__ int32_t *)workspace + length * 2 + Align(this->expertNum, sizeof(int32_t)) * 2, this->coreNum * 2);

    pipe->InitBuffer(copyInQueue, 1, AlignBytes(this->perLoopRows, sizeof(int32_t)) * 2);
    pipe->InitBuffer(copyOutQueue, 1, AlignBytes(INT32_ONE_BLOCK_NUM, sizeof(int32_t)));
    if constexpr (IsSameType<T, int8_t>::value) {
        pipe->InitBuffer(copyOutZeroQueue, 1, AlignBytes(this->perLoopCols, sizeof(int16_t)));
    } else {
        pipe->InitBuffer(copyOutZeroQueue, 1, AlignBytes(this->perLoopCols, sizeof(T)));
    }
    if (this->needScaleCopy) {
        pipe->InitBuffer(scaleOutZeroQueue, 1, BLOCK_BYTES);
    }
}

template <typename T, typename TilingData>
__aicore__ inline void MoeCustomSrcToDstWithCapacity<T, TilingData>::Process()
{
    if (this->blockIdx < this->srcToDstTilingData->needCoreNum) {
        AssistInit();
        this->outTmpLocal = copyOutZeroQueue.DeQue<T>();
        if (this->needScaleCopy) {
            this->scaleLocal = scaleOutZeroQueue.DeQue<float>();
        }
        currentLoopRows = perLoopRows;
        for (int64_t loop = 0; loop < this->rowLoops; loop++) {
            if (loop == this->rowLoops - 1) {
                currentLoopRows = lastLoopRows;
            }
            CopyIn(loop);
            CopyOut(loop);
        }
        CopyOutRemain();
    }
    this->SyncAll();
}
} // namespace MoeInitRoutingCustom
#endif // MOE_CUSTOM_ROW_IDX_GATHER_DROPPAD_H