#ifndef SPATIAL_DIALECT_H
#define SPATIAL_DIALECT_H

include "mlir/IR/OpBase.td"
include "mlir/IR/BuiltinTypes.td"
include "mlir/IR/AttrTypeBase.td"

def SpatialDialect : Dialect {
  let name = "spat";
  let summary = "Dialect designed for deep learning computation in a spatial architecture";
  let cppNamespace = "::onnx_mlir::spatial";
  let useDefaultTypePrinterParser = 1;
}

class SpatOp<string mnemonic, list<Trait> traits = []> :
    Op<SpatialDialect, mnemonic, traits>;

// TODO maybe remove and use AnyRankedTensor directly
def SpatTensor:
  AnyTypeOf<[AnyMemRef, AnyRankedTensor], "", "::mlir::ShapedType">;

class SpatType<string name, string typeMnemonic, list<Trait> traits = []>
    : TypeDef<SpatialDialect, name, traits> {
  let mnemonic = typeMnemonic;
}

def SpatChannelType : SpatType<"SpatChannel", "ch"> {
  let summary = "Virtual channel type";
}

def SpatWeightedCompute: SpatOp<"compute", [SingleBlock, AttrSizedOperandSegments]> {
  let summary = "Compute operation, with constant weights already attached";

  let arguments = (ins
    Variadic<SpatTensor>:$weights,
    Variadic<SpatTensor>:$inputs
  );

  let results = (outs
    Variadic<SpatTensor>:$outputs
  );

  let regions = (region SizedRegion<1>:$body);

  let hasVerifier = 1;

  let assemblyFormat = [{
    `[` $weights `]` `(` $inputs `)` attr-dict `:` `[` type($weights) `]` `(` type($inputs) `)` `->` type($outputs) $body
  }];
}

def SpatYieldOp: SpatOp<"yield", [Terminator]> {
  let arguments = (ins
    Variadic<SpatTensor>:$outputs
  );

  let assemblyFormat = [{
    $outputs attr-dict `:` type($outputs)
  }];
}

//===----------------------------------------------------------------------===//
// Data movement operations
//===----------------------------------------------------------------------===//

def SpatChannelNewOp: SpatOp<"channel_new", []> {
  let results = (outs
    SpatChannelType:$new_channel
  );

  let builders = [
    OpBuilder<(ins ), [{
      $_state.addTypes(SpatChannelType());
    }]>
  ];

  let assemblyFormat = [{
    attr-dict
  }];
}

def SpatChannelSendOp: SpatOp<"channel_send", []> {
  let arguments = (ins
    SpatChannelType: $channel,
    SpatTensor: $data
  );

  let assemblyFormat = [{
    $data `to` $channel attr-dict `:` `(` type($data) `->` type($channel) `)`
  }];
}

def SpatChannelReceiveOp: SpatOp<"channel_receive", []> {
  let arguments = (ins
    SpatChannelType: $channel
  );

  let results = (outs
    SpatTensor: $data
  );

  let assemblyFormat = [{
    $channel attr-dict `:` `(` type($channel) `->` type($data) `)`
  }];
}

def SpatChannelBroadcastSendOp : SpatOp<"channel_broadcast_send", []> {
  let arguments = (ins
    SpatChannelType: $channel,
    SpatTensor: $data
  );
}

def SpatChannelBroadcastReceiveOp : SpatOp<"channel_broadcast_receive", []> {
  let arguments = (ins
    SpatChannelType: $channel
  );

  let results = (outs
    SpatTensor: $data
  );
}

//===----------------------------------------------------------------------===//
// Math operations
//===----------------------------------------------------------------------===//

def SpatConstantOp: SpatOp<"constant", []> {
  let description = [{
    "Constant value, should be used for weights and biases"
  }];

  let arguments = (ins
    AnyAttr: $value,
    BoolAttr: $shouldAllocate
  );

  let results = (outs
    SpatTensor: $out
  );
}

def SpatWeightedVMMOp: SpatOp<"Wvmm", []> {
  let summary = "Vector-matrix-Multiplication within a WeightedCompute operation. The matrix is found in the weights of the WeightedCompute operation, indexed by the weightIndex attribute.";

  let arguments = (ins
    I32Attr: $weightIndex,
    SpatTensor:$vector
  );

  let results = (outs
    SpatTensor:$output
  );

  // TODO: Verifier that checks it is within a WeightedCompute operation,
  // that the weightIndex is valid, and that the matrix is of the right size.
  let hasVerifier = 1;
}

def SpatWeightedMVMOp: SpatOp<"Wmvm", []> {
  let summary = "Matrix-vector multiplication within a WeightedCompute operation. The matrix is found in the weights of the WeightedCompute operation, indexed by the weightIndex attribute.";

  let arguments = (ins
    I32Attr: $weightIndex,
    SpatTensor:$vector
  );

  let results = (outs
    SpatTensor:$output
  );

  // TODO: Verifier that checks it is within a WeightedCompute operation,
  // that the weightIndex is valid, and that the matrix is of the right size.
  let hasVerifier = 1;
}


