//===- SpaceToDepth.h -----------------------------------------------------===//
//
//                             The ONNC Project
//
// See LICENSE.TXT for details.
//
//===----------------------------------------------------------------------===//
#ifndef ONNC_IR_COMPUTE_OPERATOR_SPACETODEPTH_H
#define ONNC_IR_COMPUTE_OPERATOR_SPACETODEPTH_H
#include <onnc/IR/ComputeOperator.h>
#include <onnc/IR/ComputeVisitor.h>
#include <onnc/IR/Compute/Attributes.h>
#include <onnc/Support/IOStream.h>

namespace onnc {

class SpaceToDepth : public ComputeOperator
{
public:
  enum IOConst {
    kInput = 0,
    kOutput = 0
  };

  static char ID;

public:
  SpaceToDepth(const IntAttr& pBlocksize);

  // clang-format off
  
  // clang-format on

  // shallow copy constructor.
  SpaceToDepth(const SpaceToDepth &pCopy);

  virtual ~SpaceToDepth() { }

  // clang-format off
  // Attributes getters
  const IntAttr& getBlocksize() const { return m_Blocksize; }


  // Attributes setters
  void setBlocksize(const IntAttr& pBlocksize) { m_Blocksize = pBlocksize; }

  // clang-format on

  Tensor* getInput(unsigned int pIdx) override { return static_cast<Tensor*>(m_Inputs[pIdx]); }

  const Tensor* getInput(unsigned int pIdx) const override { return static_cast<Tensor*>(m_Inputs[pIdx]); }

  Tensor* getOutput(unsigned int pIdx) override { return static_cast<Tensor*>(m_Outputs[pIdx]); }

  const Tensor* getOutput(unsigned int pIdx) const override { return static_cast<Tensor*>(m_Outputs[pIdx]); }

  // clang-format off
  // Inputs getters
  const Tensor* getInput() const { return getInput(kInput); }

  Tensor* getInput() { return getInput(kInput); }


  // Outputs getters
  const Tensor* getOutput() const { return getOutput(kOutput); }

  Tensor* getOutput() { return getOutput(kOutput); }


  // Inputs setters
  void setInput(Tensor& pTensor) { m_Inputs[kInput] = &pTensor; }


  // Outputs setters
  void setOutput(Tensor& pTensor) { m_Outputs[kOutput] = &pTensor; }

  // clang-format on

  void printAttributes(std::ostream& pOS) const override;

  void accept(ComputeVisitor& pVisitor) override { pVisitor.visit(*this); }

  void accept(ComputeVisitor& pVisitor) const override { pVisitor.visit(*this); }

  static bool classof(const ComputeOperator* pOp);

protected:
  // clang-format off
  IntAttr m_Blocksize;
  // clang-format on
};

} // namespace of onnc

#endif
