#include "mlir/IR/BuiltinAttributes.h"
#include "mlir/IR/BuiltinTypes.h"

#include "llvm/ADT/SmallVector.h"

#include "src/Accelerators/PIM/Common/IR/ConstantUtils.hpp"
#include "src/Accelerators/PIM/Common/IR/ShapeUtils.hpp"
#include "src/Accelerators/PIM/Conversion/ONNXToSpatial/CompileTime.hpp"
#include "src/Accelerators/PIM/Conversion/ONNXToSpatial/Common/BiasAddUtils.hpp"

using namespace mlir;

namespace onnx_mlir {

LogicalResult isSupportedBiasAddShape(RankedTensorType biasType, RankedTensorType resultType) {
  if (!biasType || !resultType || !biasType.hasStaticShape() || !resultType.hasStaticShape())
    return failure();
  if (resultType.getRank() != 4)
    return failure();
  if (biasType.getElementType() != resultType.getElementType())
    return failure();

  const int64_t channels = resultType.getDimSize(1);
  ArrayRef<int64_t> shape = biasType.getShape();
  if (shape.empty())
    return success();
  if (shape.size() == 1)
    return success(shape[0] == channels);
  if (shape.size() == 2)
    return success(shape[0] == 1 && shape[1] == channels);
  if (shape.size() == 4)
    return success(shape[0] == 1 && shape[1] == channels && shape[2] == 1 && shape[3] == 1);
  return failure();
}

FailureOr<SmallVector<Attribute>> getBiasChannelValues(DenseElementsAttr denseAttr, RankedTensorType resultType) {
  auto biasType = dyn_cast<RankedTensorType>(denseAttr.getType());
  if (!biasType || failed(isSupportedBiasAddShape(biasType, resultType)))
    return failure();

  const int64_t channels = resultType.getDimSize(1);
  if (denseAttr.isSplat()) {
    return SmallVector<Attribute>(channels, denseAttr.getSplatValue<Attribute>());
  }

  SmallVector<Attribute> flattened(denseAttr.getValues<Attribute>());
  if (biasType.getRank() == 1)
    return flattened;
  if (biasType.getRank() == 2)
    return flattened;

  SmallVector<Attribute> channelValues;
  channelValues.reserve(channels);
  const int64_t channelStride = biasType.getDimSize(2) * biasType.getDimSize(3);
  for (int64_t channel = 0; channel < channels; ++channel)
    channelValues.push_back(flattened[channel * channelStride]);
  return channelValues;
}

bool isSupportedBiasAddValue(Value bias, RankedTensorType resultType, DenseElementsAttr* denseAttr) {
  auto attr = getHostConstDenseElementsAttr(bias);
  if (!attr)
    return false;
  auto biasType = dyn_cast<RankedTensorType>(attr.getType());
  if (!biasType || failed(isSupportedBiasAddShape(biasType, resultType)))
    return false;
  if (failed(getBiasChannelValues(attr, resultType)))
    return false;
  if (denseAttr)
    *denseAttr = attr;
  return true;
}

FailureOr<BiasAddPlanCandidate> classifyBiasAddPlanCandidate(Value lhs, Value rhs, RankedTensorType resultType) {
  auto lhsType = dyn_cast<RankedTensorType>(lhs.getType());
  auto rhsType = dyn_cast<RankedTensorType>(rhs.getType());
  if (!lhsType || !rhsType)
    return failure();
  if (lhsType == resultType && isSupportedBiasAddValue(rhs, resultType))
    return BiasAddPlanCandidate {lhs, rhs};
  if (rhsType == resultType && isSupportedBiasAddValue(lhs, resultType))
    return BiasAddPlanCandidate {rhs, lhs};
  return failure();
}

FailureOr<Value>
materializeDenseBiasAddTensor(Value bias, RankedTensorType resultType, RewriterBase& rewriter, Location loc) {
  DenseElementsAttr denseAttr;
  if (!isSupportedBiasAddValue(bias, resultType, &denseAttr))
    return failure();

  FailureOr<SmallVector<Attribute>> channelValues = getBiasChannelValues(denseAttr, resultType);
  if (failed(channelValues))
    return failure();

  SmallVector<Attribute> resultValues;
  resultValues.reserve(resultType.getNumElements());
  const int64_t batches = resultType.getDimSize(0);
  const int64_t channels = resultType.getDimSize(1);
  const int64_t height = resultType.getDimSize(2);
  const int64_t width = resultType.getDimSize(3);
  for (int64_t n = 0; n < batches; ++n)
    for (int64_t c = 0; c < channels; ++c)
      for (int64_t h = 0; h < height; ++h)
        for (int64_t w = 0; w < width; ++w)
          resultValues.push_back((*channelValues)[c]);

  auto resultAttr = DenseElementsAttr::get(resultType, resultValues);
  return getOrCreateConstant(rewriter, rewriter.getInsertionBlock()->getParentOp(), resultAttr, resultType);
}

} // namespace onnx_mlir