/*
 *  ******************************************************************************
 *  *
 *  *
 *  * This program and the accompanying materials are made available under the
 *  * terms of the Apache License, Version 2.0 which is available at
 *  * https://www.apache.org/licenses/LICENSE-2.0.
 *  *
 *  * See the NOTICE file distributed with this work for additional
 *  * information regarding copyright ownership.
 *  * Unless required by applicable law or agreed to in writing, software
 *  * distributed under the License is distributed on an "AS IS" BASIS, WITHOUT
 *  * WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the
 *  * License for the specific language governing permissions and limitations
 *  * under the License.
 *  *
 *  * SPDX-License-Identifier: Apache-2.0
 *  *****************************************************************************
 */

//
// @author Paul Dubs
// @author Adam Gibson
//

#ifndef LIBND4J_ATTENTIONHELPER_CPP
#define LIBND4J_ATTENTIONHELPER_CPP
#include "../AttentionHelper.h"
#include <indexing/NDIndexUtils.h>
#include <helpers/AttentionHelper.h>
#include <ops/declarable/CustomOperations.h>
#include <ops/declarable/helpers/batched_gemm.h>
#if NOT_EXCLUDED(OP_multi_head_dot_product_attention)

namespace sd {

NDArray AttentionHelper::multiHeadProject(NDArray *input, NDArray *projectionMatrix,
                                          LaunchContext *context) {
  auto miniBatchSize = input->sizeAt(0);
  auto seqLength = input->sizeAt(2);
  auto numHeads = projectionMatrix->sizeAt(0);
  auto projectedSize = projectionMatrix->sizeAt(1);

  std::vector<sd::LongType> epsPermVec = {1, 0,2};
  auto inputPerm = input->permute(epsPermVec, false, false);  //[batch, nIn, timeSteps] -> [nIn, batch, timeSteps]
  std::vector<sd::LongType> inputPermShape = {input->sizeAt(1), (miniBatchSize * seqLength)};
  auto inputPrep = inputPerm->reshape('c', inputPermShape);  //[nIn, batch*timeSteps]
  std::vector<sd::LongType> projectionMatrixShape = {numHeads * projectionMatrix->sizeAt(1), projectionMatrix->sizeAt(2)};
  auto projectionPrep = projectionMatrix->reshape(
      'c',
      projectionMatrixShape);  //[nHeads, hS, nIn] -> [nHeads*hS, nIn]

  std::vector<LongType> projectedShape = {numHeads * projectionMatrix->sizeAt(1), (miniBatchSize * seqLength)};
  NDArray projected('c',projectedShape, input->dataType(),
                    context);  //[nHeads*hS, batch*timeSteps]
  ops::matmul mmul;
  mmul.execute({&projectionPrep, &inputPrep}, {&projected});

  projected.reshapei({numHeads, projectedSize, miniBatchSize, seqLength});
  projected.permutei({2, 0, 1, 3}, false, false);  //[minibatch, numHeads, projectedSize, seqLength]

  return projected;
}


/**
 * @param shape
 * @return
 */
NDArray * AttentionHelper::lowerTriangularMask(std::vector<LongType> *shape) {
  auto rowIndexOnes = NDArrayFactory::valueOf(*shape,1,'c');
  auto colIndexOnes = NDArrayFactory::valueOf(*shape, 1, 'c');
  ops::cumsum cumsum;
  auto rowCumSum = cumsum.evaluate({rowIndexOnes},{},{-2,0},{});
  auto colsCumSum = cumsum.evaluate({colIndexOnes}, {}, {-1, 0}, {});
  ops::greater_equal greaterEqual;
  auto ret = greaterEqual.evaluate({rowCumSum.at(0),colsCumSum.at(0)});
  return ret[0];
}

/**
 * @param query
 * @param value
 * @return
 */
NDArray *AttentionHelper::computeCasualMask(NDArray *query, NDArray *value, bool multiHead) {
  if(multiHead) {
    auto qSeqLength = query->sizeAt(1);
    auto vSeqLength = value != nullptr ? value->sizeAt(1) : qSeqLength;
    ops::matrix_band_part matrixBandPart;
    auto ones = NDArrayFactory::create('c',{1,qSeqLength,vSeqLength}, INT32);
    int assignVal = 1;
    ones->assign(assignVal);
    auto lower = matrixBandPart.evaluate({ones},{},{-1,0});
    auto ret = lower.at(0)->cast(BOOL);
    delete ones;
    return ret;

  } else {
    std::vector<LongType> causalMaskShape2;
    causalMaskShape2.push_back(query->sizeAt(0));
    //4d
    if(query->rankOf() > 3)
      causalMaskShape2.push_back(query->sizeAt(1));

    causalMaskShape2.push_back(query->sizeAt(-2));
    causalMaskShape2.push_back(value->sizeAt(-2));

    auto ret  = lowerTriangularMask(&causalMaskShape2);
    return ret;

  }

}


/**
 * @param query
 * @param value
 * @param attentionMask
 * @param useCausalMask
 * @return
 */
NDArray *AttentionHelper::computeAttentionMask(NDArray *query, NDArray *value, NDArray *queryMask, NDArray *valueMask,
                                               NDArray *attentionMask, bool useCausalMask) {
  auto internalQueryMask = queryMask;
  auto internalValueMask = valueMask;
  NDArray *autoMask = nullptr;
  ops::create_view createView;
  ops::boolean_and booleanAnd;
  auto all = NDIndexUtils::createAll();
  auto newAxis = NDIndexUtils::createNewAxis();

  if (internalQueryMask != nullptr && !internalQueryMask->isEmpty()) {
    internalQueryMask = queryMask->cast(BOOL);
    if (autoMask != nullptr && !autoMask->isEmpty()) {
      autoMask = createView.evaluate({internalQueryMask, all, all, newAxis}).at(0);
    }
  }

  if (valueMask != nullptr && !valueMask->isEmpty()) {
    internalValueMask = valueMask->cast(BOOL);
    auto mask = createView.evaluate({internalValueMask, all, newAxis, all}).at(0);
    if (autoMask == nullptr || autoMask->isEmpty()) {
      autoMask = mask;
    } else {
      autoMask = booleanAnd.evaluate({autoMask, mask}).at(0);
    }
  }

  if (useCausalMask) {
    auto mask = computeCasualMask(query, value, false);
    if (autoMask == nullptr) {
      autoMask = mask;
    } else {
      autoMask = booleanAnd.evaluate({autoMask, mask}).at(0);
    }
  }

  if (autoMask != nullptr && !autoMask->isEmpty()) {
    if (attentionMask == nullptr || attentionMask->isEmpty()) {
      return autoMask;
    } else {
      auto ret = booleanAnd.evaluate({attentionMask, autoMask}).at(0);
      return ret;
    }
  }

  delete all;
  delete newAxis;

  return autoMask;
}

NDArray * AttentionHelper::mergeMasks(NDArray *x, NDArray *y) {
  if(x == nullptr || x->isEmpty()) {
    return y;
  }

  if (y == nullptr || y->isEmpty()) {
    return x;
  }

  ops::boolean_and booleanAnd;
  auto ret = booleanAnd.evaluate({x,y});
  return ret.at(0);
}

void AttentionHelper::applyAttentionScores(NDArray *scores, NDArray *value, NDArray *scoresMask,
                                           double dropout, int randomSeed, NDArray *applyScoresOut, NDArray *attentionLogits,
                                           NDArray *dropoutMask) {
  ops::boolean_not booleanNot;
  ops::softmax softmax;
  ops::dropout dropoutOp;
  ops::matmul matmul;

  int softmaxDim = -1;
  if (scoresMask != nullptr && !scoresMask->isEmpty()) {
    REQUIRE_TRUE(scoresMask->sizeAt(-2) == 1 || scoresMask->sizeAt(-2) == scores->sizeAt(-2),0,
                 "Scores mask must be either broadcastable or equal to scores shape. scores size at -2: was: %i scores size at -2 was: %i",scoresMask->sizeAt(-2),scores->sizeAt(-2));

    REQUIRE_TRUE(scoresMask->sizeAt(-1) == scores->sizeAt(-1),0,
                 "Scores mask must be either broadcastable or equal to scores shape. scores size at -1: was: %i scores size at -1 was: %i",scoresMask->sizeAt(-1),scores->sizeAt(-1));

    auto castedScoresMask = scoresMask->cast(BOOL);
    auto paddingMask = booleanNot.evaluate({castedScoresMask}).at(0);
    auto paddingMaskCast = paddingMask->cast(scores->dataType());
    if (attentionLogits->dataType() == BFLOAT16) {
      auto minus =  65504 * *paddingMaskCast;
      *attentionLogits -= *minus;
      delete minus;
    } else {
      auto minus = 1.0e9 * *paddingMask;
      *attentionLogits -= *minus;
      delete minus;
    }

    if(paddingMaskCast != paddingMask) {
      delete paddingMaskCast;
    }

    if(scoresMask != castedScoresMask) {
      delete castedScoresMask;
    }
  }



  softmax.execute({attentionLogits},{scores},{},{softmaxDim});
  auto weights = scores;

  if (dropout > 0) {
    dropoutOp.execute({weights},{weights,dropoutMask},{dropout},{randomSeed});
  }

  //batch size, tq tv
  //batch size tv dim
  //output: batch size, tq dim
  matmul.execute({weights,value},{applyScoresOut});

}

void AttentionHelper::dotProductAttentionBpHelper(NDArray *query, NDArray *key, NDArray *values,
                                                  double scale,
                                                  NDArray *dLdq, NDArray *dLdk, NDArray *dLdv, NDArray *eps, LongType dropoutSeed, NDArray *qMask, NDArray *vMask, bool useCausalMask, double dropout, bool training,
                                                  NDArray *attentionScoresWeights, NDArray *attentionLogits,
                                                  NDArray *dropoutMask) {
  ops::matmul_bp matMulBp;
  ops::softmax_bp softmaxBp;
  NDArray dldW(attentionScoresWeights->shapeInfo());
  NDArray dldS(attentionScoresWeights->shapeInfo());
  NDArray * mask = nullptr;
  NDArray *causalPointer = nullptr;

  if(useCausalMask) {
    std::vector<LongType> causalMaskShape2;
    causalMaskShape2.push_back(attentionLogits->sizeAt(0));
    //4d
    if(attentionLogits->rankOf() > 3)
      causalMaskShape2.push_back(attentionLogits->sizeAt(1));

    for(int i = attentionLogits->rankOf() - 2; i < attentionLogits->rankOf(); i++) {
      causalMaskShape2.push_back(attentionLogits->sizeAt(i));
    }
    causalPointer = lowerTriangularMask(&causalMaskShape2);
  }

  mask = mergeMasks(vMask,causalPointer);



  matMulBp.execute({attentionScoresWeights,values,eps},{&dldW,dLdv},{},{});
  if(dropout > 0.0 && training) {
    ops::dropout_bp dropoutOp;
    auto inputs = {attentionScoresWeights,dropoutMask,&dldW};
    dropoutOp.execute(inputs,{&dldW},{dropout},{dropoutSeed},{false});
  }


  softmaxBp.execute({attentionLogits,&dldW,attentionScoresWeights},{&dldS},{},{-1},{});



  if(scale != 0.0 && scale != 1.0) {
    dldS *= scale;
  }

  // Initialize times as a scalar placeholder (will be reassigned if mask is present)
  NDArray times(query->dataType(), query->getContext(), true);
  if(mask != nullptr && !mask->isEmpty()) {
    ops::expand_dims expandDims;
    auto maskCast = mask->cast(query->dataType());
    auto mask2 = *maskCast * 1e9;
    times = *mask2;
    dldS *= times;
    delete mask2;
    delete maskCast;

  }



  matMulBp.execute({query,key,&dldS},{dLdq,dLdk},{},{0,1,0});
}




/**
   *
   * @param query
   * @param key
   * @param scoreMode
   * @param scale
   * @return
 */
void AttentionHelper::attentionBpHelper(NDArray *query, NDArray *key, NDArray *values, double scale, NDArray *dLdq,
                                        NDArray *dLdk, NDArray *dLdv, NDArray *eps,
                                        LongType dropoutSeed,
                                        NDArray *qMask, NDArray *vMask,
                                        bool useCausalMask, double dropout, bool training, NDArray *attentionScoresOut,
                                        NDArray *attentionScoresWeights,
                                        NDArray *attentionScoresLogits,
                                        NDArray *dropoutMask) {
  dotProductAttentionBpHelper(query, key, values, scale, dLdq, dLdk, dLdv, eps, dropoutSeed, qMask, vMask,
                              useCausalMask, dropout, training, attentionScoresWeights, attentionScoresLogits,
                              dropoutMask);


}

/**
   *
   * @param query
   * @param key
   * @param scoreMode
   * @param scale
   * @return
 */
void AttentionHelper::attentionHelper(NDArray *query, NDArray *key, double scale, NDArray *attentionLogits) {
  ops::matmul matmul3;
  matmul3.execute({query,key},{attentionLogits},{},{0,1});
  if(scale != 0.0 && scale != 1.0) {
    *attentionLogits *= scale;
  }

  // Clamp attention logits to prevent numerical overflow in subsequent softmax
  // Values beyond this range would produce Inf in exp() which leads to NaN
  // Use clipbyvalue op for proper clamping
  ops::clipbyvalue clipOp;
  clipOp.execute({attentionLogits}, {attentionLogits}, {-1e4, 1e4}, {});
}




/**
 * @param inputs
 * @param mask
 * @param training
 * @param returnAttentionScores
 * @param useCausalMask
 */
void AttentionHelper::doAttentionBp(std::vector<NDArray *> &inputs, std::vector<NDArray *> &masks, bool training,
                                    bool useCausalMask, double dropout, double scale, std::vector<NDArray *> outputs,
                                    LongType dropoutSeed) {
  auto q = inputs[0];
  auto v = inputs[1];
  auto k = inputs[2];
  auto attentionScoresOut = inputs[3];
  auto attentionScoresWeights = inputs[4];
  auto attentionScoresLogits = inputs[5];
  auto eps = inputs[6];

  auto dropoutMask = inputs.size() > 7 ? inputs[7] : inputs[7];

  ops::expand_dims expandDims;
  ops::ones_as onesAs;
  ops::shape_of shapeOf;
  ops::concat concatOp;
  ops::create_view createView;
  auto qMask = masks.size() > 0 ? masks[0] : nullptr;
  auto vMask = masks.size() > 1 ? masks[1] : nullptr;
  auto vmaskInternal = vMask;
  auto qMaskInternal = qMask;
  if(vMask != nullptr && !vMask->isEmpty() && vMask->rankOf() < v->rankOf()) {
    vmaskInternal = expandDims.evaluate({vMask},{},{-2}).at(0);
  }

  if(qMask != nullptr && !qMask->isEmpty()) {
    qMaskInternal = expandDims.evaluate({qMaskInternal},{},{-1}).at(0);
  }


  auto dLdq = outputs[0];
  auto dLdv = outputs[1];
  auto dLdk = outputs[2];
  attentionBpHelper(q, k, v, scale, dLdq, dLdk, dLdv, eps, dropoutSeed, qMaskInternal, vmaskInternal, useCausalMask,
                    dropout, training, attentionScoresOut, attentionScoresWeights, attentionScoresLogits, dropoutMask);

}


/**
 * @param inputs
 * @param mask
 * @param training
 * @param returnAttentionScores
 * @param useCausalMask
 */
void AttentionHelper::doAttention(std::vector<NDArray *> &inputs, std::vector<NDArray *> &masks, bool training,
                                  bool useCausalMask, double dropout, double scale, NDArray *attentionScores,
                                  int dropoutSeed, NDArray *applyScoresOut, NDArray *attentionLogits,
                                  NDArray *dropoutMask) {
  auto q = inputs[0];
  auto v = inputs[1];
  auto k = inputs.size() > 2 ? inputs[2]  : v;
  auto concatWeights = inputs.size() > 3 ? inputs[3] : nullptr;

  ops::expand_dims expandDims;
  ops::ones_as onesAs;
  ops::shape_of shapeOf;
  ops::concat concatOp;
  ops::create_view createView;
  auto qMask = masks.size() > 0 ? masks[0] : nullptr;
  auto vMask = masks.size() > 1 ? masks[1] : nullptr;
  auto vmaskInternal = vMask;
  auto qMaskInternal = qMask;

  NDArray *casualPointer = nullptr;
  //inputs: query and value
  //shape: batch_size Tq dim (batch_size Tv dim)
  //note this does not apply softmax yet, we are just computing logits here
  attentionHelper(q, k, scale, attentionLogits);

  if(vMask != nullptr && !vMask->isEmpty() && vMask->rankOf() < v->rankOf()) {
    vmaskInternal = expandDims.evaluate({vMask},{},{-2}).at(0);
  }

  if(useCausalMask) {
    std::vector<LongType> causalMaskShape2;
    causalMaskShape2.push_back(attentionScores->sizeAt(0));
    //4d
    if(attentionScores->rankOf() > 3)
      causalMaskShape2.push_back(attentionScores->sizeAt(1));

    for(int i = attentionScores->rankOf() - 2; i < attentionScores->rankOf(); i++) {
      causalMaskShape2.push_back(attentionScores->sizeAt(i));
    }
    casualPointer = lowerTriangularMask(&causalMaskShape2);
  }

  auto scoresMask = mergeMasks(vmaskInternal,casualPointer);

  //compute actual softmax now
  if(training) {
    applyAttentionScores(attentionScores, v, scoresMask, dropout, dropoutSeed, applyScoresOut, attentionLogits,
                         dropoutMask);
  } else {
    applyAttentionScores(attentionScores, v, scoresMask, 0, dropoutSeed, applyScoresOut, attentionLogits, dropoutMask);
  }
  //inputs: scores:  batch size tq tv value:batch size, tv,dim scoresmask: batch size 1 tv or batch size tq tv
  if(qMask != nullptr && !qMask->isEmpty()) {
    qMaskInternal = expandDims.evaluate({qMaskInternal},{},{-1}).at(0);
    auto casted = qMaskInternal->cast(attentionScores->dataType());
    *attentionScores *= *casted;
  }

}


void AttentionHelper::multiHeadProjectBp(NDArray *input, NDArray *projectionMatrix,
                                         NDArray *eps,
                                         NDArray *dLdInput, NDArray *dLdProjectionMatrix, LaunchContext *context) {
  auto miniBatchSize = input->sizeAt(0);
  auto seqLength = input->sizeAt(2);
  auto numHeads = projectionMatrix->sizeAt(0);
  auto projectedSize = projectionMatrix->sizeAt(1);

  std::vector<sd::LongType> epsPermVec = {1, 2, 0, 3};
  auto epsPerm = eps->permute(epsPermVec, false, false);
  std::vector<sd::LongType> epsReshapeVec = {numHeads * projectedSize, miniBatchSize * seqLength};
  auto epsReshaped = epsPerm->reshape('c', epsReshapeVec);

  std::vector<sd::LongType> inputPermVec = {1, 0, 2};
  auto inputPerm = input->permute(inputPermVec, false, false);
  std::vector<sd::LongType> inputPermShape = {input->sizeAt(1), miniBatchSize * seqLength};
  auto inputPrep = inputPerm->reshape('c',inputPermShape,false);
  std::vector<sd::LongType> projectionMatrixShape = {numHeads * projectionMatrix->sizeAt(1), projectionMatrix->sizeAt(2)};
  auto projectionPrep =
      projectionMatrix->reshape('c', projectionMatrixShape);

  ops::matmul_bp mmulBp;
  NDArray dLdProjectionPrep(projectionPrep->shapeInfo(), false, context);
  NDArray dLdInputPrep(inputPrep->shapeInfo(), false, context);
  mmulBp.execute({projectionPrep, inputPrep, epsReshaped}, std::vector<NDArray *>{&dLdProjectionPrep, &dLdInputPrep},
                 {}, {}, {});

  dLdProjectionPrep.reshapei({numHeads, projectionMatrix->sizeAt(1), projectionMatrix->sizeAt(2)});
  dLdProjectionMatrix->assign(&dLdProjectionPrep);

  dLdInputPrep.reshapei({input->sizeAt(1), miniBatchSize, seqLength});
  dLdInputPrep.permutei({1, 0, 2}, false, false);
  dLdInput->assign(&dLdInputPrep);

  delete epsReshaped;
  delete projectionPrep;

}
}  // namespace sd
#endif

#endif