#include "ConvGeometry.hpp"

#include <algorithm>
#include <limits>

#include "src/Accelerators/PIM/Common/IR/ShapeUtils.hpp"

namespace onnx_mlir {

namespace {

static const ONNXToSpatialPlanningOptions& defaultPlanningOptions() {
  static const ONNXToSpatialPlanningOptions options {
    std::numeric_limits<uint64_t>::max(),
    std::numeric_limits<uint64_t>::max(),
    spatial::ConvLoweringStrategy::Auto,
    false,
  };
  return options;
}

} // namespace

const ONNXToSpatialPlanningOptions& ConvLoweringState::planningOptions() const {
  return options ? *options : defaultPlanningOptions();
}

bool isDepthwiseConv(int64_t group, int64_t numChannelsIn, int64_t numChannelsOut, int64_t numChannelsInPerGroup) {
  return group == numChannelsIn && numChannelsInPerGroup == 1 && numChannelsOut % group == 0;
}

void classifyConvProblem(ConvProblem& problem) {
  problem.isDepthwise = isDepthwiseConv(
    problem.group, problem.numChannelsIn, problem.numChannelsOut,
    problem.numChannelsInPerGroup);
  problem.isGrouped = problem.group > 1;
  problem.isPointwise = problem.wHeight == 1 && problem.wWidth == 1
    && problem.strideHeight == 1 && problem.strideWidth == 1
    && problem.dilationHeight == 1 && problem.dilationWidth == 1
    && problem.padHeightBegin == 0 && problem.padHeightEnd == 0
    && problem.padWidthBegin == 0 && problem.padWidthEnd == 0;
}

ConvGeometry buildConvGeometry(const ConvProblem& problem,
                               const spatial::SpatialTargetResources& target) {
  ConvGeometry geo {
    problem.numChannelsInPerGroup * problem.wHeight * problem.wWidth,
    problem.numChannelsOutPerGroup,
    problem.batchSize * problem.outHeight * problem.outWidth,
    static_cast<int64_t>(target.matrixShape.rows),
    static_cast<int64_t>(target.matrixUnitsPerProcessor),
    1,
    0,
  };
  geo.pack = std::max<int64_t>(1, geo.xbarSize / std::max<int64_t>(geo.k, geo.c));
  geo.im2colElements = static_cast<uint64_t>(std::max<int64_t>(0, geo.p)) * static_cast<uint64_t>(std::max<int64_t>(0, geo.k));
  return geo;
}

static ConvMaterializationKind getMaterializationKind(
    spatial::ConvLoweringStrategy strategy) {
  switch (strategy) {
  case spatial::ConvLoweringStrategy::Depthwise:
    return ConvMaterializationKind::StructuredDepthwise;
  case spatial::ConvLoweringStrategy::Legacy:
  case spatial::ConvLoweringStrategy::PackedIm2Col:
    return ConvMaterializationKind::PackedIm2Col;
  case spatial::ConvLoweringStrategy::StreamedPatch:
  case spatial::ConvLoweringStrategy::OutputChannelTiled:
  case spatial::ConvLoweringStrategy::Tiled2D:
    return ConvMaterializationKind::StreamedPatch;
  case spatial::ConvLoweringStrategy::StreamedPacked:
    return ConvMaterializationKind::StreamedPacked;
  case spatial::ConvLoweringStrategy::InputKTiled:
    return ConvMaterializationKind::InputKTiled;
  case spatial::ConvLoweringStrategy::Auto:
    break;
  }
  llvm_unreachable("auto is not a Conv materialization kind");
}

static bool fitsSingleCrossbar(const ConvGeometry& geo) {
  return geo.k <= geo.xbarSize && geo.c <= geo.xbarSize;
}

static bool fitsPackedIm2Col(const ConvGeometry& geo,
                             const ONNXToSpatialPlanningOptions& options) {
  return fitsSingleCrossbar(geo) && geo.pack >= 2
    && geo.im2colElements <= options.convIm2colMaxElements;
}

mlir::FailureOr<ConvPlan> makeConvPlan(const ConvProblem& problem,
                                       spatial::ConvLoweringStrategy strategy,
                                       const spatial::SpatialTargetResources& target,
                                       const ONNXToSpatialPlanningOptions& options) {
  ConvGeometry geo = buildConvGeometry(problem, target);
  auto plan = [&]() { return ConvPlan {getMaterializationKind(strategy)}; };
  auto ifApplicable = [&](bool applicable) -> mlir::FailureOr<ConvPlan> {
    return applicable ? mlir::FailureOr<ConvPlan>(plan()) : mlir::FailureOr<ConvPlan>(mlir::failure());
  };
  switch (strategy) {
  case spatial::ConvLoweringStrategy::Auto:
    return mlir::failure();
  case spatial::ConvLoweringStrategy::Legacy:
    return plan();
  case spatial::ConvLoweringStrategy::Depthwise:
    return ifApplicable(problem.isDepthwise);
  case spatial::ConvLoweringStrategy::PackedIm2Col:
    return ifApplicable(fitsPackedIm2Col(geo, options));
  case spatial::ConvLoweringStrategy::StreamedPatch:
    return ifApplicable(fitsSingleCrossbar(geo));
  case spatial::ConvLoweringStrategy::StreamedPacked:
    return ifApplicable(fitsSingleCrossbar(geo) && geo.pack >= 2);
  case spatial::ConvLoweringStrategy::OutputChannelTiled:
    return ifApplicable(geo.k <= geo.xbarSize && geo.c > geo.xbarSize);
  case spatial::ConvLoweringStrategy::InputKTiled:
    return ifApplicable(geo.k > geo.xbarSize && geo.c <= geo.xbarSize);
  case spatial::ConvLoweringStrategy::Tiled2D:
    return ifApplicable(geo.k > geo.xbarSize && geo.c > geo.xbarSize);
  }
  llvm_unreachable("unknown Conv lowering strategy");
}

uint64_t chooseStreamChunkPositions(const ConvGeometry& geo,
                                    int64_t packFactor,
                                    const ONNXToSpatialPlanningOptions& options) {
  const uint64_t patchElements = static_cast<uint64_t>(std::max<int64_t>(1, geo.k));
  uint64_t chunkPositions = std::max<uint64_t>(1, options.convIm2colMaxElements / patchElements);
  chunkPositions = std::min<uint64_t>(chunkPositions, static_cast<uint64_t>(std::max<int64_t>(1, geo.p)));
  chunkPositions = std::min<uint64_t>(chunkPositions, std::max<uint64_t>(1, options.convStreamChunkPositions));

  if (packFactor > 1 && chunkPositions > static_cast<uint64_t>(packFactor)) {
    chunkPositions -= chunkPositions % static_cast<uint64_t>(packFactor);
    chunkPositions = std::max<uint64_t>(chunkPositions, static_cast<uint64_t>(packFactor));
  }
  return std::max<uint64_t>(1, chunkPositions);
}

RowInterval computeConvInputRowsForOutputRows(RowInterval outputRows, const ConvProblem& problem) {
  const int64_t rawBegin = outputRows.begin * problem.strideHeight - problem.padHeightBegin;
  const int64_t rawEnd =
    (outputRows.end - 1) * problem.strideHeight - problem.padHeightBegin
    + problem.dilationHeight * (problem.wHeight - 1) + 1;
  return {std::max<int64_t>(0, rawBegin), std::min<int64_t>(problem.xHeight, rawEnd)};
}

ConvRowDemand buildConvRowDemand(RowInterval outputRows, const ConvProblem& problem) {
  ConvRowDemand demand;
  demand.outputRows = outputRows;
  demand.neededInputRows = computeConvInputRowsForOutputRows(outputRows, problem);
  demand.acquiredInputRows = demand.neededInputRows;

  const int64_t rawBegin = outputRows.begin * problem.strideHeight - problem.padHeightBegin;
  const int64_t rawEnd =
    (outputRows.end - 1) * problem.strideHeight - problem.padHeightBegin
    + problem.dilationHeight * (problem.wHeight - 1) + 1;
  demand.topHaloRows = std::max<int64_t>(0, -rawBegin);
  demand.bottomHaloRows = std::max<int64_t>(0, rawEnd - problem.xHeight);
  demand.acquiredInputRows = demand.neededInputRows;
  return demand;
}

} // namespace onnx_mlir