# https://github.com/tensorflow/tensorflow/blob/master/tensorflow/compiler/mlir/tensorflow/ir/tf_generated_ops.td#L965
def TF_AvgPoolOp : TF_Op<"AvgPool", [Pure]> {
  let summary = "Performs average pooling on the input.";

  let description = [{
Each entry in `output` is the mean of the corresponding size `ksize`
window in `value`.
  }];

  let arguments = (ins
    Arg<TF_FloatTensor, [{4-D with shape `[batch, height, width, channels]`.}]>:$value,

    ConfinedAttr<I64ArrayAttr, [ArrayMinCount<4>]>:$ksize,
    ConfinedAttr<I64ArrayAttr, [ArrayMinCount<4>]>:$strides,
    TF_AnyStrAttrOf<["SAME", "VALID"]>:$padding,
    DefaultValuedOptionalAttr<TF_ConvnetDataFormatAttr, "\"NHWC\"">:$data_format
  );

  let results = (outs
    Res<TF_FloatTensor, [{The average pooled output tensor.}]>:$output
  );

  TF_DerivedOperandTypeAttr T = TF_DerivedOperandTypeAttr<0>;
}

Type convertToSignless(MLIRContext* context, Type type) {
  if (auto ranked_tensor = type.dyn_cast<RankedTensorType>()) {
    if (auto intTy = ranked_tensor.getElementType().dyn_cast<IntegerType>()) {
      if (!intTy.isSignless()) {
        return RankedTensorType::get(
            ranked_tensor.getShape(),
            IntegerType::get(context, intTy.getWidth(),
                             mlir::IntegerType::SignednessSemantics::Signless));
      }
    }
  }
  return type;
}

FunctionType convertToSignlessFuncType(MLIRContext* context, FunctionType funcType) {
  llvm::SmallVector<Type, 4> inputs;
  llvm::SmallVector<Type, 4> results;
  for (auto arg : funcType.getInputs()) { inputs.push_back(convertToSignless(context, arg)); }
  for (auto res : funcType.getResults()) { results.push_back(convertToSignless(context, res)); }
  return FunctionType::get(context, inputs, results);
}

bool isSignLessTensorOrOther(Type type) {
  if (auto ranked_tensor = type.dyn_cast<RankedTensorType>()) {
    if (auto intTy = ranked_tensor.getElementType().dyn_cast<IntegerType>()) {
      if (intTy.isUnsigned()) { return false; }
      if (intTy.isSigned()) { return false; }
    }
  }
  return true;
}
bool allSignless(mlir::TypeRange types) {
  for (auto type : types) {
    if (!isSignLessTensorOrOther(type)) { return false; }
  }
  return true;
}

bool allSignless(FunctionType funcType) {
  for (auto arg : funcType.getInputs()) {
    if (!isSignLessTensorOrOther(arg)) { return false; }
  }
  for (auto res : funcType.getResults()) {
    if (!isSignLessTensorOrOther(res)) { return false; }
  }
  return true;
}

Value CreateTransposeValue(Location& loc, ConversionPatternRewriter& rewriter, Value input,
                           ArrayRef<int32_t> perms) {
  int perms_size = perms.size();
  auto transpose_perms = rewriter.create<tosa::ConstOp>(
      loc, RankedTensorType::get({perms_size}, rewriter.getI32Type()),
      rewriter.getI32TensorAttr(perms));
  const auto shape_type = input.getType().cast<ShapedType>();
  std::vector<int64_t> ranked_type;
  for (const auto& index : perms) ranked_type.push_back(shape_type.getDimSize(index));
  return rewriter.create<tosa::TransposeOp>(
      loc, RankedTensorType::get(ranked_type, shape_type.getElementType()), input, transpose_perms);
};

RankedTensorType CreateTransposeType(ShapedType output, ArrayRef<int32_t> perms) {
  std::vector<int64_t> ranked_type;
  for (auto index : perms) ranked_type.push_back(output.getDimSize(index));
  return RankedTensorType::get(ranked_type, output.getElementType());
};

