import org.gradle.integtests.fixtures.SourceFile

import static org.gradle.nativeplatform.fixtures.app.SourceFileElement.ofFile

class CppSum extends CppSourceFileElement implements SumElement {
    final SourceFileElement header
    final SourceFileElement source

    CppSum(String publicHeadersDir = "headers") {
        header = ofFile(new SourceFile(publicHeadersDir, "sum.h", """
#ifdef _WIN32

// RUN: tf-opt -split-input-file -verify-diagnostics -tf-einsum %s | FileCheck %s

func.func @unary_einsum_reduce_sum_transpose(%arg0: tensor<3x4x5x6xf32>) -> tensor<3x5x4xf32> {
  %0 = "tf.Einsum"(%arg0) {T = "tfdtype$DT_FLOAT", equation = "...gse->...sg"}: (tensor<3x4x5x6xf32>) -> tensor<3x5x4xf32>
  func.return %0 : tensor<3x5x4xf32>
  // CHECK-LABEL: unary_einsum_reduce_sum_transpose
  // CHECK-DAG: %[[cst:.*]] = arith.constant dense<3> : tensor<1xi32>

}

struct EinsumDimensionNumbers {
  // Each field contains the list of dimensions appearing only in the specifed
  // arguments of the einsum op with natural ordering. For example `rhs_out`
  // contains the dimensions appearing in the RHS and the OUTPUT of the einsum
  // but not in the LHS.
  std::vector<int64_t> lhs;
  std::vector<int64_t> rhs;
  std::vector<std::tuple<int64_t, int64_t>> lhs_rhs;

// RUN: tf-opt %s -tf-batch-matmul-to-tf-einsum | FileCheck %s

func.func @test_batch_matmul_to_einsum(%arg0: tensor<1x2x3xf32>, %arg1: tensor<3x4xf32>) -> tensor<1x2x4xf32> {
  // CHECK-LABEL: test_batch_matmul_to_einsum
  // CHECK: "tf.Einsum"(%arg0, %arg1) <{equation = "...mk,...kn->...mn"}> : (tensor<1x2x3xf32>, tensor<3x4xf32>) -> tensor<1x2x4xf32>

### Format

The format for ADR should follow this template:

```markdown
# ADR-000X - Title

## Date

20YY-MM-DD

## Context

Lorem ipsum dolor sit amet, consectetur adipiscing elit, sed do eiusmod tempor incididunt ut labore et dolore magna
aliqua. Ut enim ad minim veniam, quis nostrud exercitation ullamco laboris nisi ut aliquip ex ea commodo consequat.

  func.return %0 : tensor<?x2x4x5xf32>
}

// CHECK: func @xla_dot_v2
// CHECK: %[[einsum:.*]] = "tf.Einsum"(%arg0, %arg1) <{equation = "abc,cde->abde"}> : (tensor<?x2x3xf32>, tensor<3x4x5xf32>) -> tensor<?x2x4x5xf32>
// CHECK: return %[[einsum]] : tensor<?x2x4x5xf32>

// -----

// dimension_numbers: {
//   offset_dims: 0
//   collapsed_slice_dims: 1
//   start_index_map: 1

                    "\""# func_name #"\", $0...)", returns>;

// Add the second argument to the first argument, which is expected to be an
// argument list.
// bias(einsum(inputs), bias) --> einsum_with_bias(AppendToVector(inputs, bias))
// Since inputs is a vector in case of einsum, we cannot use ArgumentList here.
def AppendToVector : NativeCodeCall<"AppendToVector($0, $1)">;

// The list of arguments of the composite function.

  [], (addBenefit 10)>;

// Converts inlined Einsum pattern to TF XlaDotV2 op.
def ConvertTFEinsumToXLADotV2Op : Pat<
  (TF_EinsumOp:$einsum
    $args, $equation),
  (CreateXlaDotV2OpFromTfEinsumOp
    $equation, $args, $einsum),
  [(IsInt32ElementType $einsum),
   // Constraint to check:
   // 1. The einsum has two inputs and one output.
   // 2. The einsum is not created by the convert function itself.

def IsPrecisionEmpty :
  Constraint<CPred<"IsPrecisionEmpty($0)">>;

// Creates Einsum Op from XlaDotV2 Op by generating equation.
def CreateEinsumOpFromXlaDotV2Op : NativeCodeCall<
  "CreateEinsumOpFromXlaDotV2Op($_builder, $_loc, $0...)">;

// Convert XlaDotV2 Op to Einsum Op with above two functions.
def ConvertXlaDotV2OpToEinsumOp : Pat<

/**
 * A Cpp app that remove all sources.
 */
class IncrementalCppStaleLinkOutputApp extends IncrementalCppElement implements AppElement {
    private final greeter = new CppGreeter()
    private final sum = new CppSum()
    private final main = new CppMain(greeter, sum)

    final List<IncrementalElement.Transform> incrementalChanges = [
        delete(greeter), delete(sum), delete(main)]