#include "mlir/Conversion/AffineToStandard/AffineToStandard.h"
#include "mlir/Transforms/Passes.h"

#include "llvm/Support/Error.h"
#include "llvm/Support/ErrorHandling.h"
#include "llvm/Support/JSON.h"
#include "llvm/Support/MemoryBuffer.h"
#include "llvm/Support/Path.h"

#include "llvm/ADT/SmallString.h"
#include <cmath>
#include <limits>
#include <tuple>

#include "src/Accelerators/PIM/Common/PimCommon.hpp"
#include "src/Accelerators/PIM/Compiler/PimCompilerOptions.hpp"
#include "PimPipelineSelection.hpp"
#include "PimInstructionTiming.hpp"
#include "PimPeriodicAnalysis.hpp"
#include "src/Accelerators/PIM/Compiler/PimCompilerUtils.hpp"
#include "src/Accelerators/PIM/Conversion/ONNXToSpatial/ONNXToSpatialOptions.hpp"
#include "src/Accelerators/PIM/Dialect/Pim/PimOps.hpp"
#include "src/Accelerators/PIM/Dialect/Spatial/SpatialTargetResources.hpp"
#include "src/Accelerators/PIM/Dialect/Spatial/Passes/Transforms/MergeComputeNodes/SpatialDataflowCsvExporter.hpp"
#include "src/Accelerators/PIM/Dialect/Spatial/Passes/Transforms/MergeComputeNodes/Scheduling/SchedulingTarget.hpp"
#include "src/Accelerators/PIM/Passes/PIMPasses.h"
#include "src/Compiler/CompilerPasses.hpp"

#define DEBUG_TYPE "PimCompilerUtils"

using namespace mlir;
using namespace onnx_mlir;

