#pragma once

#include "mlir/IR/BuiltinTypes.h"
#include "mlir/IR/Value.h"

#include "src/Accelerators/PIM/Conversion/ONNXToSpatial/ONNXToSpatialOptions.hpp"
#include "src/Accelerators/PIM/Dialect/Spatial/SpatialTargetResources.hpp"
#include "src/Accelerators/PIM/Dialect/Spatial/SpatialOps.hpp"

#include <cstdint>

namespace mlir {
class Operation;
} // namespace mlir

namespace onnx_mlir {

struct ConvProblem {
  mlir::RankedTensorType xType;
  mlir::RankedTensorType wType;
  mlir::RankedTensorType outType;
  int64_t batchSize;
  int64_t numChannelsIn;
  int64_t xHeight;
  int64_t xWidth;
  int64_t numChannelsOut;
  int64_t wHeight;
  int64_t wWidth;
  int64_t outHeight;
  int64_t outWidth;
  int64_t group;
  int64_t numChannelsInPerGroup;
  int64_t numChannelsOutPerGroup;
  int64_t padHeightBegin;
  int64_t padHeightEnd;
  int64_t padWidthBegin;
  int64_t padWidthEnd;
  int64_t strideHeight;
  int64_t strideWidth;
  int64_t dilationHeight;
  int64_t dilationWidth;
  bool hasBias;
  bool isDepthwise = false;
  bool isGrouped = false;
  bool isPointwise = false;
};

struct ConvLoweringState {
  ConvProblem problem;
  mlir::Operation* diagnosticAnchor = nullptr;
  mlir::Value x;
  mlir::Value w;
  mlir::Value b;
  const spatial::SpatialTargetResources* target = nullptr;
  const ONNXToSpatialPlanningOptions* options = nullptr;

  const spatial::SpatialTargetResources& targetInfo() const { return *target; }
  const ONNXToSpatialPlanningOptions& planningOptions() const;
};

struct ConvGeometry {
  int64_t k;
  int64_t c;
  int64_t p;
  int64_t xbarSize;
  int64_t matrixUnitsPerProcessor;
  int64_t pack;
  uint64_t im2colElements;
};

struct RowInterval {
  int64_t begin = 0;
  int64_t end = 0;
};

struct ConvRowDemand {
  RowInterval outputRows;
  RowInterval neededInputRows;
  RowInterval acquiredInputRows;
  int64_t topHaloRows = 0;
  int64_t bottomHaloRows = 0;
};

enum class ConvMaterializationKind : uint8_t {
  StructuredDepthwise,
  PackedIm2Col,
  StreamedPatch,
  StreamedPacked,
  InputKTiled,
};

struct ConvPlan {
  ConvMaterializationKind kind = ConvMaterializationKind::PackedIm2Col;
};

bool isDepthwiseConv(int64_t group, int64_t numChannelsIn, int64_t numChannelsOut, int64_t numChannelsInPerGroup);

void classifyConvProblem(ConvProblem& problem);

ConvGeometry buildConvGeometry(const ConvProblem& problem,
                               const spatial::SpatialTargetResources& target);

mlir::FailureOr<ConvPlan> makeConvPlan(const ConvProblem& problem,
                                       spatial::ConvLoweringStrategy strategy,
                                       const spatial::SpatialTargetResources& target,
                                       const ONNXToSpatialPlanningOptions& options);

uint64_t chooseStreamChunkPositions(const ConvGeometry& geo,
                                    int64_t packFactor,
                                    const ONNXToSpatialPlanningOptions& options);

RowInterval computeConvInputRowsForOutputRows(RowInterval outputRows, const ConvProblem& problem);

ConvRowDemand buildConvRowDemand(RowInterval outputRows, const ConvProblem& problem);

} // namespace onnx_mlir