#include "tsl/platform/statusor.h"

namespace mlir {
namespace TFL {

absl::StatusOr<TypedAttr> CreateTypedAttr(ShapedType shaped_type, int value) {
  Type element_type = shaped_type.getElementType();
  if (element_type.isF16()) {
    auto floatType = mlir::FloatType::getF16(element_type.getContext());
    auto floatAttr = mlir::FloatAttr::get(floatType, static_cast<float>(value));

    auto const_type = mlir::dyn_cast<ShapedType>(op.getOperand(0).getType());
    auto result_type = mlir::dyn_cast<ShapedType>(op.getResult().getType());
    if (!const_type || !const_type.getElementType().isSignlessInteger(4) ||
        !result_type || !result_type.getElementType().isSignlessInteger(8)) {
      return failure();
    }

    auto input_op =
        dyn_cast<arith::ConstantOp>(op.getOperand(0).getDefiningOp());
    if (!input_op) {

  int num_elements = type.getNumElements();
  int num_zeros = 0;

  if (mlir::isa<FloatType>(type.getElementType())) {
    for (const auto val : attr.getValues<APFloat>()) {
      if (val.isZero()) {
        num_zeros++;
      }
    }
  } else if (mlir::isa<quant::QuantizedType>(type.getElementType())) {
    for (const auto val : attr.getValues<int8_t>()) {
      if (val == 0) {
        num_zeros++;
      }

      return failure();
    }

    if (!IsI8F32UniformQuantizedPerAxisType(filter_type.getElementType())) {
      LLVM_DEBUG(llvm::dbgs() << "Expected a per-channel uniform quantized "
                                 "(i8->f32) type filter. Got: "
                              << filter_type.getElementType() << "\n");
      return failure();
    }

  // Skip int8 x int8 => int32.
  if (lhs_type.getElementType().isInteger(8) &&
      rhs_type.getElementType().isInteger(8)) {
    return rewriter.notifyMatchFailure(op,
                                       "skip unrolling for int8 BatchMatMulV3");
  }

  auto element_type = lhs_type.getElementType();

  if (element_type != rhs_type.getElementType()) {

      // Ignore types with undefined bit widths.
      if (!tensor_type.getElementType().isIntOrFloat()) continue;
      if (!tensor_type.hasStaticShape()) {
        has_unknown_shape = true;
        return size;
      }
      size += tensor_type.getNumElements() *
              tensor_type.getElementType().getIntOrFloatBitWidth();
    }
    return size;
  };

    if (dot.getLhs().getType().getElementType() !=
            first_dot.getLhs().getType().getElementType() ||
        dot.getRhs().getType().getElementType() !=
            first_dot.getRhs().getType().getElementType() ||
        dot.getType().getElementType() != first_dot.getType().getElementType())
      return rewriter.notifyMatchFailure(

    if (!keys_ty || !keys_ty.hasStaticShape() ||
        !keys_ty.getElementType().isIntOrFloat())
      return rewriter.notifyMatchFailure(
          op,
          "only match for the case where the first operand has a static "
          "int/float shapeType");
    if (!indices_ty || !indices_ty.hasStaticShape() ||
        !indices_ty.getElementType().isInteger(32))
      return rewriter.notifyMatchFailure(
          op,

        RankedTensorType::get(lhs_flattened_shape, lhs_type.getElementType()),
        lhs_transposed.getResult());
  } else {
    auto lhs_flattend_shape_op = BuildDotOperandFlattenedShapeOp(
        lhs, lhs_dot_dimensions_info, builder, /*is_lhs=*/true);
    lhs_flattend = rewriter.create<mhlo::DynamicReshapeOp>(
        loc,
        RankedTensorType::get(lhs_flattened_shape, lhs_type.getElementType()),
        lhs_transposed, lhs_flattend_shape_op);

  if (!tensorType) return false;
  return tensorType.getElementType().isF32();
}

// Returns true iff the given value is a bf16 tensor.
inline bool TFTypeIsBFloat16Tensor(Value value) {
  auto tensorType = mlir::dyn_cast<TensorType>(value.getType());
  if (!tensorType) return false;
  return tensorType.getElementType().isBF16();
}

// Returns true iff the given value is a f16 tensor.