/* ******************************************************************************
 *
 *
 * 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
 ******************************************************************************/

//
//  xw_plus_b op. Created by GS <george@skymind.io> 31.01.2018
//  @author Oleg Semeniv <oleg.semeniv@gmail.com>
//
//  Fixed to handle higher-rank inputs (e.g., from ONNX Gemm with batched input)
//  and corrected bias addition to always use row vector broadcasting.
//

#include <system/op_boilerplate.h>
#if NOT_EXCLUDED(OP_xw_plus_b)

#include <helpers/MmulHelper.h>
#include <ops/declarable/CustomOperations.h>
#include <ops/declarable/helpers/matmul.h>

namespace sd {
namespace ops {
CUSTOM_OP_IMPL(xw_plus_b, 3, 1, false, 0, 0) {
  bool aTranspose = (block.getIArguments()->size() > 0 ? INT_ARG(0) == 1 : false);
  bool bTranspose = (block.getIArguments()->size() > 1 ? INT_ARG(1) == 1 : false);
  bool cTranspose = (block.getIArguments()->size() > 2 ? INT_ARG(2) == 1 : false);

  auto x = INPUT_VARIABLE(0);
  auto w = INPUT_VARIABLE(1);
  auto b = INPUT_VARIABLE(2);
  auto z = OUTPUT_VARIABLE(0);

  if (x->isEmpty() || w->isEmpty() || b->isEmpty()) return Status::OK;

  // Handle higher rank inputs by reshaping to 2D for matmul
  // This supports inputs like [batch, 1, 1, hidden] from ONNX pooler operations
  NDArray* xEffective = nullptr;
  NDArray* wEffective = nullptr;
  NDArray* zEffective = nullptr;
  NDArray* bEffective = nullptr;

  bool deleteX = false;
  bool deleteW = false;
  bool deleteZ = false;
  bool deleteB = false;

  sd::LongType batchSize = 1;
  bool needsReshape = x->rankOf() > 2;

  if (needsReshape) {
    // Calculate the 2D shape: flatten all but last dimension
    for (int i = 0; i < x->rankOf() - 1; i++) {
      batchSize *= x->sizeAt(i);
    }
    sd::LongType inputLastDim = x->sizeAt(-1);
    sd::LongType outputLastDim = z->sizeAt(-1);

    // Reshape to 2D - these create new NDArray objects
    std::vector<sd::LongType> xShape2D = {batchSize, inputLastDim};
    std::vector<sd::LongType> zShape2D = {batchSize, outputLastDim};
    xEffective = x->reshape('c', xShape2D);
    zEffective = z->reshape('c', zShape2D);
    deleteX = true;
    deleteZ = true;
  } else {
    xEffective = x;
    zEffective = z;
  }

  // Apply transposes if needed
  if (aTranspose) {
    auto xTransposed = new NDArray(xEffective->transpose());
    if (deleteX) delete xEffective;
    xEffective = xTransposed;
    deleteX = true;
  }
  if (cTranspose) {
    auto zTransposed = new NDArray(zEffective->transpose());
    if (deleteZ) delete zEffective;
    zEffective = zTransposed;
    deleteZ = true;
  }

  // Handle weight transpose
  if (bTranspose) {
    wEffective = new NDArray(w->transpose());
    deleteW = true;
  } else {
    wEffective = w;
  }

  REQUIRE_TRUE(xEffective->rankOf() == 2, 0,
               "xw_plus_b: After reshaping, input x array should have rank equal 2, but got instead %i!",
               xEffective->rankOf());
  REQUIRE_TRUE(wEffective->rankOf() == 2, 0,
               "xw_plus_b: Input weights array should have rank equal 2, but got instead %i!",
               wEffective->rankOf());
  REQUIRE_TRUE(zEffective->rankOf() == 2, 0,
               "xw_plus_b: After reshaping, output array should have rank equal 2, but got instead %i!",
               zEffective->rankOf());

  // Perform matrix multiplication: z = x @ w
  MmulHelper::mmul(xEffective, wEffective, zEffective, 1.0, 0.0);

  // Add bias - ALWAYS as a row vector since output is [batch, features]
  // The bias vector has shape [features] and should broadcast across the batch dimension
  // This is the standard behavior for ONNX Gemm: Y = X @ W + B where B broadcasts
  if (b->rankOf() == 1) {
    std::vector<sd::LongType> bShape2D = {1, b->lengthOf()};
    bEffective = b->reshape('c', bShape2D);
    deleteB = true;
  } else {
    bEffective = b;
  }

  if (zEffective->isMatrix()) {
    zEffective->addiRowVector(bEffective);
  } else {
    *zEffective += *bEffective;
  }

  // Cleanup heap-allocated arrays
  if (deleteB) delete bEffective;
  if (deleteW) delete wEffective;
  if (deleteX) delete xEffective;
  if (deleteZ) delete zEffective;

  return Status::OK;
}

DECLARE_SHAPE_FN(xw_plus_b) {
  auto xShape = inputShape->at(0);
  auto weights = INPUT_VARIABLE(1);
  bool aTranspose = (block.getIArguments()->size() > 0 ? INT_ARG(0) == 1 : false);
  bool bTranspose = (block.getIArguments()->size() > 1 ? INT_ARG(1) == 1 : false);
  bool cTranspose = (block.getIArguments()->size() > 2 ? INT_ARG(2) == 1 : false);

  int nWeightsFormat = block.getIArguments()->size() > 0 ? INT_ARG(0) : 0;

  auto weightsShape =
      (1 == nWeightsFormat) ? ShapeUtils::evalTransposeShapeInfo(*weights, block.getWorkspace()) : inputShape->at(1);

  // Handle higher rank inputs
  if (shape::rank(xShape) > 2) {
    // Calculate 2D shapes for matmul
    sd::LongType batchSize = 1;
    for (int i = 0; i < shape::rank(xShape) - 1; i++) {
      batchSize *= shape::sizeAt(xShape, i);
    }
    sd::LongType lastDim = shape::sizeAt(xShape, shape::rank(xShape) - 1);
    
    // Create temporary 2D shape for x
    std::vector<sd::LongType> x2dShape = {batchSize, lastDim};
    auto x2dShapeInfo = ConstantShapeHelper::getInstance().createShapeInfo(ArrayOptions::dataType(xShape), 
                                                                         'c', x2dShape);
    
    // Get the output shape from matmul
    auto matmulOutput = ShapeUtils::matrixProductShape(x2dShapeInfo, const_cast<sd::LongType *>(weightsShape), 
                                                       aTranspose, bTranspose,
                                                       ArrayOptions::dataType(xShape), block.getWorkspace());
    
    // Calculate final output shape
    std::vector<sd::LongType> outputShape;
    for (int i = 0; i < shape::rank(xShape) - 1; i++) {
      outputShape.push_back(shape::sizeAt(xShape, i));
    }
    // Add the output dimension from the weights
    outputShape.push_back(shape::sizeAt(matmulOutput, 1));
    
    auto finalShape = ConstantShapeHelper::getInstance().createShapeInfo(ArrayOptions::dataType(xShape), 
                                                                         'c', outputShape);
    return SHAPELIST(finalShape);
  } else {
    // Original behavior for rank 2 inputs
    auto outputShape = ShapeUtils::matrixProductShape(xShape, const_cast<sd::LongType *>(weightsShape), aTranspose,
                                                      bTranspose,
                                                      ArrayOptions::dataType(xShape), block.getWorkspace());
    return SHAPELIST(outputShape);
  }
}

DECLARE_TYPES(xw_plus_b) {
  getOpDescriptor()->setAllowedInputTypes(ANY)->setAllowedOutputTypes({ALL_FLOATS});
}

CUSTOM_OP_IMPL(xw_plus_b_bp, 4, 3, false, 0, 0) {

  bool aTranspose = (block.getIArguments()->size() > 0 ? INT_ARG(0) == 1 : false);
  bool bTranspose = (block.getIArguments()->size() > 1 ? INT_ARG(1) == 1 : false);
  auto x = aTranspose ? INPUT_VARIABLE(0)->transpose() : INPUT_VARIABLE(0);  // transpose() already returns NDArray*
  auto b = INPUT_VARIABLE(2);
  auto dLdz = INPUT_VARIABLE(3);

  if (x->isEmpty() || INPUT_VARIABLE(1)->isEmpty() || b->isEmpty() || dLdz->isEmpty()) return Status::OK;

  auto w = bTranspose ? INPUT_VARIABLE(1)->transpose() : INPUT_VARIABLE(1);  // transpose() already returns NDArray*

  REQUIRE_TRUE(x->rankOf() == 2, 0, "xw_plus_b BP: Input x array should have rank equal 2, but got instead %i!",
               x->rankOf());
  REQUIRE_TRUE(w->rankOf() == 2, 0, "xw_plus_b BP: Input weights array should have rank equal 2, but got instead %i!",
               w->rankOf());
  REQUIRE_TRUE(dLdz->rankOf() == 2, 0, "xw_plus_b BP: Output array should have rank equal 2, but got instead %i!",
               dLdz->rankOf());

  auto dLdx = aTranspose ? OUTPUT_VARIABLE(0)->transpose() : OUTPUT_VARIABLE(0);  // transpose() already returns NDArray*
  auto dLdb = OUTPUT_VARIABLE(2);

  auto dLdw = (bTranspose) ? OUTPUT_VARIABLE(1)->transpose() : OUTPUT_VARIABLE(1);  // transpose() already returns NDArray*

  // dLdb - reduceAlongDimension returns pointer
  std::vector<LongType> dims({0});
  auto* assign = dLdz->reduceAlongDimension(reduce::Sum, &dims);
  dLdb->assign(assign);
  delete assign;

  matmul_bp mmul_bp;
  mmul_bp.execute({x, w, dLdz}, std::vector<NDArray*>{dLdx, dLdw}, {}, {}, {});

  // Transpose views are managed by parent arrays - no deletion needed
  // x is from INPUT_VARIABLE(0)->transpose() if aTranspose
  // w is from INPUT_VARIABLE(1)->transpose() if bTranspose
  // dLdx is from OUTPUT_VARIABLE(0)->transpose() if aTranspose
  // dLdw is from OUTPUT_VARIABLE(1)->transpose() if bTranspose
  // All are views managed by their parent arrays

  return Status::OK;
}

DECLARE_SHAPE_FN(xw_plus_b_bp) {
  return SHAPELIST(CONSTANT(inputShape->at(0)), CONSTANT(inputShape->at(1)), CONSTANT(inputShape->at(2)));
}

DECLARE_TYPES(xw_plus_b_bp) {
  getOpDescriptor()->setAllowedInputTypes(ANY)->setAllowedOutputTypes({ALL_FLOATS});
}
}  // namespace ops
}  // namespace sd

#endif