#ifndef LCF_LAYER_NORM_H_
#define LCF_LAYER_NORM_H_
#include "tensorOperate.h"

class LayerNormLayer : public TensorOperate
{
    public:
        LayerNormLayer();
        LayerNormLayer(int ifType, int iID, int iMaxSeqNum, DevicesBatchDisptach *iDevicesBatchDisptach, float initGamma, \
                       float initBeta, float inLambda,  float inAlpha, float lb, float ld, int ls);
        void SetLayerConfig(int ifType, int iID, int iMaxSeqNum, DevicesBatchDisptach *iDevicesBatchDisptach, float initGamma, \
                            float initBeta, float inLambda, float inAlpha, float lb, float ld, int ls);
        virtual ~LayerNormLayer();
        void ComputeAndCreateOutYInDTensorInfo();
        void SaveLayerWeight(int n);
        void LoadLayerWeight();
        void ComputerWeightUpdateGrad(int devID, MyVoid *iSumBeta, MyVoid *iSumGamma, MyVoid *inD, int seqNum);
        void CompuateOutY(int devID, int seqNum);
        void CompuateOutD(int devID, int seqNum);
        void UpdateWeightEnter(int devID, MyVoid *iSumBeta, MyVoid *iSumGamma);
        void ComputerYMemRealease(int devID);
        void ComputerDMemRealease(int devID);
        void ComputerPreLayerDAndGradKeepD(int devID);
        void UpdateWeightMain(int devID);
        void CreateParaAndInstance();
        void CompGrad(int devID, int seqNum);

    private:
        void SetChildLayerMomeryPointInitStatus();
        void RealeaseChildLayerMemory();
        void InitBetaGamma();
        TensorData *GetMathExpect(int devID, int seqNum);
        TensorData *GetVariance(int devID, int seqNum);
        TensorData *GetXNorm(int devID, int seqNum);
        
        TensorData *mXNorm;                                        // [batchNum][inChannels][inHeight][inWidth]
        TensorData *mMathExpect;                                   // mathExpect[batchNum]
        TensorData *mVariance;                                     // variance[batchNum]

        TensorData *gamma;                                         // Scale coff
        TensorData *beta;                                          // Shift coff
        TensorData *momentumGm;                                    // momentumGa
        TensorData *momentumBt;                                    // momentumBt
        TensorData *gammaGrad;
        TensorData *betaGrad;
        float mInitBeta, mInitGamma;
        
        int inChannels, inHeight, inWidth;
        int outChannels, outHeight, outWidth;
};

#endif
