/* ******************************************************************************
 *
 *
 * 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 Yurii Shyrma (iuriish@yahoo.com), created on 07.03.2019
//
#include <execution/Threads.h>
#include <helpers/ConstantTadHelper.h>
#include <helpers/ShapeUtils.h>
#include <ops/declarable/helpers/gather.h>
#include <legacy/NativeOpExecutioner.h>

#include <numeric>
#if NOT_EXCLUDED(OP_gather)
namespace sd {
namespace ops {
namespace helpers {

////////////////////////////////////////////////////////////////////////
void gather(sd::LaunchContext* context, NDArray* input, NDArray* indices, NDArray* output,
            const std::vector<LongType>& intArgs) {
  sd::LongType axis = intArgs.size() > 0 ? intArgs[0] : 0;
  const sd::LongType inputRank = input->rankOf();
  if (axis < 0) axis += inputRank;

  const sd::LongType numOfIntArgs = intArgs.size();

  // Special handling for 1D input with axis=0
  // This handles cases like gathering from shape arrays where we want flat indexing
  bool is1DFlatGather = (inputRank == 1 && axis == 0);

  if (indices != nullptr) {
    // Validate indices
    for (sd::LongType i = 0; i < indices->lengthOf(); ++i) {
      auto idx = indices->e<sd::LongType>(i);
      
      if (is1DFlatGather) {
        // For 1D arrays with axis=0, treat as flat array access
        if (idx >= input->lengthOf() || idx < 0) {
          std::string error = "helpers::gather function: invalid flat index ";
          error += std::to_string(idx);
          error += " at position ";
          error += std::to_string(i);
          error += ". Input is 1D with length ";
          error += std::to_string(input->lengthOf());
          error += ", valid range is [0, ";
          error += std::to_string(input->lengthOf() - 1);
          error += "]";
          THROW_EXCEPTION(error.c_str());
        }
      } else {
        // Standard axis-based validation
        if (idx >= input->sizeAt(axis) || idx < 0) {
          std::string error = "helpers::gather function: invalid index ";
          error += std::to_string(idx);
          error += " at position ";
          error += std::to_string(i);
          error += ". Input shape ";
          error += ShapeUtils::shapeAsString(input->shapeInfo());
          error += ", axis ";
          error += std::to_string(axis);
          error += ", valid range is [0, ";
          error += std::to_string(input->sizeAt(axis) - 1);
          error += "]";
          THROW_EXCEPTION(error.c_str());
        }
      }
    }

    if (is1DFlatGather) {
      // Special case: 1D input with axis=0 - treat as flat array gather
      // This handles gathering from shape arrays like [1, 512] -> gather index 1 -> get 512
      auto func = PRAGMA_THREADS_FOR {
        for (auto i = start; i < stop; i++) {
          auto idx = indices->e<sd::LongType>(i);
          auto value = input->e<double>(idx);  // Get value at flat index
          output->p(i, value);  // Put in output at position i
        }
      };
      samediff::Threads::parallel_for(func, 0, indices->lengthOf());
      
    } else {
      // Standard gather implementation
      //
      // For gather with axis=A on input shape [..., dimA, ...] and indices shape [I1, I2, ...]:
      // - Output shape is: input[0:A] + indices_shape + input[A+1:]
      // - Input TADs: iterate along axis A, each TAD has shape input[A+1:]
      // - Output TADs: iterate along indices dimensions, each TAD has same shape as input TAD
      //
      // tadForDimensions takes dimensions to KEEP in each TAD (not to exclude)
      // It then internally calls evalDimsToExclude to find which dims to iterate over

      std::vector<sd::LongType> axesVec = {axis};
      auto dimensions = ShapeUtils::evalDimsToExclude(input->rankOf(), 1, axesVec.data());

      // For output TADs, we want the same shape as input TADs
      // Input TAD shape = all dims except axis
      // Output shape = input[0:axis] + indices_shape + input[axis+1:]
      // Output TAD dims should be: dims 0 to axis-1, then dims axis+indicesRank to end
      // This gives TAD shape matching input's TAD shape
      std::vector<sd::LongType> outputTadDims;
      sd::LongType indicesRank = indices->rankOf();

      // Add dimensions before the indices dimensions (0 to axis-1)
      for (sd::LongType d = 0; d < axis; d++) {
        outputTadDims.push_back(d);
      }
      // Add dimensions after the indices dimensions (axis+indicesRank to outputRank-1)
      for (sd::LongType d = axis + indicesRank; d < output->rankOf(); d++) {
        outputTadDims.push_back(d);
      }

      // If outputTadDims is empty, it means each TAD is a scalar - handle this case
      // by using the same approach as input (which would also have empty TAD dims)

      // Get TAD packs - these are cached and should not be deleted
      auto tadPack = sd::ConstantTadHelper::getInstance().tadForDimensions(input->shapeInfo(), dimensions);
      auto tadPackOut = sd::ConstantTadHelper::getInstance().tadForDimensions(output->shapeInfo(), &outputTadDims);

      // Validate TAD packs before use
      if (tadPack == nullptr || tadPackOut == nullptr) {
        if (dimensions) delete dimensions;
        THROW_EXCEPTION("gather: Failed to create TAD packs");
      }

      // Now safe to delete dimensions as TAD helper has made internal copy
      delete dimensions;

      auto tadShapeInfo = tadPack->primaryShapeInfo();
      auto tadOffsets = tadPack->primaryOffsets();
      auto tadShapeInfoOut = tadPackOut->primaryShapeInfo();
      auto tadOffsetsOut = tadPackOut->primaryOffsets();

      // Validate that input and output TAD shapes match
      auto inputTadLength = shape::length(tadShapeInfo);
      auto outputTadLength = shape::length(tadShapeInfoOut);
      if (inputTadLength != outputTadLength) {
        std::string error = "gather: TAD shape mismatch - input TAD length ";
        error += std::to_string(inputTadLength);
        error += " != output TAD length ";
        error += std::to_string(outputTadLength);
        error += ". Input shape: ";
        error += ShapeUtils::shapeAsString(input->shapeInfo());
        error += ", Output shape: ";
        error += ShapeUtils::shapeAsString(output->shapeInfo());
        error += ", Indices shape: ";
        error += ShapeUtils::shapeAsString(indices->shapeInfo());
        error += ", axis: ";
        error += std::to_string(axis);
        THROW_EXCEPTION(error.c_str());
      }

      auto tadShapeInfoCast = const_cast<sd::LongType *>(tadShapeInfo);
      auto tadShapeInfoOutCast = const_cast<sd::LongType *>(tadShapeInfoOut);

      // Calculate the number of gather operations (equal to indices length)
      const sd::LongType numGatherOps = indices->lengthOf();

      // Validate bounds before parallel execution
      if (numGatherOps > tadPackOut->numberOfTads()) {
        std::string error = "gather: indices length ";
        error += std::to_string(numGatherOps);
        error += " exceeds output TAD count ";
        error += std::to_string(tadPackOut->numberOfTads());
        THROW_EXCEPTION(error.c_str());
      }

      auto func = PRAGMA_THREADS_FOR {
        for (auto i = start; i < stop; i++) {
          auto idx = indices->e<sd::LongType>(i);

          // Bounds check for input TAD access
          if (idx >= tadPack->numberOfTads() || idx < 0) {
            continue;
          }

          // Bounds check for output TAD access
          if (i >= tadPackOut->numberOfTads()) {
            continue;
          }

          auto offsetIn = tadOffsets[idx];
          auto offsetOut = tadOffsetsOut[i];

          NativeOpExecutioner::execTransformAny(input->getContext(),
                                                transform::Assign,
                                                input->bufferWithOffset(offsetIn), tadShapeInfoCast,
                                                nullptr, nullptr,
                                                output->bufferWithOffset(offsetOut), tadShapeInfoOutCast,
                                                nullptr, nullptr,
                                                nullptr, false);
        }
      };
      samediff::Threads::parallel_tad(func, 0, numGatherOps);
    }
      
  } else {
    // Integer arguments case
    for (int i = 1; i < numOfIntArgs; ++i) {
      if (is1DFlatGather) {
        // For 1D arrays with axis=0, validate against total length
        if (intArgs[i] >= input->lengthOf() || intArgs[i] < 0) {
          std::string error = "helpers::gather function: invalid flat index ";
          error += std::to_string(intArgs[i]);
          error += " at position ";
          error += std::to_string(i-1);
          error += ". Input is 1D with length ";
          error += std::to_string(input->lengthOf());
          error += ", valid range is [0, ";
          error += std::to_string(input->lengthOf() - 1);
          error += "]";
          THROW_EXCEPTION(error.c_str());
        }
      } else {
        // Standard validation
        if (intArgs[i] >= input->sizeAt(axis) || intArgs[i] < 0) {
          std::string error = "helpers::gather function: invalid index ";
          error += std::to_string(intArgs[i]);
          error += " at position ";
          error += std::to_string(i-1);
          error += ". Input shape ";
          error += ShapeUtils::shapeAsString(input->shapeInfo());
          error += ", axis ";
          error += std::to_string(axis);
          error += ", valid range is [0, ";
          error += std::to_string(input->sizeAt(axis) - 1);
          error += "]";
          THROW_EXCEPTION(error.c_str());
        }
      }
    }

    if (numOfIntArgs == 2) {
      if (is1DFlatGather) {
        // For 1D flat gather with single index
        auto value = input->e<double>(intArgs[1]);
        output->assign(value);
      } else {
        // Standard single index gather
        NDArray *copy = (*input)(intArgs[1], {axis});
        output->assign(copy);
        delete copy;
      }
    } else {
      if (is1DFlatGather) {
        // Multiple indices for 1D flat gather
        for (int i = 1; i < numOfIntArgs; ++i) {
          auto idx = intArgs[i];
          auto value = input->e<double>(idx);
          output->p(i - 1, value);
        }
      } else {
        // Standard multiple indices gather
        // Use the same dimension calculation for input and output TADs
        std::vector<sd::LongType> axesVec = {axis};
        auto dimensions = ShapeUtils::evalDimsToExclude(input->rankOf(), 1, axesVec.data());

        // Get TAD packs - these are cached and should not be deleted
        auto tadPack = sd::ConstantTadHelper::getInstance().tadForDimensions(input->shapeInfo(), dimensions);
        auto tadPackOut = sd::ConstantTadHelper::getInstance().tadForDimensions(output->shapeInfo(), dimensions);

        // Validate TAD packs before use
        if (tadPack == nullptr || tadPackOut == nullptr) {
          if (dimensions) delete dimensions;
          THROW_EXCEPTION("gather: Failed to create TAD packs");
        }

        // Now safe to delete dimensions as TAD helper has made internal copy
        delete dimensions;

        auto tadShapeInfo = tadPack->primaryShapeInfo();
        auto tadOffsets = tadPack->primaryOffsets();
        auto tadShapeInfoOut = tadPackOut->primaryShapeInfo();
        auto tadOffsetsOut = tadPackOut->primaryOffsets();

        // Validate that input and output TAD shapes match
        auto inputTadLength = shape::length(tadShapeInfo);
        auto outputTadLength = shape::length(tadShapeInfoOut);
        if (inputTadLength != outputTadLength) {
          std::string error = "gather: TAD shape mismatch - input TAD length ";
          error += std::to_string(inputTadLength);
          error += " != output TAD length ";
          error += std::to_string(outputTadLength);
          error += ". Input shape: ";
          error += ShapeUtils::shapeAsString(input->shapeInfo());
          error += ", Output shape: ";
          error += ShapeUtils::shapeAsString(output->shapeInfo());
          error += ", axis: ";
          error += std::to_string(axis);
          THROW_EXCEPTION(error.c_str());
        }

        // Number of gather operations (number of indices provided as int args)
        const sd::LongType numGatherOps = numOfIntArgs - 1;

        // Validate bounds before parallel execution
        if (numGatherOps > tadPackOut->numberOfTads()) {
          std::string error = "gather: number of indices ";
          error += std::to_string(numGatherOps);
          error += " exceeds output TAD count ";
          error += std::to_string(tadPackOut->numberOfTads());
          THROW_EXCEPTION(error.c_str());
        }

        auto func = PRAGMA_THREADS_FOR {
          for (auto i = start; i < stop; i++) {
            auto idx = intArgs[i + 1];

            // Bounds check for input TAD access
            if (idx >= tadPack->numberOfTads() || idx < 0) {
              continue;
            }

            // Bounds check for output TAD access
            if (i >= tadPackOut->numberOfTads()) {
              continue;
            }

            auto offsetIn = tadOffsets[idx];
            auto offsetOut = tadOffsetsOut[i];

            NativeOpExecutioner::execTransformAny(input->getContext(),
                                                  transform::Assign,
                                                  input->bufferWithOffset(offsetIn), const_cast<sd::LongType*>(tadShapeInfo),
                                                  nullptr, nullptr,
                                                  output->bufferWithOffset(offsetOut), const_cast<sd::LongType*>(tadShapeInfoOut),
                                                  nullptr, nullptr,
                                                  nullptr, false);
          }
        };
        samediff::Threads::parallel_tad(func, 0, numGatherOps);
      }
    }
  }
}

}  // namespace helpers
}  // namespace ops
}  // namespace sd
#endif