namespace onnx_mlir {

namespace {

void setDefaultPimInterProcessorLatencies(
    spatial::SchedulingTarget& target) {
  size_t rows = static_cast<size_t>(
    std::sqrt(static_cast<long double>(target.processorCount)));
  while (rows > 1 && target.processorCount % rows != 0)
    --rows;
  size_t columns = (target.processorCount + rows - 1) / rows;

  target.interProcessorLatencyNs.assign(
    target.processorCount * target.processorCount, 0);
  Cost latencySum = 0;
  size_t pairCount = 0;
  for (size_t source = 0; source < target.processorCount; ++source) {
    for (size_t destination = 0;
         destination < target.processorCount; ++destination) {
      if (source == destination)
        continue;
      size_t sourceRow = source / columns;
      size_t sourceColumn = source % columns;
      size_t destinationRow = destination / columns;
      size_t destinationColumn = destination % columns;
      size_t rowDistance = sourceRow > destinationRow
                         ? sourceRow - destinationRow
                         : destinationRow - sourceRow;
      size_t columnDistance = sourceColumn > destinationColumn
                            ? sourceColumn - destinationColumn
                            : destinationColumn - sourceColumn;
      Cost latency = static_cast<Cost>(2 + rowDistance + columnDistance);
      target.interProcessorLatencyNs[
        source * target.processorCount + destination] = latency;
      latencySum = checkedAdd(latencySum, latency);
      ++pairCount;
    }
  }
  target.averageInterProcessorLatencyNs =
    pairCount == 0
      ? 0
      : (latencySum + static_cast<Cost>(pairCount) - 1)
          / static_cast<Cost>(pairCount);
}

spatial::SchedulingTarget getDefaultPimSchedulingTarget() {
  spatial::SchedulingTarget target;
  target.processorCount = static_cast<size_t>(coresCount.getValue());
  target.residentWeightCapacity = crossbarCountInCore.getValue();
  target.matrixRows = crossbarSize.getValue();
  // The oracle charges eight input reads for each matrix operation.
  target.matrixInputReadMultiplier = 8;
  target.matrixColumns = crossbarSize.getValue();
  target.synchronizationRegisterCount = kPimEventRegisterCount;
  target.sharedMemoryWidthBytes = 16;
  target.sharedMemoryReadLatencyNs = 8;
  target.sharedMemoryWriteLatencyNs = 10;
  target.zeroSplatUsesSharedMemory = false;

  setDefaultPimInterProcessorLatencies(target);
  return target;
}

spatial::ConvLoweringStrategy getSpatialConvLoweringStrategy(PimConvLoweringType strategy) {
  switch (strategy) {
  case PimConvLoweringAuto: return spatial::ConvLoweringStrategy::Auto;
  case PimConvLoweringLegacy: return spatial::ConvLoweringStrategy::Legacy;
  case PimConvLoweringDepthwise: return spatial::ConvLoweringStrategy::Depthwise;
  case PimConvLoweringPackedIm2Col: return spatial::ConvLoweringStrategy::PackedIm2Col;
  case PimConvLoweringStreamedPatch: return spatial::ConvLoweringStrategy::StreamedPatch;
  case PimConvLoweringStreamedPacked: return spatial::ConvLoweringStrategy::StreamedPacked;
  case PimConvLoweringOutputChannelTiled: return spatial::ConvLoweringStrategy::OutputChannelTiled;
  case PimConvLoweringInputKTiled: return spatial::ConvLoweringStrategy::InputKTiled;
  case PimConvLoweringTiled2D: return spatial::ConvLoweringStrategy::Tiled2D;
  }
  llvm_unreachable("unknown Pim Conv lowering strategy");
}

spatial::SpatialDataflowExportStage getPimSpatialDataflowExportStage(
    PimSpatialDataflowExportType stage) {
  switch (stage) {
  case SpatialDataflowExportNone:     return spatial::SpatialDataflowExportStage::None;
  case SpatialDataflowExportSpatial1: return spatial::SpatialDataflowExportStage::Spatial1;
  case SpatialDataflowExportSpatial2: return spatial::SpatialDataflowExportStage::Spatial2;
  case SpatialDataflowExportSpatial3: return spatial::SpatialDataflowExportStage::Spatial3;
  case SpatialDataflowExportSpatial4: return spatial::SpatialDataflowExportStage::Spatial4;
  case SpatialDataflowExportAll:      return spatial::SpatialDataflowExportStage::All;
  }
  llvm_unreachable("unknown Pim Spatial dataflow export stage");
}

spatial::SpatialTargetResources getPimSpatialTargetResources(const spatial::SchedulingTarget& target) {
  spatial::SpatialTargetResources resources;
  resources.matrixShape = {target.matrixRows, target.matrixColumns};
  resources.matrixUnitsPerProcessor = target.residentWeightCapacity;
  resources.processorCount = target.processorCount;
  resources.vectorWidth = target.vectorWidth;
  if (failed(resources.verify()))
    llvm::report_fatal_error("Pim target resources are incomplete");
  return resources;
}

ONNXToSpatialPlanningOptions getPimONNXToSpatialPlanningOptions() {
  ONNXToSpatialPlanningOptions options;
  options.convIm2colMaxElements = pimConvIm2colMaxElements.getValue();
  options.convStreamChunkPositions = pimConvStreamChunkPositions.getValue();
  options.forcedConvStrategy = getSpatialConvLoweringStrategy(pimConvLowering.getValue());
  options.reportConvLowering = pimReportConvLowering.getValue();
  return options;
}

const llvm::json::Object& requireObject(const llvm::json::Object& object,
                                        llvm::StringRef key,
                                        llvm::StringRef path) {
  const llvm::json::Object* nested = object.getObject(key);
  if (!nested)
    llvm::report_fatal_error("Pim target config is missing object '" + path + "." + key + "'");
  return *nested;
}

Cost getConfigCost(const llvm::json::Object& object,
                   llvm::StringRef key,
                   Cost fallback,
                   bool allowZero = false) {
  std::optional<double> number = object.getNumber(key);
  if (!number)
    return fallback;
  if (!std::isfinite(*number) || *number < 0.0 || (!allowZero && *number == 0.0)
      || *number > static_cast<double>(std::numeric_limits<Cost>::max()))
    llvm::report_fatal_error("Pim target config field '" + key + "' must be a valid positive number");
  return static_cast<Cost>(std::ceil(*number));
}

std::pair<size_t, size_t> getConfigPair(const llvm::json::Object& object,
                                        llvm::StringRef key) {
  const llvm::json::Array* values = object.getArray(key);
  if (!values || values->size() != 2)
    llvm::report_fatal_error("Pim target config field '" + key + "' must contain two integers");
  std::optional<int64_t> first = (*values)[0].getAsInteger();
  std::optional<int64_t> second = (*values)[1].getAsInteger();
  if (!first || !second || *first <= 0 || *second <= 0)
    llvm::report_fatal_error("Pim target config field '" + key + "' must contain two positive integers");
  return {static_cast<size_t>(*first), static_cast<size_t>(*second)};
}

void loadPimInterProcessorLatencies(
    PimInstructionTiming& target,
    const llvm::json::Object& network,
    std::optional<int64_t> memorySwitch) {
  std::optional<llvm::StringRef> filename =
    network.getString("net_config_file_path");
  if (!filename)
    llvm::report_fatal_error(
      "Pim target config is missing network latency file path");

  llvm::SmallString<256> networkPath(*filename);
  if (!llvm::sys::path::is_absolute(networkPath)) {
    llvm::SmallString<256> configDirectory(pimTargetConfig.getValue());
    llvm::sys::path::remove_filename(configDirectory);
    llvm::sys::path::append(configDirectory, networkPath);
    networkPath = configDirectory;
  }

  auto buffer = llvm::MemoryBuffer::getFile(networkPath);
  if (!buffer)
    llvm::report_fatal_error(
      llvm::Twine("failed to read Pim network config '")
      + networkPath + "': " + buffer.getError().message());
  auto parsed = llvm::json::parse(buffer.get()->getBuffer());
  if (!parsed)
    llvm::report_fatal_error(
      llvm::Twine("failed to parse Pim network config '")
      + networkPath + "': " + llvm::toString(parsed.takeError()));
  const llvm::json::Object* root = parsed->getAsObject();
  const llvm::json::Object* latencies =
    root ? root->getObject("latency") : nullptr;
  if (!latencies)
    llvm::report_fatal_error(
      "Pim network config is missing its latency matrix");

  target.interProcessorLatencyNs.assign(
    target.processorCount * target.processorCount, 0);
  Cost latencySum = 0;
  size_t pairCount = 0;
  for (size_t source = 0; source < target.processorCount; ++source) {
    std::string sourceKey = std::to_string(source);
    const llvm::json::Object* row = latencies->getObject(sourceKey);
    if (!row)
      llvm::report_fatal_error(
        llvm::Twine("Pim network config is missing latency row ")
        + sourceKey);
    for (size_t destination = 0;
         destination < target.processorCount; ++destination) {
      if (source == destination)
        continue;
      std::string destinationKey = std::to_string(destination);
      std::optional<double> latency = row->getNumber(destinationKey);
      if (!latency || !std::isfinite(*latency) || *latency <= 0.0)
        llvm::report_fatal_error(
          llvm::Twine("Pim network config is missing latency ")
          + sourceKey + " -> " + destinationKey);
      Cost roundedLatency = static_cast<Cost>(std::ceil(*latency));
      target.interProcessorLatencyNs[
        source * target.processorCount + destination] = roundedLatency;
      latencySum = checkedAdd(latencySum, roundedLatency);
      ++pairCount;
    }
  }
  target.averageInterProcessorLatencyNs =
    pairCount == 0
      ? 0
      : (latencySum + static_cast<Cost>(pairCount) - 1)
          / static_cast<Cost>(pairCount);
  if (memorySwitch) {
    const auto *memoryRow = latencies->getObject(std::to_string(*memorySwitch));
    for (size_t core = 0; core < target.processorCount; ++core) {
      const auto *row = latencies->getObject(std::to_string(core));
      auto request = row ? row->getNumber(std::to_string(*memorySwitch)) : std::nullopt;
      auto response = memoryRow ? memoryRow->getNumber(std::to_string(core)) : std::nullopt;
      if (!request || !response || !std::isfinite(*request)
          || !std::isfinite(*response) || *request < 0 || *response < 0
          || *request >= static_cast<double>(std::numeric_limits<Cost>::max())
          || *response >= static_cast<double>(std::numeric_limits<Cost>::max())) {
        target.sharedMemoryRequestLatencyNs.clear();
        target.sharedMemoryResponseLatencyNs.clear();
        break;
      }
      target.sharedMemoryRequestLatencyNs.push_back(std::ceil(*request));
      target.sharedMemoryResponseLatencyNs.push_back(std::ceil(*response));
    }
  }
}

PimInstructionTiming loadPimInstructionTiming() {
  PimInstructionTiming target;
  static_cast<spatial::SchedulingTarget &>(target) = getDefaultPimSchedulingTarget();
  if (pimTargetConfig.empty())
    return target;

  auto buffer = llvm::MemoryBuffer::getFile(pimTargetConfig);
  if (!buffer)
    llvm::report_fatal_error(
      llvm::Twine("failed to read Pim target config '")
      + pimTargetConfig.getValue() + "': " + buffer.getError().message());
  auto parsed = llvm::json::parse(buffer.get()->getBuffer());
  if (!parsed)
    llvm::report_fatal_error(
      llvm::Twine("failed to parse Pim target config '")
      + pimTargetConfig.getValue() + "': "
      + llvm::toString(parsed.takeError()));
  const llvm::json::Object* root = parsed->getAsObject();
  if (!root)
    llvm::report_fatal_error("Pim target config must contain a JSON object");

  const llvm::json::Object& chip = requireObject(*root, "chip_config", "root");
  const llvm::json::Object& core = requireObject(chip, "core_config", "chip_config");
  const llvm::json::Object& matrix =
    requireObject(core, "matrix_config", "chip_config.core_config");
  const llvm::json::Object& localMemory =
    requireObject(core, "local_memory_config", "chip_config.core_config");
  const llvm::json::Object& network =
    requireObject(chip, "network_config", "chip_config");
  // Older target configurations omit global memory and use adapter defaults.
  if (chip.get("global_memory_config")) {
    const llvm::json::Object& sharedMemory =
      requireObject(chip, "global_memory_config", "chip_config");
    Cost memoryPeriod = getConfigCost(sharedMemory, "period", Cost{1});
    target.sharedMemoryWidthBytes = getConfigCost(
      sharedMemory, "data_width", target.sharedMemoryWidthBytes);
    target.sharedMemoryReadLatencyNs = checkedMultiply(memoryPeriod,
      getConfigCost(sharedMemory, "read_latency_cycle", Cost{8}, true));
    target.sharedMemoryWriteLatencyNs = checkedMultiply(memoryPeriod,
      getConfigCost(sharedMemory, "write_latency_cycle", Cost{10}, true));
  }

  std::optional<int64_t> coreCount = chip.getInteger("core_cnt");
  if (!coreCount || *coreCount <= 0)
    llvm::report_fatal_error("Pim target config field 'core_cnt' must be a positive integer");
  target.processorCount = static_cast<size_t>(*coreCount);
  target.residentWeightCapacity =
    getConfigCost(matrix, "xbar_array_count", target.residentWeightCapacity);
  std::tie(target.matrixRows, target.matrixColumns) =
    getConfigPair(matrix, "xbar_size");

  if (target.processorCount != static_cast<size_t>(coresCount.getValue())
      || target.residentWeightCapacity != crossbarCountInCore.getValue()
      || target.matrixRows != crossbarSize.getValue()
      || target.matrixColumns != crossbarSize.getValue())
    llvm::report_fatal_error("Pim target config resources do not match --core-count, "
                             "--crossbar-count, and --crossbar-size");
  loadPimInterProcessorLatencies(target, network,
    core.getInteger("global_memory_switch_id"));

  target.processorPeriodNs =
    getConfigCost(core, "period", target.processorPeriodNs);
  target.signalIssueNs = target.processorPeriodNs;
  target.controlMessageBytes = 1;
  target.sharedMemoryReadReturnsPayload = false;
  target.peerPermissionThenSourceRead = true;
  target.localMemoryWidthBytes =
    getConfigCost(localMemory, "data_width", target.localMemoryWidthBytes);
  target.localMemoryPeriodNs =
    getConfigCost(localMemory, "period", target.localMemoryPeriodNs);
  target.localMemoryReadLatencyCycles =
    getConfigCost(localMemory, "read_latency_cycle", target.localMemoryReadLatencyCycles);
  target.localMemoryWriteLatencyCycles =
    getConfigCost(localMemory, "write_latency_cycle", target.localMemoryWriteLatencyCycles);
  target.transferWidthBytes =
    getConfigCost(network, "bus_width", target.transferWidthBytes);
  target.vectorWidth = getConfigCost(core, "vector_width", target.vectorWidth);
  target.vectorLatencyCycles =
    getConfigCost(core, "vector_latency_cycle", target.vectorLatencyCycles);

  target.matrixPeriodNs =
    getConfigCost(matrix, "period", target.matrixPeriodNs);
  target.matrixInputResolutionBits =
    getConfigCost(matrix, "dac_resolution", target.matrixInputResolutionBits);
  target.matrixInputLatencyCycles =
    getConfigCost(matrix, "dac_latency_cycle", target.matrixInputLatencyCycles);
  target.matrixInputParallelism =
    getConfigCost(matrix, "dac_count", target.matrixInputParallelism);
  target.matrixReadLatencyNs =
    getConfigCost(matrix, "xbar_latency", target.matrixReadLatencyNs);
  target.matrixSampleLatencyCycles =
    getConfigCost(matrix, "sample_hold_latency_cycle", target.matrixSampleLatencyCycles);
  target.matrixOutputLatencyCycles =
    getConfigCost(matrix, "adc_latency_cycle", target.matrixOutputLatencyCycles);
  target.matrixOutputParallelism =
    getConfigCost(matrix, "adc_count", target.matrixOutputParallelism);
  target.matrixShiftLatencyCycles =
    getConfigCost(matrix, "shift_adder_latency_cycle", target.matrixShiftLatencyCycles);
  target.matrixBufferLatencyCycles =
    getConfigCost(matrix, "output_buffer_latency_cycle", target.matrixBufferLatencyCycles);
  target.matrixInputBufferLatencyCycles =
    getConfigCost(matrix,
                  "input_buffer_latency_cycle",
                  target.matrixInputBufferLatencyCycles,
                  /*allowZero=*/true);
  target.matrixPipeline = matrix.getBoolean("pipeline_mode").value_or(target.matrixPipeline);
  return target;
}

} // namespace

PimInstructionTiming getPimInstructionTiming() {
  return loadPimInstructionTiming();
}

spatial::SchedulingTarget getPimSchedulingTarget() {
  auto timing = getPimInstructionTiming();
  auto target = static_cast<spatial::SchedulingTarget>(timing);
  configurePimScheduleCosts(target, std::move(timing));
  return target;
}

llvm::Expected<int64_t> getPimPeriodicTimingContract() {
  auto buffer = llvm::MemoryBuffer::getFile(pimTargetConfig);
  if (!buffer) return llvm::createStringError("queue selection requires --pim-target-config with instruction timing");
  auto parsed = llvm::json::parse(buffer.get()->getBuffer());
  if (!parsed) return parsed.takeError();
  const auto *root = parsed->getAsObject();
  const auto *chip = root ? root->getObject("chip_config") : nullptr;
  const auto *core = chip ? chip->getObject("core_config") : nullptr;
  const auto *matrix = core ? core->getObject("matrix_config") : nullptr;
  if (!matrix || core->getInteger("rob_size").value_or(1) != 1)
    return llvm::createStringError("queue selection requires a target with ROB capacity one and matrix timing");
  auto bits = matrix->getInteger("cell_precision").value_or(2);
  if (bits <= 0) return llvm::createStringError("matrix cell precision must be positive");
  return bits;
}

void addPimLoweringPasses(PassManager &pm,
    const spatial::SpatialTargetResources &targetResources, PimEmissionTargetType target,
    bool reportProgress) {
  auto message = [&](llvm::StringRef text) {
    if (reportProgress) pm.addPass(createMessagePass(text.str()));
  };
  if (target >= EmitPim) {
    pm.addPass(createSpatialToPimPass(targetResources));
    message("Spatial lowered to Pim");
  }

  if (target >= EmitPimBufferized) {
    pm.addPass(createPimBufferizationPreparationPass(pimVerifyBufferizationCopyFreedom.getValue()));
    pm.addPass(createPimOneShotBufferizationPass());
    pm.addPass(createPimMemoryNormalizationPass());
    pm.addPass(createPimBufferizationVerificationPass());
    message("Pim bufferized");
  }

  if (target >= EmitPimCodegen) {
    pm.addPass(mlir::createLowerAffinePass());
    pm.addPass(createPimHostConstantFoldingPass());
    message("Pim host constants folded");
    pm.addPass(createPimInstructionSelectionPass());
    message("Pim instructions selected");
    pm.addPass(createPimLocalMemoryPlanningPass());
    message("Pim local memory planned");
    pm.addPass(createPimVerificationPass(targetResources, pimDetectCommunicationDeadlock.getValue()));
    message("Pim verified");
  }
}

void addPassesPim(OwningOpRef<ModuleOp>& module,
                  PassManager& pm,
                  EmissionTargetType& emissionTarget,
                  std::string outputNameNoExt) {
  verifyPimCompilerOptions();
  spatial::SchedulingTarget schedulingTarget = getPimSchedulingTarget();
  spatial::SpatialTargetResources targetResources = getPimSpatialTargetResources(schedulingTarget);

  if (pimOnlyCodegen) {
    pm.addPass(createPimInstructionSelectionPass());
    pm.addPass(createPimLocalMemoryPlanningPass());
    pm.addPass(createPimVerificationPass(targetResources, pimDetectCommunicationDeadlock.getValue()));
    pm.addPass(createEmitPimCodePass());
    return;
  }

  if (emissionTarget >= EmitONNXIR)
    addONNXToMLIRPasses(pm, /*target CPU*/ false);

  if (pimEmissionTarget >= EmitSpatial) {
    ONNXToSpatialPlanningOptions planningOptions = getPimONNXToSpatialPlanningOptions();
    spatial::SpatialDataflowExportStage exportStage =
      getPimSpatialDataflowExportStage(pimExportSpatialDataflow.getValue());
    pm.addPass(createONNXToSpatialPass(targetResources, planningOptions));
    pm.addPass(createSpatialLayoutPlanningPass(
      targetResources, pimDisableSpatialPlanning.getValue()));
    pm.addPass(createLowerSpatialPlansPass(targetResources, planningOptions, exportStage));
    pm.addPass(createTrivialGraphComputeMergePass(
      schedulingTarget.residentWeightCapacity, exportStage));
    if (pipelineStages.getValue() > 1 && pimEmissionTarget >= EmitPimCodegen) {
      pm.addPass(createPimPipelineSelectionPass(schedulingTarget, targetResources,
        pipelineStages.getValue(), exportStage));
      pm.addPass(createEmitPimCodePass());
      return;
    } else {
      if (pipelineStages.getValue() > 1)
        pm.addPass(createMessagePass("Queue-aware pipeline: intermediate emission retains the original schedule without backend evaluation"));
      pm.addPass(spatial::createScheduleAndRealizeSpatialPass(
        schedulingTarget, exportStage, pipelineStages.getValue(), pimPeriodicDiagnostics.getValue()));
    }
    pm.addPass(createMessagePass("Onnx lowered to Spatial"));
  }

  addPimLoweringPasses(pm, targetResources, pimEmissionTarget);
  if (pimEmissionTarget >= EmitPimCodegen) {
    pm.addPass(createEmitPimCodePass());
    pm.addPass(createMessagePass("Pim code emitted"));
  }
}

} // namespace onnx_mlir