//
//  LayerNormTest.cpp
//  MNNTests
//
//  Created by MNN on 2023/07/05.
//  Copyright © 2018, Alibaba Group Holding Limited
//
#include <cmath>
#include <MNN/expr/Expr.hpp>
#include <MNN/expr/ExprCreator.hpp>
#include "MNNTestSuite.h"
#include "TestUtils.h"

using namespace MNN;
using namespace MNN::Express;

static VARP _LayerNorm(VARP x, std::vector<int32_t> axis, float epsilon, std::vector<float> gamma, std::vector<float> beta, int group = 1, bool useRMS = false) {
    std::unique_ptr<OpT> op(new OpT);
    op->main.type                         = OpParameter_LayerNorm;
    op->type                              = OpType_LayerNorm;
    op->main.value                        = new LayerNormT;
    if(gamma.size() != 0){
        op->main.AsLayerNorm()->gamma         = gamma;
    }
    if(beta.size() != 0){
        op->main.AsLayerNorm()->beta          = beta;
    }
    op->main.AsLayerNorm()->epsilon       = epsilon;
    op->main.AsLayerNorm()->axis          = axis;
    op->main.AsLayerNorm()->group         = group;
    op->main.AsLayerNorm()->useRMSNorm    = useRMS;
    return (Variable::create(Expr::create(std::move(op), {x})));
}

std::vector<float> inputdata_0 = {0.6, 1.6, 2.6, 3.6, 4.6, 5.6, 6.6, 7.6, 8.6, 9.6, 0.6, 1.6, 2.6, 3.6, 4.6, 5.6, 6.6, 7.6,
                                8.6, 9.6, 0.6, 1.6, 2.6, 3.6, 4.6, 5.6, 6.6, 7.6, 8.6, 9.6, 0.6, 1.6, 2.6, 3.6, 4.6, 5.6,
                                6.6, 7.6, 8.6, 9.6, 0.6, 1.6, 2.6, 3.6, 4.6, 5.6, 6.6, 7.6, 8.6, 9.6, 0.6, 1.6, 2.6, 3.6,
                                4.6, 5.6, 6.6, 7.6, 8.6, 9.6, 0.6, 1.6, 2.6, 3.6, 4.6, 5.6, 6.6, 7.6, 8.6, 9.6, 0.6, 1.6,
                                2.6, 3.6, 4.6, 5.6, 6.6, 7.6, 8.6, 9.6, 0.6, 1.6, 2.6, 3.6, 4.6, 5.6, 6.6, 7.6, 8.6, 9.6,
                                0.6, 1.6, 2.6, 3.6, 4.6, 5.6, 6.6, 7.6, 8.6, 9.6, 0.6, 1.6, 2.6, 3.6, 4.6, 5.6, 6.6, 7.6,
                                8.6, 9.6, 0.6, 1.6, 2.6, 3.6, 4.6, 5.6, 6.6, 7.6, 8.6, 9.6, 0.6, 1.6, 2.6, 3.6, 4.6, 5.6,
                                6.6, 7.6, 8.6, 9.6, 0.6, 1.6, 2.6, 3.6, 4.6, 5.6, 6.6, 7.6, 8.6, 9.6, 0.6, 1.6, 2.6, 3.6,
                                4.6, 5.6, 6.6, 7.6, 8.6, 9.6, 0.6, 1.6, 2.6, 3.6, 4.6, 5.6, 6.6, 7.6, 8.6, 9.6, 0.6, 1.6,
                                2.6, 3.6, 4.6, 5.6, 6.6, 7.6, 8.6, 9.6, 0.6, 1.6, 2.6, 3.6, 4.6, 5.6, 6.6, 7.6, 8.6, 9.6,
                                0.6, 1.6, 2.6, 3.6, 4.6, 5.6, 6.6, 7.6, 8.6, 9.6, 0.6, 1.6, 2.6, 3.6, 4.6, 5.6, 6.6, 7.6,
                                8.6, 9.6, 0.6, 1.6, 2.6, 3.6, 4.6, 5.6, 6.6, 7.6, 8.6, 9.6, 0.6, 1.6, 2.6, 3.6, 4.6, 5.6,
                                6.6, 7.6, 8.6, 9.6, 0.6, 1.6, 2.6, 3.6, 4.6, 5.6, 6.6, 7.6, 8.6, 9.6, 0.6, 1.6, 2.6, 3.6,
                                4.6, 5.6, 6.6, 7.6, 8.6, 9.6};