Value CreateBNOp(Location loc, ConversionPatternRewriter& rewriter, Type output_type, Value x,
                 Value mean, Value variance, Value epsilon, Value gamma, Value beta) {
  // sub_op = sub(input, mean)
  auto sub_op0 = rewriter.create<tosa::SubOp>(loc, output_type, x, mean);
  // add_op0 = add(var, epsilon)
  auto add_op0 = rewriter.create<tosa::AddOp>(loc, variance.getType(), variance, epsilon);
  // rsqrt_op = rsqrt(add_op0)
  auto rsqrt_op = rewriter.create<tosa::RsqrtOp>(loc, variance.getType(), add_op0);
  // op4 = mul(sub_op, rsqrt_op)
  auto mul_op0 = rewriter.create<tosa::MulOp>(loc, output_type, sub_op0, rsqrt_op, 0);
  // op5 = mul(mul_op0, gamma)
  auto mul_op1 = rewriter.create<tosa::MulOp>(loc, output_type, mul_op0, gamma, 0);
  // op6 = add(mul_op1, beta)
  Value batch_norm = rewriter.create<tosa::AddOp>(loc, output_type, mul_op1, beta);
  return batch_norm;
};

struct Conv2DOpLowering final : public OpConversionPattern<Conv2DOp> {
 public:
  using OpConversionPattern<Conv2DOp>::OpConversionPattern;
  LogicalResult matchAndRewrite(Conv2DOp op, OpAdaptor adaptor,
                                ConversionPatternRewriter& rewriter) const override {
    auto get_pair_int64_from_array = [](ArrayAttr arr) -> std::pair<int64_t, int64_t> {
      return {arr.getValue()[0].cast<IntegerAttr>().getSInt(),
              arr.getValue()[1].cast<IntegerAttr>().getSInt()};
    };

    auto stride_pairs = get_pair_int64_from_array(op.getStrides());
    auto pad_pairs = get_pair_int64_from_array(op.getPaddingBeforeAttr());
    auto dilation_pairs = get_pair_int64_from_array(op.getDilationRate());

    const auto pad = rewriter.getDenseI64ArrayAttr(
        {pad_pairs.first, pad_pairs.second, pad_pairs.first, pad_pairs.second});
    const auto stride = rewriter.getDenseI64ArrayAttr({stride_pairs.first, stride_pairs.second});
    const auto dilation =
        rewriter.getDenseI64ArrayAttr({dilation_pairs.first, dilation_pairs.second});

    auto bias = op.getBias();
    auto loc = op.getLoc();
    if (!bias) {
      const auto output_shape = op.getOut().getType().cast<ShapedType>();
      // support nhwc
      const auto output_channels = output_shape.getDimSize(op.IsNCHW() ? 1 : 3);
      const auto bias_elem_type = output_shape.getElementType();
      const auto type = RankedTensorType::get(output_channels, bias_elem_type);
      bias = rewriter.create<tosa::ConstOp>(
          op.getLoc(), type, DenseElementsAttr::get(type, rewriter.getZeroAttr(bias_elem_type)));
    }

    auto in = op.getIn();
    auto weight = op.getWeight();
    auto out_type = op.getOut().getType().cast<ShapedType>();
    if (out_type.getRank() != 4) {
      LOG(FATAL) << "Failed to lowering oneflow op";
      op->dump();
    }
    // support nhwc
    if (op.IsNCHW()) {
      const auto perms = {0, 2, 3, 1};
      const auto reverse_perms = {0, 3, 1, 2};
      in = CreateTransposeValue(loc, rewriter, in, perms);
      weight = CreateTransposeValue(loc, rewriter, weight, perms);
      out_type = CreateTransposeType(out_type, perms);
      auto conv2d =
          rewriter.create<tosa::Conv2DOp>(loc, out_type, in, weight, bias, pad, stride, dilation);

      auto res = CreateTransposeValue(loc, rewriter, conv2d, reverse_perms);
      rewriter.replaceOp(op, {res});
    } else {
      rewriter.replaceOpWithNewOp<tosa::Conv2DOp>(op, out_type, in, weight, bias, pad, stride,
                                                  dilation);
    }
    return success();
  }
};