#include "llvm/ADT/SmallVector.h"

#include "src/Accelerators/PIM/Common/IR/AffineUtils.hpp"
#include "src/Accelerators/PIM/Common/IR/ConstantUtils.hpp"
#include "src/Accelerators/PIM/Conversion/ONNXToSpatial/Common/BiasAddUtils.hpp"
#include "src/Accelerators/PIM/Conversion/ONNXToSpatial/Common/Common.hpp"
#include "src/Accelerators/PIM/Conversion/ONNXToSpatial/Common/ComputeRegionBuilder.hpp"
#include "src/Accelerators/PIM/Conversion/ONNXToSpatial/Common/RowStripLayoutUtils.hpp"
#include "src/Accelerators/PIM/Dialect/Spatial/SpatialOps.hpp"
#include "src/Dialect/ONNX/ONNXOps.hpp"

using namespace mlir;

namespace onnx_mlir {

FailureOr<RowStripPhysicalValue> describeRowStripPhysicalValue(Value storage, RankedTensorType logicalType) {
  auto storageType = dyn_cast<RankedTensorType>(storage.getType());
  if (!storageType || !storageType.hasStaticShape() || !logicalType || !logicalType.hasStaticShape()
      || storageType.getRank() != 5 || logicalType.getRank() != 4 || logicalType.getDimSize(0) != 1
      || storageType.getElementType() != logicalType.getElementType()
      || storageType.getDimSize(1) != 1 || storageType.getDimSize(2) != 1
      || storageType.getDimSize(3) != logicalType.getDimSize(3) || storageType.getDimSize(4) <= 0)
    return failure();
  const int64_t tilesPerRow = ceilIntegerDivide(logicalType.getDimSize(1), storageType.getDimSize(4));
  if (storageType.getDimSize(0) != logicalType.getDimSize(2) * tilesPerRow)
    return failure();
  return RowStripPhysicalValue {storage, logicalType,
    RankedTensorType::get(storageType.getShape().drop_front(), storageType.getElementType(), storageType.getEncoding()),
    tilesPerRow};
}

RankedTensorType getRowStripFragmentType(RankedTensorType logicalType) {
  return RankedTensorType::get({logicalType.getDimSize(0), 1, logicalType.getDimSize(3),
                                logicalType.getDimSize(1)},
                               logicalType.getElementType(),
                               logicalType.getEncoding());
}

RankedTensorType getRowStripStorageType(RankedTensorType logicalType) {
  return spatial::getGraphBatchPhysicalResultType(logicalType.getDimSize(2), getRowStripFragmentType(logicalType));
}

std::pair<SmallVector<int64_t>, SmallVector<int64_t>> buildRowStripMetadata(RankedTensorType type) {
  SmallVector<int64_t> offsets;
  SmallVector<int64_t> sizes;
  const int64_t channels = type.getDimSize(1);
  const int64_t height = type.getDimSize(2);
  const int64_t width = type.getDimSize(3);
  offsets.reserve(height * 4);
  sizes.reserve(height * 4);
  for (int64_t row = 0; row < height; ++row) {
    offsets.append({0, 0, row, 0});
    sizes.append({1, channels, 1, width});
  }
  return {offsets, sizes};
}

Value extractRowStripFragment(Value storage,
                              RankedTensorType logicalType,
                              OpFoldResult row,
                              PatternRewriter& rewriter,
                              Location loc) {
  return *extractGraphBatchPhysicalFragment(rewriter, loc, storage, row, getRowStripFragmentType(logicalType));
}

void insertRowStripFragment(Value fragment,
                            Value output,
                            RankedTensorType logicalType,
                            OpFoldResult row,
                            PatternRewriter& rewriter,
                            Location loc) {
  assert(fragment.getType() == getRowStripFragmentType(logicalType));
  assert(output.getType() == getRowStripStorageType(logicalType));
  auto slot = dyn_cast<Value>(row);
  assert(slot && "row-strip graph publication requires a dynamic physical slot");
  publishGraphBatchPhysicalFragment(rewriter, loc, fragment, output, slot);
}

FailureOr<Value> createPerChannelConstantFragment(DenseElementsAttr denseAttr,
                                                  RankedTensorType fragmentType,
                                                  PatternRewriter& rewriter) {
  auto logicalType = RankedTensorType::get(
    {1, fragmentType.getDimSize(3), 1, 1}, fragmentType.getElementType());
  FailureOr<SmallVector<Attribute>> channelValues = getBiasChannelValues(denseAttr, logicalType);
  if (failed(channelValues))
    return failure();

  SmallVector<Attribute> values;
  values.reserve(fragmentType.getNumElements());
  for (int64_t n = 0; n < fragmentType.getDimSize(0); ++n)
    for (int64_t h = 0; h < fragmentType.getDimSize(1); ++h)
      for (int64_t w = 0; w < fragmentType.getDimSize(2); ++w)
        for (int64_t channel = 0; channel < fragmentType.getDimSize(3); ++channel)
          values.push_back((*channelValues)[channel]);

  auto attr = DenseElementsAttr::get(fragmentType, values);
  return getOrCreateConstant(rewriter, rewriter.getInsertionBlock()->getParentOp(), attr, fragmentType);
}

FailureOr<Value> createRowStripStorageFromRows(Value rows,
                                               RankedTensorType logicalType,
                                               PatternRewriter& rewriter,
                                               Location loc) {
  auto rowsType = dyn_cast<RankedTensorType>(rows.getType());
  if (!rowsType || !rowsType.hasStaticShape() || rowsType.getRank() != 2)
    return failure();
  if (!logicalType || !logicalType.hasStaticShape() || logicalType.getRank() != 4)
    return failure();
  if (logicalType.getDimSize(0) != 1)
    return failure();
  if (rowsType.getElementType() != logicalType.getElementType())
    return failure();

  const int64_t channels = logicalType.getDimSize(1);
  const int64_t height = logicalType.getDimSize(2);
  const int64_t width = logicalType.getDimSize(3);
  if (rowsType.getDimSize(0) != height * width)
    return failure();
  if (rowsType.getDimSize(1) != channels)
    return failure();

  auto rowSliceType = RankedTensorType::get({width, channels}, logicalType.getElementType(), rowsType.getEncoding());
  auto fragmentType = getRowStripFragmentType(logicalType);
  auto storageType = getRowStripStorageType(logicalType);
  auto batchOp = createSpatComputeBatch(
    rewriter, loc, TypeRange {storageType}, height, {}, ValueRange {rows}, [&](detail::SpatComputeBatchBodyArgs args) {
      Operation* anchorOp = rewriter.getInsertionBlock()->getParentOp();
      Value rowStart = affineMulConst(rewriter, loc, args.lane, width, anchorOp);
      SmallVector<OpFoldResult> rowOffsets {rowStart, rewriter.getIndexAttr(0)};
      SmallVector<OpFoldResult> rowSizes {rewriter.getIndexAttr(width), rewriter.getIndexAttr(channels)};
      Value rowSlice = tensor::ExtractSliceOp::create(
        rewriter, loc, rowSliceType, args.inputs.front(), rowOffsets, rowSizes, getUnitStrides(rewriter, 2));
      Value fragment = tensor::ExpandShapeOp::create(
        rewriter, loc, fragmentType, rowSlice, SmallVector<ReassociationIndices> {{0, 1, 2}, {3}});
      insertRowStripFragment(fragment, args.outputs.front(), logicalType, args.lane, rewriter, loc);
      return success();
    });
  if (failed(batchOp))
    return failure();
  return batchOp->getResult(0);
}

FailureOr<Value> createRowStripAssemblyBlueprint(const RowStripPhysicalValue& value,
                                                 PatternRewriter& rewriter,
                                                 Location loc) {
  const int64_t laneCount = cast<RankedTensorType>(value.storage.getType()).getDimSize(0);
  const int64_t tileChannels = value.fragmentType.getDimSize(3);
  auto nchwFragmentType = RankedTensorType::get(
    {1, tileChannels, 1, value.logicalType.getDimSize(3)}, value.fragmentType.getElementType(),
    value.fragmentType.getEncoding());
  auto nchwStorageType = spatial::getGraphBatchPhysicalResultType(laneCount, nchwFragmentType);
  auto transposed = createSpatComputeBatch(
    rewriter, loc, TypeRange {nchwStorageType}, laneCount, {}, ValueRange {value.storage},
    [&](detail::SpatComputeBatchBodyArgs args) {
      FailureOr<Value> fragment = extractGraphBatchPhysicalFragment(
        rewriter, loc, args.inputs.front(), args.lane, value.fragmentType);
      if (failed(fragment))
        return failure();
      Value nchw = ONNXTransposeOp::create(
        rewriter, loc, nchwFragmentType, *fragment, rewriter.getI64ArrayAttr({0, 3, 1, 2}));
      publishGraphBatchPhysicalFragment(rewriter, loc, nchw, args.outputs.front(), args.lane);
      return success();
    });
  if (failed(transposed))
    return failure();

  SmallVector<FragmentAssemblyEntry> entries;
  for (int64_t row = 0; row < value.logicalType.getDimSize(2); ++row)
    for (int64_t tile = 0; tile < value.tilesPerRow; ++tile) {
      const int64_t channelOffset = tile * tileChannels;
      entries.push_back({row * value.tilesPerRow + tile, 0, {0, channelOffset, row, 0},
        {1, std::min(tileChannels, value.logicalType.getDimSize(1) - channelOffset), 1,
         value.logicalType.getDimSize(3)}});
    }
  return createFragmentAssemblyBlueprint(transposed->getResult(0), value.logicalType, entries, "nhwc_row_strip",
                                         kRowStripIndexMap, rewriter, loc);
}

FailureOr<Value> applyRowStripRelu(const RowStripPhysicalValue& value, PatternRewriter& rewriter, Location loc) {
  auto storageType = cast<RankedTensorType>(value.storage.getType());
  const int64_t laneCount = storageType.getDimSize(0);
  auto batchOp = createSpatComputeBatch(rewriter,
                                        loc,
                                        TypeRange {storageType},
                                        laneCount,
                                        {},
                                        ValueRange {value.storage},
                                        [&](detail::SpatComputeBatchBodyArgs args) {
                                          FailureOr<Value> fragment = extractGraphBatchPhysicalFragment(
                                            rewriter, loc, args.inputs.front(), args.lane, value.fragmentType);
                                          if (failed(fragment)) return failure();
                                          Value relu = spatial::SpatReluOp::create(
                                            rewriter, loc, value.fragmentType, *fragment).getResult();
                                          publishGraphBatchPhysicalFragment(
                                            rewriter, loc, relu, args.outputs.front(), args.lane);
                                          return success();
                                        });
  if (failed(batchOp))
    return failure();
  return batchOp->getResult(0);
}

FailureOr<Value> applyRowStripBiasAdd(const RowStripPhysicalValue& value,
                                     Value bias,
                                     PatternRewriter& rewriter,
                                     Location loc) {
  DenseElementsAttr denseAttr;
  if (!isSupportedBiasAddValue(bias, value.logicalType, &denseAttr))
    return failure();
  FailureOr<SmallVector<Attribute>> channelValues = getBiasChannelValues(denseAttr, value.logicalType);
  if (failed(channelValues)) return failure();
  auto storageType = cast<RankedTensorType>(value.storage.getType());
  auto biasStorageType = spatial::getGraphBatchPhysicalResultType(value.tilesPerRow, value.fragmentType);
  SmallVector<Attribute> biasValues(
    biasStorageType.getNumElements(), cast<Attribute>(rewriter.getZeroAttr(value.fragmentType.getElementType())));
  const int64_t tileChannels = value.fragmentType.getDimSize(3);
  const int64_t width = value.fragmentType.getDimSize(2);
  for (int64_t channel = 0; channel < value.logicalType.getDimSize(1); ++channel)
    for (int64_t w = 0; w < width; ++w)
      biasValues[(channel / tileChannels) * width * tileChannels + w * tileChannels
                 + channel % tileChannels] =
        (*channelValues)[channel];
  Value biasStorage = getOrCreateConstant(rewriter, rewriter.getInsertionBlock()->getParentOp(),
    DenseElementsAttr::get(biasStorageType, biasValues), biasStorageType);
  const int64_t laneCount = storageType.getDimSize(0);
  auto batchOp = createSpatComputeBatch(rewriter,
                                        loc,
                                        TypeRange {storageType},
                                        laneCount,
                                        {},
                                        ValueRange {value.storage, biasStorage},
                                        [&](detail::SpatComputeBatchBodyArgs args) {
                                          Operation* anchorOp = rewriter.getInsertionBlock()->getParentOp();
                                          FailureOr<Value> fragment = extractGraphBatchPhysicalFragment(
                                            rewriter, loc, args.inputs[0], args.lane, value.fragmentType);
                                          Value tile = affineModConst(rewriter, loc, args.lane,
                                            value.tilesPerRow, anchorOp);
                                          FailureOr<Value> constant = extractGraphBatchPhysicalFragment(
                                            rewriter, loc, args.inputs[1], tile, value.fragmentType);
                                          if (failed(fragment) || failed(constant)) return failure();
                                          Value added = spatial::SpatVAddOp::create(
                                            rewriter, loc, value.fragmentType, *fragment, *constant).getResult();
                                          publishGraphBatchPhysicalFragment(
                                            rewriter, loc, added, args.outputs.front(), args.lane);
                                          return success();
                                        });
  if (failed(batchOp))
    return failure();
  return batchOp->getResult(0);
}

} // namespace onnx_mlir