std::vector<float> tgdata_0 = {-1.34164079, -0.4472136 ,  0.4472136 ,  1.34164079, -1.34164079,
                                -0.4472136 ,  0.4472136 ,  1.34164079,  0.86824314,  1.11631261,
                                -1.11631261, -0.86824314, -1.34164079, -0.4472136 ,  0.4472136 ,
                                 1.34164079, -1.34164079, -0.4472136 ,  0.4472136 ,  1.34164079,
                                -1.34164079, -0.4472136 ,  0.4472136 ,  1.34164079, -1.34164079,
                                -0.4472136 ,  0.4472136 ,  1.34164079,  0.86824314,  1.11631261,
                                -1.11631261, -0.86824314, -1.34164079, -0.4472136 ,  0.4472136 ,
                                 1.34164079, -1.34164079, -0.4472136 ,  0.4472136 ,  1.34164079,
                                -1.34164079, -0.4472136 ,  0.4472136 ,  1.34164079, -1.34164079,
                                -0.4472136 ,  0.4472136 ,  1.34164079,  0.86824314,  1.11631261,
                                -1.11631261, -0.86824314, -1.34164079, -0.4472136 ,  0.4472136 ,
                                 1.34164079, -1.34164079, -0.4472136 ,  0.4472136 ,  1.34164079,
                                -1.34164079, -0.4472136 ,  0.4472136 ,  1.34164079, -1.34164079,
                                -0.4472136 ,  0.4472136 ,  1.34164079,  0.86824314,  1.11631261,
                                -1.11631261, -0.86824314, -1.34164079, -0.4472136 ,  0.4472136 ,
                                 1.34164079, -1.34164079, -0.4472136 ,  0.4472136 ,  1.34164079,
                                -1.34164079, -0.4472136 ,  0.4472136 ,  1.34164079, -1.34164079,
                                -0.4472136 ,  0.4472136 ,  1.34164079,  0.86824314,  1.11631261,
                                -1.11631261, -0.86824314, -1.34164079, -0.4472136 ,  0.4472136 ,
                                 1.34164079, -1.34164079, -0.4472136 ,  0.4472136 ,  1.34164079,
                                -1.34164079, -0.4472136 ,  0.4472136 ,  1.34164079, -1.34164079,
                                -0.4472136 ,  0.4472136 ,  1.34164079,  0.86824314,  1.11631261,
                                -1.11631261, -0.86824314, -1.34164079, -0.4472136 ,  0.4472136 ,
                                 1.34164079, -1.34164079, -0.4472136 ,  0.4472136 ,  1.34164079,
                                -1.34164079, -0.4472136 ,  0.4472136 ,  1.34164079, -1.34164079,
                                -0.4472136 ,  0.4472136 ,  1.34164079,  0.86824314,  1.11631261,
                                -1.11631261, -0.86824314, -1.34164079, -0.4472136 ,  0.4472136 ,
                                 1.34164079, -1.34164079, -0.4472136 ,  0.4472136 ,  1.34164079,
                                -1.34164079, -0.4472136 ,  0.4472136 ,  1.34164079, -1.34164079,
                                -0.4472136 ,  0.4472136 ,  1.34164079,  0.86824314,  1.11631261,
                                -1.11631261, -0.86824314, -1.34164079, -0.4472136 ,  0.4472136 ,
                                 1.34164079, -1.34164079, -0.4472136 ,  0.4472136 ,  1.34164079,
                                -1.34164079, -0.4472136 ,  0.4472136 ,  1.34164079, -1.34164079,
                                -0.4472136 ,  0.4472136 ,  1.34164079,  0.86824314,  1.11631261,
                                -1.11631261, -0.86824314, -1.34164079, -0.4472136 ,  0.4472136 ,
                                 1.34164079, -1.34164079, -0.4472136 ,  0.4472136 ,  1.34164079,
                                -1.34164079, -0.4472136 ,  0.4472136 ,  1.34164079, -1.34164079,
                                -0.4472136 ,  0.4472136 ,  1.34164079,  0.86824314,  1.11631261,
                                -1.11631261, -0.86824314, -1.34164079, -0.4472136 ,  0.4472136 ,
                                 1.34164079, -1.34164079, -0.4472136 ,  0.4472136 ,  1.34164079,
                                -1.34164079, -0.4472136 ,  0.4472136 ,  1.34164079, -1.34164079,
                                -0.4472136 ,  0.4472136 ,  1.34164079,  0.86824314,  1.11631261,
                                -1.11631261, -0.86824314, -1.34164079, -0.4472136 ,  0.4472136 ,
                                 1.34164079, -1.34164079, -0.4472136 ,  0.4472136 ,  1.34164079,
                                -1.34164079, -0.4472136 ,  0.4472136 ,  1.34164079, -1.34164079,
                                -0.4472136 ,  0.4472136 ,  1.34164079,  0.86824314,  1.11631261,
                                -1.11631261, -0.86824314, -1.34164079, -0.4472136 ,  0.4472136 ,
                                 1.34164079, -1.34164079, -0.4472136 ,  0.4472136 ,  1.34164079};
