#include "ArgMinExecution.hpp"
#include "core/MusaBackend.hpp"

namespace MNN {
namespace MUSA {

template<typename T>
__global__ void ArgMinKernel(const T* input, int* output,
                              int outerSize, int axisSize, int innerSize) {
    int outerIdx = blockIdx.x * blockDim.x + threadIdx.x;
    
    if (outerIdx < outerSize) {
        T minVal = input[outerIdx * axisSize * innerSize];
        int minIdx = 0;
        
        for (int i = 0; i < axisSize; i++) {
            for (int j = 0; j < innerSize; j++) {
                int idx = (outerIdx * axisSize + i) * innerSize + j;
                if (input[idx] < minVal) {
                    minVal = input[idx];
                    minIdx = i;
                }
            }
        }
        
        output[outerIdx] = minIdx;
    }
}

ArgMinExecution::ArgMinExecution(const std::vector<Tensor*>& inputs, const MNN::Op* op, Backend* backend)
    : Execution(inputs, {}, backend) {
    mBackend = static_cast<MusaBackend*>(backend);
    mOp = op->main_as_ArgMin();
}

ErrorCode ArgMinExecution::onResize(const std::vector<Tensor*>& inputs, const std::vector<Tensor*>& outputs) {
    auto input = inputs[0];
    auto output = outputs[0];
    
    mAxis = mOp->axis();
    if (mAxis < 0) {
        mAxis += input->dimensions();
    }
    
    mOuterSize = 1;
    for (int i = 0; i < mAxis; i++) {
        mOuterSize *= input->length(i);
    }
    
    mAxisSize = input->length(mAxis);
    
    mInnerSize = 1;
    for (int i = mAxis + 1; i < input->dimensions(); i++) {
        mInnerSize *= input->length(i);
    }
    
    int threads = 256;
    int blocks = (mOuterSize + threads - 1) / threads;
    
    mDim3Grid = {blocks, 1, 1};
    mDim3Block = {threads, 1, 1};
    
    return NO_ERROR;
}

ErrorCode ArgMinExecution::onExecute(const std::vector<Tensor*>& inputs, const std::vector<Tensor*>& outputs) {
    auto input = inputs[0];
    auto output = outputs[0];
    
    auto inputPtr = input->host<float>();
    auto outputPtr = output->host<int>();
    
    ArgMinKernel<<<mDim3Grid, mDim3Block>>>(
        inputPtr, outputPtr,
        mOuterSize, mAxisSize, mInnerSize
    );
    
    musaError_t err = musaGetLastError();
    if (err != musaSuccess) {
        return COMPUTE_NO_SUPPORT;
    }
    
    return NO_ERROR;
}

class ArgMinCreator : public Creator {
public:
    virtual Execution* onCreate(const std::vector<Tensor*>& inputs, const MNN::Op* op, Backend* backend) const override {
        return new ArgMinExecution(inputs, op, backend);
    }
};

MNNCreatorRegister<ArgMinCreator> gArgMinRegistration(OpType_ArgMin);

} // namespace MUSA
} // namespace MNN