def SpatVAddOp: SpatOp<"vadd", []> {
  let summary = "Element-wise add between tensors a and b. Tensor b must have the same size of tensor b or be a 1x1";

  let arguments = (ins
    SpatTensor: $a,
    SpatTensor: $b
  );

  let results = (outs
    SpatTensor:$output
  );

  let hasVerifier = 1;

  let assemblyFormat = [{
    $a `,` $b attr-dict `:` `(` type($a) `,` type($b) `)` `->` type($output)
  }];
}

def SpatVMulOp: SpatOp<"vmul", []> {
  let summary = "Element-wise multiplication between tensors a and b. Tensor b must have the same size of tensor b or be a 1x1";

    let arguments = (ins
      SpatTensor: $a,
      SpatTensor: $b
    );

    let results = (outs
      SpatTensor:$output
    );

    //let hasVerifier = 1;

    let assemblyFormat = [{
      $a `,` $b attr-dict `:` `(` type($a) `,` type($b) `)` `->` type($output)
    }];
}

def SpatVDivOp: SpatOp<"vdiv", []> {
  let summary = "Element-wise division between tensors a and b. Tensor b must have the same size of tensor b or be a 1x1";

  let arguments = (ins
    SpatTensor:$a,
    SpatTensor:$b
  );

  let results = (outs
    SpatTensor:$output
  );

  //let hasVerifier = 1;

  let assemblyFormat = [{
    $a `,` $b attr-dict `:` `(` type($a) `,` type($b) `)` `->` type($output)
  }];
}

//TODO: remove
def SpatVSDivOp: SpatOp<"vsdiv", []> {

  let summary = "Element-wise division between each element of a vector, and a scalar (wrapped in a tensor for convenience)";

  let arguments = (ins
    SpatTensor:$dividend,
    SpatTensor:$divisor
  );

  let results = (outs
    SpatTensor:$output
  );
}

def SpatSumOp: SpatOp<"sum", []> {
  let summary = "Sum all the elements in the input tensors into a single scalar wrapped in tensor for convenience";

  let arguments = (ins
    SpatTensor: $input
  );

  let results = (outs
    SpatTensor:$output
  );
}

def SpatSigmoidOp: SpatOp<"sigmoid", []> {
  let arguments = (ins
    SpatTensor:$input
  );

  let results = (outs
    SpatTensor:$output
  );
}

def SpatReluOp: SpatOp<"relu", []> {
  let arguments = (ins
    SpatTensor:$input
  );

  let results = (outs
    SpatTensor:$output
  );
}

def SpatVMaxOp: SpatOp<"vmax", []> {

  let summary = "Element-wise max function";

  let arguments = (ins
    SpatTensor: $a,
    SpatTensor: $b
  );

  let results = (outs
    SpatTensor:$output
  );

  let hasVerifier = 1;
}

def SpatApplyFiltersOp : SpatOp<"apply_filters", []> {
  let summary = "Apply multiple crossbar weights to a convolutional input tile.";
  let description = [{
    Applies a variable number of crossbar weights to a single large image tensor tile,
    producing a corresponding output tile. This essentially encapsulates a big for loop
    over all pixels in the input tile, where each pixel is multiplied by all the weights
    in the operation.
  }];

  let arguments = (ins
    I64ArrayAttr: $weightIndices,
    I64ArrayAttr: $xKernelPositions,
    I64ArrayAttr: $yKernelPositions,
    SpatTensor: $input
  );
  let results = (outs SpatTensor);

  let assemblyFormat = [{
    $input attr-dict `:` type($input) `->` type(results)
  }];
}

//===----------------------------------------------------------------------===//
// Other operations
//===----------------------------------------------------------------------===//

def SpatImgConcatOp: SpatOp<"img_concat", []> {

  let summary = "Concatenate pixel tiles into a single image";

  let description = [{
    Concatenate pixel tiles into a single image:
    1. First, concatenate the pixel tiles along the "channel" axis (axis 1).
    2. Next, concatenate the pixel tiles along the "width" axis (axis 2).
    3. Finally, concatenate the pixel tiles along the "height" axis (axis 3).

    The input tiles should be provided in a specific order:
      start from the top left pixel,
      then continue with the pixel on its right,
      and once you finish the first row of pixels, go to the next row.
  }];

  let arguments = (ins
    Variadic<SpatTensor>:$inputs
  );

  let results = (outs
    SpatTensor:$output
  );

  let hasVerifier = 1;

  let extraClassDeclaration = [{
      mlir::Value getInputTile(size_t x, size_t y, size_t tile);
  }];
}

#endif // SPATIAL_DIALECT_H