std::vector<float> inputdata_1 = {0.7742,  0.5332, -0.8799, -1.0676, -0.7402, -1.5074,  0.2052,  0.3648,
                                  1.5056, -0.2804,  1.2785,  1.3410,  0.5395,  0.1665, -1.4594,  0.1158,
                                  -1.8436, -0.1581, -1.5055,  0.3366,  0.4938,  0.0630,  0.5465,  0.9264,
                                  -1.0491,  2.4297,  1.9942,  0.4256,  1.3902, -0.1021, -0.9883,  0.4500};
std::vector<float> tgdata_1 = {1.1381,  0.8445, -0.8770, -1.1056, -0.4238, -1.4374,  0.8252,  1.0360,
                               0.7544, -1.7206,  0.4397,  0.5264,  0.9098,  0.4242, -1.6923,  0.3583,
                               -1.1587,  0.6996, -0.7859,  1.2451, -0.0446, -1.4518,  0.1277,  1.3688,
                               -1.4550,  1.0769,  0.7599, -0.3818,  1.3930, -0.3354, -1.3618,  0.3041};
float eps = 0.000001f;

static bool testKernel (std::vector<float> inputdata, std::vector<float> targetdata, std::vector<int> dimensions, std::vector<int> reduceAxis, float epsilon, std::vector<float> gamma, std::vector<float> beta, std::vector<float> inputQuan, std::vector<float> outputQuan, std::string testName, int precision, int group = 1) {
    int size = 1;
    for (int i = 0; i < dimensions.size(); ++i) {
        size *= dimensions[i];
    }
    int reducesize = 1;
    for (int i = 0; i < reduceAxis.size(); ++i) {
        reducesize *= dimensions[reduceAxis[i]];
    }
    MNN_ASSERT(gamma.size() == 0 || (gamma.size() >0 && reducesize == gamma.size()));
    MNN_ASSERT(gamma.size() == beta.size());
    auto input = _Input(dimensions, NCHW);
    if (inputQuan.size() > 0) {
        input->writeScaleMap(inputQuan[0], inputQuan[1]);
    }
    auto inputPtr = input->writeMap<float>();
    ::memcpy(inputPtr, inputdata.data(), input->getInfo()->size * sizeof(float));
    auto output = _LayerNorm(input, reduceAxis, epsilon, gamma, beta, group, false);
    if (outputQuan.size() > 0) {
        output->writeScaleMap(outputQuan[0], outputQuan[1]);
    }
    const float* outputPtr = output->readMap<float>();
    float ratio = 0.02;
    if (!checkVector<float>(outputPtr, targetdata.data(), size, ratio)) {
        MNN_ERROR("%s failed: data dimension=[", testName.c_str());
        for (int i = 0; i < dimensions.size(); ++i) {
            if (i < dimensions.size() - 1) {MNN_PRINT("%d, ", dimensions[i]);}
            else {MNN_PRINT("%d], reduce axis=[", dimensions[i]);};
        }
        for (int i = 0; i < reduceAxis.size(); ++i) {
            if (i < reduceAxis.size() - 1) {MNN_PRINT("%d, ", reduceAxis[i]);}
            else {MNN_PRINT("%d]\n", reduceAxis[i]);};
        }
        return false;
    }
    return true;
}

