#include "ConvGeometry.hpp"

#include <algorithm>

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

namespace onnx_mlir {

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

ConvGeometry buildConvGeometry(const ConvLoweringState& state) {
  ConvGeometry geo {
    state.batchSize,
    state.numChannelsIn,
    state.xHeight,
    state.xWidth,
    state.numChannelsOut,
    state.wHeight,
    state.wWidth,
    state.outHeight,
    state.outWidth,
    state.group,
    state.numChannelsInPerGroup,
    state.numChannelsOutPerGroup,
    state.numChannelsInPerGroup * state.wHeight * state.wWidth,
    state.numChannelsOutPerGroup,
    state.batchSize * state.outHeight * state.outWidth,
    static_cast<int64_t>(crossbarSize.getValue()),
    1,
    0,
    state.hasBias,
    isDepthwiseConv(state.group, state.numChannelsIn, state.numChannelsOut, state.numChannelsInPerGroup),
  };
  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;
}

uint64_t chooseStreamChunkPositions(const ConvGeometry& geo, int64_t packFactor) {
  const uint64_t patchElements = static_cast<uint64_t>(std::max<int64_t>(1, geo.k));
  uint64_t chunkPositions = std::max<uint64_t>(1, pimConvIm2colMaxElements / 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, pimConvStreamChunkPositions));

  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 ConvLoweringState& state) {
  const int64_t rawBegin = outputRows.begin * state.strideHeight - state.padHeightBegin;
  const int64_t rawEnd =
    (outputRows.end - 1) * state.strideHeight - state.padHeightBegin + state.dilationHeight * (state.wHeight - 1) + 1;
  return {std::max<int64_t>(0, rawBegin), std::min<int64_t>(state.xHeight, rawEnd)};
}

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

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

} // namespace onnx_mlir