#pragma once

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

#include "src/Accelerators/PIM/Conversion/ONNXToSpatial/Common/ContractionPlanning.hpp"
#include "src/Accelerators/PIM/Dialect/Spatial/SpatialTargetInfo.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 {
  mlir::Operation* diagnosticAnchor = nullptr;
  mlir::Value x;
  mlir::Value w;
  mlir::Value b;
  const spatial::SpatialTargetInfo* target = nullptr;

  const spatial::SpatialTargetInfo& targetInfo() const { return *target; }
};

struct ConvGeometry {
  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 k;
  int64_t c;
  int64_t p;
  int64_t xbarSize;
  int64_t matrixUnitsPerProcessor;
  int64_t pack;
  uint64_t im2colElements;
  bool hasBias;
  bool isDepthwise;
};

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,
  PointwiseContraction,
  PackedIm2Col,
  StreamedPatch,
  StreamedPacked,
  OutputChannelTiled,
  InputKTiled,
  Tiled2D,
};

struct ConvPlan {
  ConvGeometry geometry;
  spatial::ConvLoweringStrategy strategy = spatial::ConvLoweringStrategy::Auto;
  ConvMaterializationKind materializationKind = ConvMaterializationKind::PackedIm2Col;
  int64_t laneCount = 0;
  int64_t mvmCount = 0;
  int64_t vectorCount = 0;
  int64_t reductionCount = 0;
  uint64_t weightElements = 0;
  uint64_t scratchElements = 0;
  uint64_t materializationElements = 0;
  uint64_t communicationElements = 0;
  spatial::PhysicalLayout resultLayout = spatial::PhysicalLayout::DenseNCHW;
  bool consumesRowStrip = false;
  bool producesRowStrip = false;
  bool requiresInputMaterialization = false;
  bool requiresOutputMaterialization = false;
  bool usesContraction = false;
  bool hasContractionPlan = false;
  ContractionPlan contraction;
};

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::SpatialTargetInfo& target);

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

llvm::SmallVector<ConvPlan, 8> buildConvPlanCandidates(
    const ConvProblem& problem, const spatial::SpatialTargetInfo& target);

uint64_t chooseStreamChunkPositions(const ConvGeometry& geo,
                                    int64_t packFactor,
                                    const spatial::SpatialTargetInfo& target);

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

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

} // namespace onnx_mlir