static int nc4hw4Offset(int n, int c, int plane, int batch) {
    return (c % 4) + 4 * plane * n + 4 * plane * batch * (c / 4);
}

static void computeChannelLayerNorm(const std::vector<float>& input, std::vector<float>& output, int batch, int channel,
                                    const std::vector<float>& gamma, const std::vector<float>& beta) {
    output.resize(batch * channel);
    for (int n = 0; n < batch; ++n) {
        float mean = 0.0f;
        for (int c = 0; c < channel; ++c) {
            mean += input[n * channel + c];
        }
        mean /= channel;
        float variance = 0.0f;
        for (int c = 0; c < channel; ++c) {
            float v = input[n * channel + c] - mean;
            variance += v * v;
        }
        variance /= channel;
        float inv = 1.0f / std::sqrt(variance + eps);
        for (int c = 0; c < channel; ++c) {
            float v = (input[n * channel + c] - mean) * inv;
            if (!gamma.empty()) {
                v = v * gamma[c] + (beta.empty() ? 0.0f : beta[c]);
            }
            output[n * channel + c] = v;
        }
    }
}

static bool checkNC4HW4Logical(VARP output, const std::vector<float>& expected, int batch, int channel,
                               const char* testName) {
    auto ptr = output->readMap<float>();
    std::vector<float> actual(batch * channel);
    for (int n = 0; n < batch; ++n) {
        for (int c = 0; c < channel; ++c) {
            actual[n * channel + c] = ptr[nc4hw4Offset(n, c, 1, batch)];
        }
    }
    if (!checkVector<float>(actual.data(), expected.data(), batch * channel, 0.02f)) {
        MNN_ERROR("%s failed!\n", testName);
        return false;
    }
    return true;
}

