#ifndef LCF_MULTI_HEAD_ATTENTION_H_
#define LCF_MULTI_HEAD_ATTENTION_H_
#include "tensorOperate.h"
#include "attentionLayer.h"
#include "fullConnectLayer.h"
#include "fusionLayer.h"

class MultiHeadAttentionLayer : public TensorOperate
{
    public:
        MultiHeadAttentionLayer();
        virtual ~MultiHeadAttentionLayer();

        void SetLayerConfig(int ifType, int iID, int iMaxSeqNum, DevicesBatchDisptach *iDevicesBatchDisptach, int iHVec, \
                            int KVNum, int headNum, float iLambda, float iAlpha, float lb, float ld, int ls, int ifQMethod);
                                             
        void ComputeAndCreateOutYInDTensorInfo();
        void CreateParaAndInstance();

        void ComputerYMemRealease(int devID);
        void ComputerDMemRealease(int devID);
        
        void CompuateOutY(int devID, int seqNum);
        void CompuateOutD(int devID, int seqNum);
        
        void UpdateWeightMain();
        int GetQVNum();
        
        void SetFilterMat(int devID, int seqNum, float *mat);
        void UpdateDSTD(float dStd);

    private:
        void SetQKVInInfo(AttentionLayer *inAttentionLayer, FCLayer *inQFCLayer, FCLayer *inKFCLayer, FCLayer *inVFCLayer);
        void SetChildLayerMomeryPointInitStatus();
        void RealeaseChildLayerMemory();
        void SetInOutPoint();
        
        FCLayer *mQFCLayer;
        FCLayer *mKFCLayer;
        FCLayer *mVFCLayer;
        FCLayer *mMultiQFCLayer;
        FCLayer *mMultiKFCLayer;
        FCLayer *mMultiVFCLayer;
        AttentionLayer *mAttentionLayer;
        FusionLayer *mFusionLayer;
                            
        int mHidVecDim;                   // hidden layer dim equal out dim
        int mKVNum;
        int mHeadNum;
        int mIfQMethod;                  //0:(only input 0 K+V+Q) 1:(input 0 is K+V; input 1 is Q)
};

#endif