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

class BatchNormLayer : public TensorOperate
{
    public:
        BatchNormLayer(int ifType, int iID, int iMaxSeqNum, DevicesBatchDisptach *iDevicesBatchDisptach, float initGamma, float initBeta, float inLambda, \
                       float inAlpha, float lb, float ld, int ls);
        virtual ~BatchNormLayer();
        void ComputeAndCreateOutYInDTensorInfo();
        void SetmMovingAveCoff(float userCoff);
        void SaveLayerWeight(int n);
        void LoadLayerWeight();
        void UpdateWeightEnter(int devID, MyVoid *sumGammaCMem, MyVoid *sumBetaCMem);
        void CompuateOutYInference(int devID, int seqNum);
        void CompuateOutY(int devID, int seqNum);
        void CompuateOutD(int devID, int seqNum);
        void ComputerWeightUpdateGrad(int devID, MyVoid *inDCMem, MyVoid *sumGammaCMem, MyVoid *sumBetaCMem, int seqNum);
        void ComputerYMemRealease(int devID);
        void ComputerDMemRealease(int devID);
        void UpdateWeightMain();
        void CreateParaAndInstance();
        TensorData *GetXNorm(int devID, int seqNum);
        void CompuateOutYTrain(int devID, int seqNum);
        void CompGrad(int devID, int seqNum);

    private:
        void SetChildLayerMomeryPointInitStatus();
        void RealeaseChildLayerMemory();
        void InitGammaBetaEV();
        TensorData *GetVariance(int devID, int seqNum);
        TensorData *GetMathExpect(int devID, int seqNum);

        TensorData *mXNorm;           // Norm[batchNum][inChannels][inHeight][inWidth]
        TensorData *mGamma;           // scale gamma [outChannels]
        TensorData *mBeta;            // shift beta [outChannels]
        TensorData *mGammaGrad;       // scale gamma update grad
        TensorData *mBetaGrad;        // shift beta update grad
        TensorData *mMathExpect;      // mathExpect[outChannels]
        TensorData *mVariance;        // variance [outChannels]
        TensorData *mMomentumGm;      // momentumGa[outChannels]
        TensorData *mMomentumBt;      // momentumBt[outChannels]
        TensorData *mSaveMathExpect;  // saveMathExpect [outChannels], be used while inference
        TensorData *mSaveVariance;    // saveVariance [outChannels],  be used while inference
        float mMovingAveCoff;         // moving average method inertia coff 0.9~0.99
        float mInitBeta, mInitGamma;  // init beta gamma value
        
        int inChannels, inHeight, inWidth;
        int outChannels, outHeight, outWidth;
};

#endif