class LayerNormTest : public MNNTestCase {
public:
    virtual ~LayerNormTest() = default;
    virtual bool run(int precision) {
        { // test 1.
            std::vector<int32_t> axis = {0, 1, 2};
            std::vector<int> dims = {1, 4, 1, 2};
            // set input data
            std::vector<float> inputdata = {-1.0, -2.0, 3.0, 4.0, -5.0, -6.0, 7.0, 8.0};
            std::vector<float> tgdata = {0.03527864, 0.0242914, 0.14944272, 0.1714172, -0.24249224, -0.22996021, 0.68665631, 0.66994695};
            std::vector<float> gammaData = {0.1f, 0.2f, 0.3f, 0.4f};
            std::vector<float> betaData = {0.08f, 0.06f, 0.16f, 0.15f};
            
            bool testSuc = testKernel(inputdata, tgdata, dims, axis, eps, gammaData, betaData, {}, {}, "Float LayerNorm Test", 1);
            if (!testSuc) {
                return false;
            }
        }
        { // test 2.
            std::vector<int32_t> axis = {0, 1, 2};
            std::vector<int> dims = {1, 4, 1, 2};
            // set input data
            std::vector<float> inputdata = {-1.0, -2.0, 3.0, 4.0, -5.0, -6.0, 7.0, 8.0};
            std::vector<float> tgdata = {0.03527864, 0.0242914, 0.14944272, 0.1714172, -0.24249224, -0.22996021, 0.68665631, 0.66994695};
            std::vector<float> inputQuan = {0.063745, 2.0};
            std::vector<float> outputQuan = {0.0095, 0.0};
            std::vector<float> gammaData = {0.1f, 0.2f, 0.3f, 0.4f};
            std::vector<float> betaData = {0.08f, 0.06f, 0.16f, 0.15f};
            
            bool testSuc = testKernel(inputdata, tgdata, dims, axis, eps, gammaData, betaData, inputQuan, outputQuan, "Int8 LayerNorm Test", 1);
            if (!testSuc) {
                return false;
            }
        }
        { // test 3.
            std::vector<int32_t> axis = {0, 1, 2};
            std::vector<int> dims = {1, 4, 1, 2};
            std::vector<float> gammaData = {};
            std::vector<float> betaData = {};
            // set input data
            std::vector<float> inputdata = {-1.0, -2.0, 3.0, 4.0, -5.0, -6.0, 7.0, 8.0};
            std::vector<float> tgdata = {-0.4472136, -0.55708601, 0.4472136, 0.55708601, -1.34164079, -1.29986737, 1.34164079, 1.29986737};
            
            bool testSuc = testKernel(inputdata, tgdata, dims, axis, eps, gammaData, betaData, {}, {}, "Float LayerNorm Test", 1);
            if (!testSuc) {
                return false;
            }
        }
        { // test 4.
            std::vector<int32_t> axis = {0, 1, 2};
            std::vector<int> dims = {1, 4, 1, 2};
            std::vector<float> gammaData = {};
            std::vector<float> betaData = {};
            // set input data
            std::vector<float> inputdata = {-1.0, -2.0, 3.0, 4.0, -5.0, -6.0, 7.0, 8.0};
            std::vector<float> tgdata = {-0.4472136, -0.55708601, 0.4472136, 0.55708601, -1.34164079, -1.29986737, 1.34164079, 1.29986737};
            std::vector<float> inputQuan = {0.063745, 2.0};
            std::vector<float> outputQuan = {0.0105, 0.0};
            
            bool testSuc = testKernel(inputdata, tgdata, dims, axis, eps, gammaData, betaData, inputQuan, outputQuan, "Int8 LayerNorm Test", 1);
            if (!testSuc) {
                return false;
            }
        }
        { // test 5.
            std::vector<int32_t> axis = {2, 3};
            std::vector<int> dims = {6, 10, 2, 2};
            std::vector<float> gammaData = {};
            std::vector<float> betaData = {};

            bool testSuc = testKernel(inputdata_0, tgdata_0, dims, axis, eps, gammaData, betaData, {}, {}, "Float LayerNorm Test", 1);
            if (!testSuc) {
                return false;
            }
        }
        {
            std::vector<int32_t> axis = {2, 3};
            std::vector<int> dims = {6, 10, 2, 2};
            std::vector<float> gammaData = {};
            std::vector<float> betaData = {};
            std::vector<float> inputQuan = {0.0752, 0.f};
            std::vector<float> outputQuan = {0.0105, 0.f};
            
            bool testSuc = testKernel(inputdata_0, tgdata_0, dims, axis, eps, gammaData, betaData, inputQuan, outputQuan, "Int8 LayerNorm Test", 1);
            if (!testSuc) {
                return false;
            }
        }
        { // Group Norm without axis 
            std::vector<int> dims = {2, 4, 2, 2};
            auto input = _Input(dims, NCHW);
            bool testSuc = testKernel(inputdata_1, tgdata_1, dims, {}, eps, {}, {}, {}, {}, "Flaot GroupNorm Test", 1, 4);
            if (!testSuc) {
                return false;
            }
        }
        return true;
    }
};
MNNTestSuiteRegister(LayerNormTest, "op/layernorm");

class LayerNormC4Test : public MNNTestCase {
public:
    virtual ~LayerNormC4Test() = default;
    bool runOne(int batch, int channel) {
        const int physicalSize = batch * UP_DIV(channel, 4) * 4;
        std::vector<float> logical(batch * channel);
        std::vector<float> packed(physicalSize, 0.0f);
        std::vector<float> gamma(channel);
        std::vector<float> beta(channel);
        for (int i = 0; i < (int)logical.size(); ++i) {
            logical[i] = (float)((i % 17) - 8) * 0.11f + (float)(i % 5) * 0.03f;
        }
        for (int c = 0; c < channel; ++c) {
            gamma[c] = 0.7f + (float)(c % 11) * 0.05f;
            beta[c] = -0.2f + (float)(c % 7) * 0.04f;
        }
        for (int n = 0; n < batch; ++n) {
            for (int c = 0; c < channel; ++c) {
                packed[nc4hw4Offset(n, c, 1, batch)] = logical[n * channel + c];
            }
        }
        auto input = _Input({batch, channel, 1, 1}, NC4HW4);
        ::memcpy(input->writeMap<float>(), packed.data(), packed.size() * sizeof(float));
        input->unMap();
        std::unique_ptr<OpT> op(new OpT);
        op->main.type = OpParameter_LayerNorm;
        op->type = OpType_LayerNorm;
        op->defaultDimentionFormat = MNN_DATA_FORMAT_NC4HW4;
        op->main.value = new LayerNormT;
        op->main.AsLayerNorm()->gamma = gamma;
        op->main.AsLayerNorm()->beta = beta;
        op->main.AsLayerNorm()->epsilon = eps;
        op->main.AsLayerNorm()->axis = {1};
        auto output = Variable::create(Expr::create(std::move(op), {input}));
        std::vector<float> expected;
        computeChannelLayerNorm(logical, expected, batch, channel, gamma, beta);
        return checkNC4HW4Logical(output, expected, batch, channel, "LayerNormC4Test");
    }

