#pragma once

#include "mlir/Dialect/Tensor/IR/Tensor.h"
#include "mlir/IR/Block.h"
#include "mlir/IR/BuiltinTypes.h"
#include "mlir/IR/ValueRange.h"
#include "mlir/Transforms/DialectConversion.h"

#include <cassert>
#include <cstddef>
#include <limits>
#include <type_traits>
#include <utility>

#include "src/Accelerators/PIM/Common/Support/CheckedArithmetic.hpp"
#include "src/Accelerators/PIM/Conversion/ONNXToSpatial/CompileTime.hpp"
#include "src/Accelerators/PIM/Dialect/Spatial/SpatialOps.hpp"

namespace onnx_mlir {

namespace detail {

inline mlir::ValueRange getBlockArgs(mlir::Block* block) { return mlir::ValueRange(block->getArguments()); }

inline mlir::ValueRange getInputBlockArgs(mlir::Block* block, size_t weightCount) {
  return mlir::ValueRange(block->getArguments()).drop_front(weightCount);
}

template <typename Fn, size_t... Is>
decltype(auto) invokeWithBlockArgs(Fn&& fn, mlir::Block* block, std::index_sequence<Is...>) {
  return std::forward<Fn>(fn)(block->getArgument(Is)...);
}

template <typename Fn, size_t... Is>
decltype(auto) invokeWithValues(Fn&& fn, mlir::ValueRange values, std::index_sequence<Is...>) {
  return std::forward<Fn>(fn)(values[Is]...);
}

template <size_t>
using ValueArg = mlir::Value;

template <typename Fn, typename Seq>
struct InvokeWithBlockArgsResult;

template <typename Fn, size_t... Is>
struct InvokeWithBlockArgsResult<Fn, std::index_sequence<Is...>> {
  using type = std::invoke_result_t<Fn, ValueArg<Is>...>;
};

template <typename Fn, typename Seq>
using InvokeWithBlockArgsResultT = typename InvokeWithBlockArgsResult<Fn, Seq>::type;

template <typename Fn>
using InvokeWithValueRangeResultT = std::invoke_result_t<Fn, mlir::ValueRange>;

struct SpatComputeBatchBodyArgs {
  mlir::Value lane;
  mlir::ValueRange weights;
  mlir::ValueRange inputs;
  mlir::ValueRange outputs;
};

} // namespace detail

template <typename RewriterT>
inline mlir::Value createSpatConcat(RewriterT& rewriter, mlir::Location loc, int64_t axis, mlir::ValueRange inputs) {
  assert(!inputs.empty() && "spat.concat requires at least one input");
  if (inputs.size() == 1)
    return inputs.front();

  auto firstType = mlir::cast<mlir::RankedTensorType>(inputs.front().getType());
  auto outputShape = llvm::to_vector(firstType.getShape());
  int64_t concatDimSize = 0;
  bool concatDimDynamic = false;

  for (mlir::Value input : inputs) {
    auto inputType = mlir::cast<mlir::RankedTensorType>(input.getType());
    assert(inputType.getRank() == firstType.getRank() && "spat.concat expects same-rank inputs");
    if (mlir::ShapedType::isDynamic(inputType.getDimSize(axis)))
      concatDimDynamic = true;
    else
      concatDimSize += inputType.getDimSize(axis);
  }

  outputShape[axis] = concatDimDynamic ? mlir::ShapedType::kDynamic : concatDimSize;
  auto outputType = mlir::RankedTensorType::get(outputShape, firstType.getElementType(), firstType.getEncoding());
  return spatial::SpatConcatOp::create(rewriter, loc, outputType, rewriter.getI64IntegerAttr(axis), inputs).getOutput();
}

/// Builds a `spat.graph_compute` with a fixed number of SSA inputs and erases it if
/// the body callback reports failure.
template <size_t NumInputs, typename RewriterT, typename BodyFn>
auto createSpatGraphCompute(RewriterT& rewriter,
                            mlir::Location loc,
                            mlir::TypeRange resultTypes,
                            mlir::ValueRange weights,
                            mlir::ValueRange inputs,
                            BodyFn&& body) {
  assert(inputs.size() == NumInputs && "NumInputs must match the number of input values");
  auto computeOp = spatial::SpatGraphCompute::create(rewriter, loc, resultTypes, weights, inputs);

  auto* block = new mlir::Block();
  for (mlir::Value weight : weights)
    block->addArgument(weight.getType(), loc);
  for (mlir::Value input : inputs)
    block->addArgument(input.getType(), loc);

  computeOp.getBody().push_back(block);
  rewriter.setInsertionPointToStart(block);

  using BodyResult = detail::InvokeWithBlockArgsResultT<std::decay_t<BodyFn>, std::make_index_sequence<NumInputs>>;
  if constexpr (std::is_same_v<BodyResult, void>) {
    detail::invokeWithValues(std::forward<BodyFn>(body),
                             detail::getInputBlockArgs(block, weights.size()),
                             std::make_index_sequence<NumInputs> {});

    rewriter.setInsertionPointAfter(computeOp);
    return computeOp;
  }
  else {
    auto bodyResult = detail::invokeWithValues(std::forward<BodyFn>(body),
                                               detail::getInputBlockArgs(block, weights.size()),
                                               std::make_index_sequence<NumInputs> {});
    if (mlir::failed(bodyResult)) {
      rewriter.setInsertionPointAfter(computeOp);
      rewriter.eraseOp(computeOp);
      return mlir::FailureOr<spatial::SpatGraphCompute>(mlir::failure());
    }
    rewriter.setInsertionPointAfter(computeOp);
    return mlir::FailureOr<spatial::SpatGraphCompute>(computeOp);
  }
}

/// Builds a `spat.graph_compute` whose body consumes the block arguments as a single
/// `ValueRange`, which is convenient for variadic reductions/concats.
template <typename RewriterT, typename BodyFn>
auto createSpatGraphCompute(RewriterT& rewriter,
                            mlir::Location loc,
                            mlir::TypeRange resultTypes,
                            mlir::ValueRange weights,
                            mlir::ValueRange inputs,
                            BodyFn&& body) {
  auto computeOp = spatial::SpatGraphCompute::create(rewriter, loc, resultTypes, weights, inputs);

  auto* block = new mlir::Block();
  for (mlir::Value weight : weights)
    block->addArgument(weight.getType(), loc);
  for (mlir::Value input : inputs)
    block->addArgument(input.getType(), loc);

  computeOp.getBody().push_back(block);
  rewriter.setInsertionPointToStart(block);

  using BodyResult = detail::InvokeWithValueRangeResultT<std::decay_t<BodyFn>>;
  if constexpr (std::is_same_v<BodyResult, void>) {
    std::forward<BodyFn>(body)(detail::getInputBlockArgs(block, weights.size()));

    rewriter.setInsertionPointAfter(computeOp);
    return computeOp;
  }
  else {
    auto bodyResult = std::forward<BodyFn>(body)(detail::getInputBlockArgs(block, weights.size()));
    if (mlir::failed(bodyResult)) {
      rewriter.setInsertionPointAfter(computeOp);
      rewriter.eraseOp(computeOp);
      return mlir::FailureOr<spatial::SpatGraphCompute>(mlir::failure());
    }
    rewriter.setInsertionPointAfter(computeOp);
    return mlir::FailureOr<spatial::SpatGraphCompute>(computeOp);
  }
}

template <typename RewriterT, typename BodyFn>
auto createSpatGraphComputeBatch(RewriterT& rewriter,
                                 mlir::Location loc,
                                 mlir::TypeRange resultTypes,
                                 int64_t laneCount,
                                 mlir::ValueRange weights,
                                 mlir::ValueRange inputs,
                                 BodyFn&& body) {
  if (laneCount <= 0 || laneCount > std::numeric_limits<int32_t>::max())
    return mlir::FailureOr<spatial::SpatGraphComputeBatch>(mlir::failure());

  auto laneCountAttr = pim::getCheckedI32Attr(rewriter, loc, laneCount, "spatial compute_batch lane count");
  if (mlir::failed(laneCountAttr))
    return mlir::FailureOr<spatial::SpatGraphComputeBatch>(mlir::failure());

  auto batchOp = spatial::SpatGraphComputeBatch::create(rewriter, loc, resultTypes, *laneCountAttr, weights, inputs);

  mlir::SmallVector<mlir::Type> blockArgTypes {rewriter.getIndexType()};
  mlir::SmallVector<mlir::Location> blockArgLocs {loc};
  blockArgTypes.reserve(1 + weights.size() + inputs.size() + resultTypes.size());
  blockArgLocs.reserve(1 + weights.size() + inputs.size() + resultTypes.size());
  for (mlir::Value weight : weights) {
    blockArgTypes.push_back(weight.getType());
    blockArgLocs.push_back(weight.getLoc());
  }
  for (mlir::Value input : inputs) {
    blockArgTypes.push_back(input.getType());
    blockArgLocs.push_back(input.getLoc());
  }
  for (mlir::Type resultType : resultTypes) {
    blockArgTypes.push_back(resultType);
    blockArgLocs.push_back(loc);
  }

  auto* block =
    rewriter.createBlock(&batchOp.getBody(), batchOp.getBody().end(), mlir::TypeRange(blockArgTypes), blockArgLocs);
  rewriter.setInsertionPointToStart(block);

  detail::SpatComputeBatchBodyArgs args {
    block->getArgument(0),
    mlir::ValueRange(block->getArguments()).slice(1, weights.size()),
    mlir::ValueRange(block->getArguments()).slice(1 + weights.size(), inputs.size()),
    mlir::ValueRange(block->getArguments()).drop_front(1 + weights.size() + inputs.size())};

  using BodyResult = std::invoke_result_t<BodyFn, detail::SpatComputeBatchBodyArgs>;
  if constexpr (std::is_same_v<BodyResult, void>) {
    std::forward<BodyFn>(body)(args);
    rewriter.setInsertionPointAfter(batchOp);
    return mlir::FailureOr<spatial::SpatGraphComputeBatch>(batchOp);
  }
  else {
    auto bodyResult = std::forward<BodyFn>(body)(args);
    if (mlir::failed(bodyResult)) {
      rewriter.setInsertionPointAfter(batchOp);
      rewriter.eraseOp(batchOp);
      return mlir::FailureOr<spatial::SpatGraphComputeBatch>(mlir::failure());
    }
    rewriter.setInsertionPointAfter(batchOp);
    return mlir::FailureOr<spatial::SpatGraphComputeBatch>(batchOp);
  }
}

template <size_t NumInputs, typename RewriterT, typename BodyFn>
auto createSpatCompute(RewriterT& rewriter,
                       mlir::Location loc,
                       mlir::TypeRange resultTypes,
                       mlir::ValueRange weights,
                       mlir::ValueRange inputs,
                       BodyFn&& body) {
  return createSpatGraphCompute<NumInputs>(
    rewriter, loc, resultTypes, weights, inputs, std::forward<BodyFn>(body));
}

template <typename RewriterT, typename BodyFn>
auto createSpatCompute(RewriterT& rewriter,
                       mlir::Location loc,
                       mlir::TypeRange resultTypes,
                       mlir::ValueRange weights,
                       mlir::ValueRange inputs,
                       BodyFn&& body) {
  return createSpatGraphCompute(rewriter, loc, resultTypes, weights, inputs, std::forward<BodyFn>(body));
}

template <typename RewriterT, typename BodyFn>
auto createSpatComputeBatch(RewriterT& rewriter,
                            mlir::Location loc,
                            mlir::TypeRange resultTypes,
                            int64_t laneCount,
                            mlir::ValueRange weights,
                            mlir::ValueRange inputs,
                            BodyFn&& body) {
  return createSpatGraphComputeBatch(
    rewriter, loc, resultTypes, laneCount, weights, inputs, std::forward<BodyFn>(body));
}

inline void createParallelInsertSliceIntoBatchOutput(mlir::PatternRewriter& rewriter,
                                                     mlir::Location loc,
                                                     mlir::Value source,
                                                     mlir::Value dest,
                                                     mlir::ArrayRef<mlir::OpFoldResult> offsets,
                                                     mlir::ArrayRef<mlir::OpFoldResult> sizes,
                                                     mlir::ArrayRef<mlir::OpFoldResult> strides) {
  auto inParallelOp = spatial::SpatInParallelOp::create(rewriter, loc);
  rewriter.setInsertionPointToStart(&inParallelOp.getRegion().front());
  mlir::tensor::ParallelInsertSliceOp::create(rewriter, loc, source, dest, offsets, sizes, strides);
}

template <typename BodyFn>
mlir::Value materializeOrComputeUnary(mlir::Value input,
                                      mlir::RankedTensorType resultType,
                                      mlir::PatternRewriter& rewriter,
                                      mlir::Location loc,
                                      BodyFn&& build) {
  auto&& buildFn = build;
  if (isCompileTimeComputable(input))
    return buildFn(input);

  auto computeOp = createSpatCompute<1>(
    rewriter, loc, mlir::TypeRange {resultType}, {}, mlir::ValueRange {input}, [&](mlir::Value computeInput) {
      mlir::Value result = buildFn(computeInput);
      spatial::SpatYieldOp::create(rewriter, loc, result);
    });
  return computeOp.getResult(0);
}

mlir::Value sumTensors(mlir::ArrayRef<mlir::Value> tensors, mlir::PatternRewriter& rewriter);

} // namespace onnx_mlir