    virtual bool run(int precision) {
        return runOne(2, 8) && runOne(13, 1024);
    }
};

class BinaryLayerNormC4Test : public MNNTestCase {
public:
    virtual ~BinaryLayerNormC4Test() = default;
    bool runOne(int batch, int channel) {
        const int physicalSize = batch * UP_DIV(channel, 4) * 4;
        std::vector<float> logical0(batch * channel);
        std::vector<float> logical1(batch * channel);
        std::vector<float> packed0(physicalSize, 0.0f), packed1(physicalSize, 0.0f), sumLogical(batch * channel);
        for (int i = 0; i < (int)logical0.size(); ++i) {
            logical0[i] = (float)((i % 17) - 8) * 0.11f + (float)(i % 5) * 0.03f;
            logical1[i] = (float)((i % 13) - 6) * -0.07f + (float)(i % 3) * 0.02f;
        }
        for (int n = 0; n < batch; ++n) {
            for (int c = 0; c < channel; ++c) {
                int logicalIndex = n * channel + c;
                int packedIndex = nc4hw4Offset(n, c, 1, batch);
                packed0[packedIndex] = logical0[logicalIndex];
                packed1[packedIndex] = logical1[logicalIndex];
                sumLogical[logicalIndex] = logical0[logicalIndex] + logical1[logicalIndex];
            }
        }
        std::vector<float> gamma(channel);
        std::vector<float> beta(channel, 0.0f);
        for (int c = 0; c < channel; ++c) {
            gamma[c] = 0.7f + (float)(c % 11) * 0.05f;
        }
        auto input0 = _Input({batch, channel, 1, 1}, NC4HW4);
        auto input1 = _Input({batch, channel, 1, 1}, NC4HW4);
        ::memcpy(input0->writeMap<float>(), packed0.data(), packed0.size() * sizeof(float));
        ::memcpy(input1->writeMap<float>(), packed1.data(), packed1.size() * sizeof(float));
        input0->unMap();
        input1->unMap();

        std::unique_ptr<OpT> op(new OpT);
        op->main.type = OpParameter_LayerNorm;
        op->type = OpType_LayerNorm;
        op->defaultDimentionFormat = MNN_DATA_FORMAT_NC4HW4;
        op->main.value = new LayerNormT;
        op->main.AsLayerNorm()->gamma = gamma;
        op->main.AsLayerNorm()->beta = beta;
        op->main.AsLayerNorm()->epsilon = eps;
        op->main.AsLayerNorm()->axis = {1};
        auto expr = Expr::create(std::move(op), {input0, input1}, 2);
        auto sumOutput = Variable::create(expr, 0);
        auto normOutput = Variable::create(expr, 1);

        std::vector<float> expectedNorm;
        computeChannelLayerNorm(sumLogical, expectedNorm, batch, channel, gamma, beta);
        return checkNC4HW4Logical(sumOutput, sumLogical, batch, channel, "BinaryLayerNormC4SumTest") &&
               checkNC4HW4Logical(normOutput, expectedNorm, batch, channel, "BinaryLayerNormC4NormTest");
    }

    virtual bool run(int precision) {
        return runOne(2, 8) && runOne(13, 1024);
    }
};

MNNTestSuiteRegister(LayerNormC4Test, "op/layernorm/c4");
MNNTestSuiteRegister(BinaryLayerNormC4Test, "op/layernorm/c4_binary");