#include "TestCheck.hpp"
#include "src/Accelerators/PIM/Compiler/PimPeriodicAnalysis.hpp"
#include "src/Accelerators/PIM/Compiler/PimInstructionTiming.hpp"
#include "src/Accelerators/PIM/Dialect/Spatial/Passes/Transforms/MergeComputeNodes/Scheduling/PipelineScheduling.hpp"
#include "src/Accelerators/PIM/Dialect/Spatial/Passes/Transforms/MergeComputeNodes/Scheduling/ScheduleEvaluation.hpp"
#include "src/Accelerators/PIM/Dialect/Spatial/Passes/Transforms/MergeComputeNodes/Scheduling/ResidentWeightPacking.hpp"
#include "src/Accelerators/PIM/Dialect/Spatial/Passes/Transforms/MergeComputeNodes/Scheduling/PipelinePackingCostCache.hpp"
#include "src/Accelerators/PIM/Dialect/Spatial/SpatialOps.hpp"
#include <chrono>
#include <cmath>
#include <functional>
#include <map>
#include <numeric>
#include <random>
#include <set>
#include <iostream>

using namespace onnx_mlir;
using namespace onnx_mlir::spatial;
using E = PimExecutionEvent;
static bool sawJointMove = false, sawSuccessiveMoves = false;

static void testPackingCostCacheEviction() {
  ComputeGraph graph;
  const std::vector<Cost> work{0, 11, 23, 7, 19, 3};
  for (uint32_t task = 0; task < work.size(); ++task) {
    ResidentWeight weight;
    weight.opaqueLane = task % 3;
    graph.nodes.push_back({{nullptr, task, 1}, work[task], {weight}, task});
  }
  IndexedResidentWeights weights(graph);
  PipelinePackingCostCache cache(weights, work, 1, 8);
  static_assert(PipelinePackingCostCache::kMaximumEntries == 256);
  std::vector<size_t> tasks{0, 1, 2, 3, 4, 5};
  auto check = [&](const auto &order) {
    TEST_CHECK(cache.minimumCores(order) == weights.getPackCountLowerBound(order, 1));
    for (size_t cores : {1u, 2u, 3u, 5u, 8u}) {
      Cost expected = weights.estimatePackCost(order, work, 1, cores)
        .value_or(std::numeric_limits<Cost>::max());
      TEST_CHECK(cache.estimate(order, cores) == expected);
      TEST_CHECK(cache.estimate(order, cores) == expected);
    }
    TEST_CHECK(cache.size() <= PipelinePackingCostCache::kMaximumEntries);
  };
  size_t orders = 0;
  do {
    check(tasks);
    ++orders;
  } while (std::next_permutation(tasks.begin(), tasks.end()));
  TEST_CHECK(orders == 720 && cache.size() == PipelinePackingCostCache::kMaximumEntries);
  // Revisit early entries after eviction. Ordered and constructive queries must agree.
  check(std::vector<size_t>{0, 1, 2, 3, 4, 5});
  check(std::vector<size_t>{0, 3, 1, 4, 2, 5});
  check(std::vector<size_t>{0});
  check(std::vector<size_t>{});
}

static void testPackingFeasibility() {
  ComputeGraph graph;
  for (size_t task = 0; task < 4; ++task) {
    ResidentWeight weight;
    weight.opaqueLane = task % 2;
    graph.nodes.push_back({{nullptr, static_cast<uint32_t>(task), 1}, 10, {weight}, task});
  }
  IndexedResidentWeights weights(graph);
  const std::vector<size_t> tasks{0, 1, 2, 3};
  const std::vector<Cost> work{10, 10, 10, 10};
  std::vector<size_t> assignment{99};
  TEST_CHECK(!estimateOrderedPackCost(work, weights.taskWeightIds, weights.weightCount,
    tasks, 1, 2, &assignment));
  TEST_CHECK(assignment.empty());
  TEST_CHECK(!estimateOrderedPackCost(work, weights.taskWeightIds, weights.weightCount,
    tasks, 1, 0));
  TEST_CHECK(weights.estimatePackCost(tasks, work, 1, 2) == 20);
  auto placement = weights.pack(tasks, work, 1, 2);
  TEST_CHECK(placement && (*placement)[0] == (*placement)[2]
    && (*placement)[1] == (*placement)[3] && (*placement)[0] != (*placement)[1]);
  TEST_CHECK(!weights.estimatePackCost(tasks, work, 1, 1));
  for (size_t cores = 4; cores <= 8; ++cores)
    TEST_CHECK(estimateOrderedPackCost(work, weights.taskWeightIds, weights.weightCount,
      tasks, 1, cores) == 10);
  weights.taskWeightIds[0] = {0, 1};
  TEST_CHECK(!estimateOrderedPackCost(work, weights.taskWeightIds, weights.weightCount,
    tasks, 1, 4));
  weights.taskWeightIds = {{0, 1}, {2}, {0, 1}, {2}};
  weights.weightCount = 3;
  TEST_CHECK(!estimateOrderedPackCost(work, weights.taskWeightIds, weights.weightCount,
    tasks, 2, 2));
  TEST_CHECK(weights.estimatePackCost(tasks, work, 2, 2) == 20);
}

static void testPackingCurves() {
  std::mt19937 random(1739);
  for (size_t fixture = 0; fixture < 150; ++fixture) {
    size_t count = fixture % 9, capacity = 1 + fixture % 3;
    ComputeGraph graph;
    std::vector<Cost> work;
    std::vector<size_t> tasks;
    for (uint32_t task = 0; task < count; ++task) {
      ResidentWeightSet weights;
      for (size_t id = 0; id < 4; ++id)
        if (random() % 3 == 0) { ResidentWeight w; w.opaqueLane = id; weights.push_back(w); }
      work.push_back(random() % 40);
      tasks.push_back(task);
      graph.nodes.push_back({{nullptr, task, 1}, work.back(), weights, task});
    }
    IndexedResidentWeights weights(graph);
    auto curve = weights.estimatePackCostCurve(tasks, work, capacity, 8);
    auto sparse = weights.estimatePackCostCurve(tasks, work, capacity, 8, nullptr, {1, 3, 6, 8});
    for (size_t cores : {1u, 3u, 6u, 8u}) TEST_CHECK(sparse[cores] == curve[cores]);
    for (size_t cores : {2u, 4u, 5u, 7u}) TEST_CHECK(!sparse[cores]);
    for (size_t cores = 1; cores <= 8; ++cores) {
      TEST_CHECK(curve[cores] == weights.estimatePackCost(tasks, work, capacity, cores));
      std::optional<Cost> reference;
      for (size_t mask = 0; mask < (size_t{1} << (count ? count - 1 : 0)); ++mask) {
        size_t used = count ? 1 : 0;
        Cost maximum = 0, current = 0;
        std::set<size_t> residents;
        bool fits = true;
        for (size_t task = 0; task < count; ++task) {
          if (task && ((mask >> (task - 1)) & 1)) {
            maximum = std::max(maximum, current); current = 0; residents.clear(); ++used;
          }
          current += work[task];
          for (size_t id : weights.taskWeightIds[task]) residents.insert(id);
          fits &= residents.size() <= capacity;
        }
        maximum = std::max(maximum, current);
        if (fits && used <= cores && (!reference || maximum < *reference)) reference = maximum;
      }
      TEST_CHECK(estimateOrderedPackCost(work, weights.taskWeightIds, weights.weightCount,
        tasks, capacity, cores) == reference);
    }
  }
}

static void testQuantizedFeasibility(uint32_t first) {
  mlir::MLIRContext context;
  context.loadDialect<spatial::SpatialDialect>();
  auto module = mlir::ModuleOp::create(mlir::UnknownLoc::get(&context));
  mlir::OpBuilder builder(module.getBodyRegion());
  ComputeGraph graph;
  graph.predecessors.resize(40); graph.successors.resize(40);
  for (uint32_t group = 0; group < 2; ++group) {
    const uint32_t lanes = group ? 40 - first : first;
    auto batch = spatial::SpatGraphComputeBatch::create(builder, module.getLoc(),
      mlir::TypeRange{}, builder.getI32IntegerAttr(lanes), mlir::ValueRange{}, mlir::ValueRange{});
    for (uint32_t lane = 0; lane < lanes; ++lane) {
      size_t task = graph.nodes.size();
      ResidentWeight weight; weight.opaqueLane = task;
      graph.nodes.push_back({{batch, lane, 1}, 10, {weight}, task});
      if (group) {
        graph.predecessors[task].push_back({0, {0, 0}});
        graph.successors[0].push_back({task, {0, 0}});
      }
    }
  }
  SchedulingTarget target;
  target.processorCount = 40; target.residentWeightCapacity = 1;
  target.interProcessorLatencyNs.assign(1600, 0);
  MergeScheduleResult schedule;
  std::string error;
  TEST_CHECK(mlir::succeeded(runPipelineScheduler(graph, schedule, 2, target, error)));
  std::vector<size_t> budgets(2);
  for (size_t stage : schedule.processorStages) ++budgets[stage];
  TEST_CHECK(budgets == std::vector<size_t>({first, 40 - first}));
  for (size_t task = 0; task < graph.nodes.size(); ++task)
    TEST_CHECK(schedule.processorStages[schedule.computeToCpuMap.lookup(graph.nodes[task].instance)]
      == (task < first ? 0 : 1));
  auto query = [](const MergeScheduleResult &) {
    WholeScheduleCost cost;
    cost.primary = cost.alternative = 100;
    cost.selectionMargin = kPipelineSelectionMargin;
    cost.loweringSupport = StaticLoweringSupport::Supported;
    for (auto &proof : cost.recurrence) proof = {ScheduleRecurrenceEvidence::Proved, 1, 100, {}};
    return cost;
  };
  auto depthOne = schedule;
  depthOne.processorStages.assign(target.processorCount, 0);
  auto search = optimizePipelineSchedule(graph, depthOne, 0, target, query, error, 2);
  TEST_CHECK(mlir::succeeded(search) && search->initialSchedules == 2);
  TEST_CHECK(search->schedules[1].processorStages == schedule.processorStages);
  TEST_CHECK(search->schedules[1].computeToCpuMap == schedule.computeToCpuMap);
  module.erase();
}

// Test-only workloads use stage-boundary host transfers and local direct links.
// The same queue
// engine and complete-state evidence rank reference and searched schedules.
static WholeScheduleCost queueCost(const ComputeGraph &graph,
    const MergeScheduleResult &schedule, const PimInstructionTiming &target,
    bool restartProtocol = false) {
  PimExecutionProgram program(target.processorCount);
  std::vector<std::vector<size_t>> tasks(target.processorCount);
  for (size_t task = 0; task < graph.nodes.size(); ++task)
    tasks[schedule.computeToCpuMap.lookup(graph.nodes[task].instance)].push_back(task);
  for (auto &core : tasks)
    llvm::sort(core, [&](size_t a, size_t b) {
      return schedule.computeToCpuSlotMap.lookup(graph.nodes[a].instance)
        < schedule.computeToCpuSlotMap.lookup(graph.nodes[b].instance);
    });
  struct Transfer { size_t source, target, reader, writer, reg; bool host; };
  std::vector<Transfer> transfers;
  std::vector<std::set<size_t>> readers(target.processorCount), writers(target.processorCount);
  // Branched synthetic workloads use host routing to avoid inventing a second
  // blocking-transfer scheduler in this test-only objective.
  bool branching = llvm::any_of(graph.successors, [](const auto &edges) { return edges.size() > 1; });
  for (size_t task = 0; task < graph.nodes.size(); ++task)
    for (const auto &[predecessor, dependency] : graph.predecessors[task]) {
      size_t writer = schedule.computeToCpuMap.lookup(graph.nodes[predecessor].instance);
      size_t reader = schedule.computeToCpuMap.lookup(graph.nodes[task].instance);
      if (writer == reader) continue;
      bool host = branching || schedule.processorStages[writer] != schedule.processorStages[reader];
      transfers.push_back({predecessor, task, reader, writer, 1 + transfers.size(), host});
      if (host) { readers[writer].insert(reader); writers[reader].insert(writer); }
    }
  std::vector<size_t> zero, downstream;
  if (restartProtocol)
    for (size_t core = 0; core < program.size(); ++core)
      (schedule.processorStages[core] ? downstream : zero).push_back(core);
  for (size_t core = 0; core < program.size(); ++core) {
    auto &events = program[core];
    for (size_t writer : writers[core]) events.emplace_back(E::Signal, 0, writer, 0, 1);
    size_t rank = std::find(downstream.begin(), downstream.end(), core) - downstream.begin();
    if (restartProtocol && rank < downstream.size())
      events.emplace_back(E::Signal, 0, rank ? downstream[(rank - 1) / 2] : zero.front(), rank ? 29 : 30, 1);
    if (!restartProtocol && schedule.processorStages[core] == 0) {
      size_t peers = 0;
      for (size_t peer = 0; peer < program.size(); ++peer)
        if (peer != core && schedule.processorStages[peer] == 0) {
          events.emplace_back(E::Signal, 0, peer, 100, 1); ++peers;
        }
      if (peers) events.emplace_back(E::Wait, 0, 0, 100, peers);
    }
    bool reuseWait = false;
    for (size_t task : tasks[core]) {
      for (const auto &transfer : transfers)
        if (transfer.target == task) {
          if (transfer.host) {
            events.emplace_back(E::Wait, 0, 0, transfer.reg, 1);
            events.emplace_back(E::Load, 64);
          } else events.emplace_back(E::Receive, 64, transfer.writer);
        }
      E work(E::Work, graph.nodes[task].cost); work.workClass = E::Vector;
      events.push_back(work);
      for (const auto &transfer : transfers)
        if (transfer.source == task) {
          if (!transfer.host) {
            events.emplace_back(E::Send, 64, transfer.reader); continue;
          }
          if (!reuseWait) {
            events.emplace_back(E::Wait, 0, 0, 0, readers[core].size()); reuseWait = true;
          }
          events.emplace_back(E::Store, 64);
          events.emplace_back(E::Signal, 0, transfer.reader, transfer.reg, 1);
        }
    }
    if (restartProtocol && schedule.processorStages[core] == 0) {
      size_t rankZero = std::find(zero.begin(), zero.end(), core) - zero.begin();
      for (size_t step = 1, reg = 16; step < zero.size(); step *= 2, ++reg) {
        events.emplace_back(E::Signal, 0, zero[(rankZero + step) % zero.size()], reg, 1);
        events.emplace_back(E::Wait, 0, 0, reg, 1);
      }
      if (core == zero.front() && !downstream.empty()) {
        events.emplace_back(E::Wait, 0, 0, 30, 1);
        events.emplace_back(E::Signal, 0, downstream.front(), 28, 1);
      }
    } else if (restartProtocol) {
      events.emplace_back(E::Wait, 0, 0, 28, 1);
      size_t children = (2 * rank + 1 < downstream.size()) + (2 * rank + 2 < downstream.size());
      if (children) events.emplace_back(E::Wait, 0, 0, 29, children);
      for (size_t child : {2 * rank + 1, 2 * rank + 2})
        if (child < downstream.size()) events.emplace_back(E::Signal, 0, downstream[child], 28, 1);
    }
  }
  if (restartProtocol)
    for (const auto &sequence : program)
      for (const auto &event : sequence)
        if (event.kind == E::Signal || event.kind == E::Wait)
          TEST_CHECK(event.reg >= 0 && static_cast<size_t>(event.reg) < target.synchronizationRegisterCount);
  PimPeriodicOptions options;
  options.verifyPeriodicState = true; options.searchGuidance = !restartProtocol;
  auto estimate = evaluatePimExecution(program, target, options);
  auto cost = std::move(estimate.cost);
  cost.loweringSupport = StaticLoweringSupport::Supported;
  cost.selectionMargin = 0.021;
  TEST_CHECK(cost.eligible());
  return cost;
}

static PimInstructionTiming hardware(size_t profile) {
  PimInstructionTiming target;
  target.processorCount = 5;
  target.residentWeightCapacity = 2;
  target.estimatePipelineDependencySpan = [](llvm::ArrayRef<Cost> spans) { return spans.front(); };
  target.peerPermissionThenSourceRead = true;
  target.signalIssueNs = profile + 1;
  target.controlMessageBytes = 1;
  target.sharedMemoryWidthBytes = 64;
  target.sharedMemoryReadLatencyNs = 17 + profile * 9;
  target.sharedMemoryWriteLatencyNs = 23 + profile * 7;
  target.synchronizationRegisterCount = 128;
  target.sharedMemoryRequestLatencyNs.assign(5, 2);
  target.sharedMemoryResponseLatencyNs.assign(5, 3);
  target.interProcessorLatencyNs.resize(25);
  for (size_t a = 0; a < 5; ++a)
    for (size_t b = 0; b < 5; ++b) {
      const size_t irregular[] = {0, 1, 5, 6, 20};
      size_t x = profile ? irregular[a] : a, y = profile ? irregular[b] : b;
      target.interProcessorLatencyNs[a * 5 + b] = x > y ? x - y : y - x;
    }
  return target;
}

static void testInitializerPackingWitness() {
  mlir::MLIRContext context;
  context.loadDialect<spatial::SpatialDialect>();
  auto module = mlir::ModuleOp::create(mlir::UnknownLoc::get(&context));
  mlir::OpBuilder builder(module.getBodyRegion());
  std::vector<mlir::Operation *> operations;
  for (size_t i = 0; i < 3; ++i) operations.push_back(
    spatial::SpatGraphComputeBatch::create(builder, module.getLoc(),
      mlir::TypeRange{}, builder.getI32IntegerAttr(3), mlir::ValueRange{}, mlir::ValueRange{}));
  const std::vector<std::vector<size_t>> resident{{2,1,5},{2,4},{0,2},{0},{4,6},{1},{2},{2},{1,4}};
  const std::vector<Cost> costs{4,1,5,4,1,2,5,2,2};
  ComputeGraph graph;
  graph.predecessors.resize(9); graph.successors.resize(9);
  for (uint32_t task = 0; task < 9; ++task) {
    ResidentWeightSet weights;
    for (size_t id : resident[task]) {
      ResidentWeight w; w.opaqueLane = id; weights.push_back(w);
    }
    graph.nodes.push_back({{operations[task % 3], task / 3, 1}, costs[task], weights, task});
  }
  IndexedResidentWeights indexed(graph);
  TEST_CHECK(!indexed.estimatePackCost({1,2,4,5,7,8}, costs, 3, 2));
  std::vector<size_t> witness;
  TEST_CHECK(indexed.estimatePackCost({1,4,7,2,5,8}, costs, 3, 2, &witness));
  TEST_CHECK(witness.size() == 6);
  auto target = hardware(1);
  target.processorCount = 4; target.residentWeightCapacity = 3;
  target.interProcessorLatencyNs.assign(16, 1);
  MergeScheduleResult schedule;
  std::string error;
  TEST_CHECK(mlir::succeeded(runPipelineScheduler(graph, schedule, 2, target, error)));
  std::vector<ResidentWeightSet> actual(4);
  for (const auto &node : graph.nodes) {
    size_t core = schedule.computeToCpuMap.lookup(node.instance);
    TEST_CHECK(core < 4);
    insertResidentWeights(actual[core], node.residentWeights);
    TEST_CHECK(actual[core].size() <= 3);
  }
  // Exhaust all core assignments under the same stage map. The initialized
  // assignment must belong to this physical-capacity reference set.
  size_t feasible = 0;
  for (size_t encoded = 0; encoded < (size_t{1} << 18); ++encoded) {
    size_t choices = encoded;
    std::vector<ResidentWeightSet> weights(4);
    std::vector<size_t> groupStages(3, 2);
    bool valid = true;
    for (size_t task = 0; task < 9 && valid; ++task) {
      size_t core = choices % 4; choices /= 4;
      size_t stage = schedule.processorStages[core], group = task % 3;
      if (groupStages[group] == 2) groupStages[group] = stage;
      valid &= groupStages[group] == stage;
      insertResidentWeights(weights[core], graph.nodes[task].residentWeights);
      valid &= weights[core].size() <= 3;
    }
    feasible += valid;
  }
  TEST_CHECK(feasible > 0);
  MergeScheduleResult alternative;
  TEST_CHECK(mlir::failed(packPipelineSchedule(graph, std::vector<size_t>(9),
    {3, 1}, target, alternative, error)));
  TEST_CHECK(error.find("nonempty compute stages") != std::string::npos);
  std::cout << "initializer_physical_feasible_assignments=" << feasible << '\n';
  module.erase();
}

static void testPhysicalReference() {
  auto target = hardware(1);
  target.processorCount = 3;
  target.sharedMemoryRequestLatencyNs.resize(3);
  target.sharedMemoryResponseLatencyNs.resize(3);
  target.interProcessorLatencyNs = {0, 2, 11, 2, 0, 7, 11, 7, 0};
  ComputeGraph graph;
  graph.predecessors.resize(3); graph.successors.resize(3);
  for (uint32_t task = 0; task < 3; ++task) {
    ResidentWeight weight; weight.opaqueLane = task % 2;
    graph.nodes.push_back({{nullptr, task, 1}, task == 1 ? 700u : 80u, {weight}, task});
    if (task) {
      graph.predecessors[task].push_back({task - 1, {64, 1}});
      graph.successors[task - 1].push_back({task, {64, 1}});
    }
  }
  double optimum = std::numeric_limits<double>::infinity();
  size_t evaluations = 0;
  for (size_t mask = 1; mask < 7; ++mask)
    for (size_t placement = 0; placement < 27; ++placement) {
      MergeScheduleResult schedule;
      for (size_t core = 0; core < 3; ++core) schedule.processorStages.push_back((mask >> core) & 1);
      size_t encoded = placement;
      std::vector<size_t> stages, slots(3);
      std::vector<ResidentWeightSet> resident(3);
      bool valid = true;
      for (const auto &node : graph.nodes) {
        size_t core = encoded % 3; encoded /= 3;
        schedule.computeToCpuMap[node.instance] = core;
        schedule.computeToCpuSlotMap[node.instance] = slots[core]++;
        stages.push_back(schedule.processorStages[core]);
        insertResidentWeights(resident[core], node.residentWeights);
        valid &= resident[core].size() <= target.residentWeightCapacity;
      }
      if (!valid || stages[0] != 0 || stages[1] < stages[0] || stages[2] < stages[1]
          || !llvm::is_contained(stages, 1)) continue;
      auto cost = queueCost(graph, schedule, target);
      optimum = std::min(optimum, std::max(*cost.primary, *cost.alternative));
      ++evaluations;
    }
  std::string error;
  MergeScheduleResult initial;
  TEST_CHECK(mlir::succeeded(runPipelineScheduler(graph, initial, 2, target, error)));
  auto evaluate = [&](const MergeScheduleResult &schedule) { return queueCost(graph, schedule, target); };
  auto search = optimizePipelineSchedule(graph, initial, 2, target, evaluate, error);
  TEST_CHECK(mlir::succeeded(search));
  size_t winner = search->costs.empty() ? 0 : selectPipelineCandidate(search->costs, error);
  auto selected = evaluate(search->schedules[winner]);
  double value = std::max(*selected.primary, *selected.alternative);
  TEST_CHECK(value + 1e-6 >= optimum);
  std::cout << "physical_reference_ns=" << optimum << " selected_ns=" << value
    << " gap_percent=" << 100 * (value / optimum - 1)
    << " physical_evaluations=" << evaluations
    << " search_evaluations=" << search->costs.size()
    << " packing_attempts=" << search->packingAttempts << '\n';
}

static void testAutoPhysicalReference() {
  for (size_t profile = 0; profile < 2; ++profile) {
    auto target = hardware(profile);
    target.processorCount = 3;
    target.sharedMemoryRequestLatencyNs.resize(3); target.sharedMemoryResponseLatencyNs.resize(3);
    target.interProcessorLatencyNs = {0, 2, 11, 2, 0, 7, 11, 7, 0};
    ComputeGraph graph;
    graph.predecessors.resize(3); graph.successors.resize(3);
    for (uint32_t task = 0; task < 3; ++task) {
      ResidentWeight weight; weight.opaqueLane = task % 2;
      graph.nodes.push_back({{nullptr, task, 1}, task == 1 ? 700u : 80u, {weight}, task});
      if (task) {
        graph.predecessors[task].push_back({task - 1, {64, 1}});
        graph.successors[task - 1].push_back({task, {64, 1}});
      }
    }
    double reference = std::numeric_limits<double>::infinity();
    size_t referenceDepth = 0, evaluations = 0;
    auto start = std::chrono::steady_clock::now();
    for (size_t depth = 1; depth <= 3; ++depth) {
      size_t maps = depth * depth * depth;
      for (size_t encodedMap = 0; encodedMap < maps; ++encodedMap)
        for (size_t placement = 0; placement < 27; ++placement) {
          MergeScheduleResult schedule;
          size_t map = encodedMap, encoded = placement;
          schedule.processorCount = 3;
          for (size_t core = 0; core < 3; ++core) { schedule.processorStages.push_back(map % depth); map /= depth; }
          std::vector<size_t> stages, slots(3);
          std::vector<ResidentWeightSet> resident(3);
          bool valid = true;
          for (const auto &node : graph.nodes) {
            size_t core = encoded % 3; encoded /= 3;
            schedule.computeToCpuMap[node.instance] = core;
            schedule.computeToCpuSlotMap[node.instance] = slots[core]++;
            stages.push_back(schedule.processorStages[core]);
            insertResidentWeights(resident[core], node.residentWeights);
            valid &= resident[core].size() <= target.residentWeightCapacity;
          }
          std::set<size_t> active(stages.begin(), stages.end());
          if (!valid || active.size() != depth || stages[0] != 0 || stages[1] < stages[0]
              || stages[1] > stages[0] + 1 || stages[2] < stages[1] || stages[2] > stages[1] + 1) continue;
          auto cost = queueCost(graph, schedule, target);
          double value = std::max(*cost.primary, *cost.alternative);
          if (value < reference) { reference = value; referenceDepth = depth; }
          ++evaluations;
        }
    }
    double referenceSeconds = std::chrono::duration<double>(std::chrono::steady_clock::now() - start).count();
    MergeScheduleResult initial;
    std::string error;
    TEST_CHECK(mlir::succeeded(packPipelineSchedule(graph, std::vector<size_t>(3), {3}, target, initial, error)));
    auto query = [&](const MergeScheduleResult &schedule) { return queueCost(graph, schedule, target); };
    start = std::chrono::steady_clock::now();
    auto search = optimizePipelineSchedule(graph, initial, 0, target, query, error);
    TEST_CHECK(mlir::succeeded(search));
    auto repeat = optimizePipelineSchedule(graph, initial, 0, target, query, error);
    TEST_CHECK(mlir::succeeded(repeat) && search->costs.size() == repeat->costs.size());
    size_t winner = selectPipelineCandidate(search->costs, error);
    TEST_CHECK(winner == selectPipelineCandidate(repeat->costs, error));
    for (size_t i = 0; i < search->schedules.size(); ++i) {
      TEST_CHECK(search->schedules[i].computeToCpuMap == repeat->schedules[i].computeToCpuMap);
      TEST_CHECK(search->costs[i].primary == repeat->costs[i].primary);
    }
    double best = std::numeric_limits<double>::infinity();
    for (const auto &cost : search->costs) best = std::min(best, std::max(*cost.primary, *cost.alternative));
    double selected = std::max(*search->costs[winner].primary, *search->costs[winner].alternative);
    TEST_CHECK(best + 1e-6 >= reference);
    std::cout << "auto_physical_profile=" << profile << " reference=" << reference
      << " reference_depth=" << referenceDepth << " best_evaluated=" << best << " selected=" << selected
      << " selected_depth=" << getPipelineStageCount(search->schedules[winner])
      << " gap_percent=" << 100 * (selected / reference - 1) << " reference_domain=all_physical_assignments_for_chain"
      << " reference_evaluations=" << evaluations << " evaluations=" << search->costs.size()
      << " packing_attempts=" << search->packingAttempts
      << " initialization_seconds=" << search->segmentCostSeconds + search->partitionStateSeconds
      << " reference_seconds=" << referenceSeconds
      << " search_repeat_seconds=" << std::chrono::duration<double>(std::chrono::steady_clock::now() - start).count() << '\n';
  }
}

static void testNonunimodalPhysicalDepthCurve() {
  TEST_CHECK(std::numeric_limits<double>::is_iec559 && std::numeric_limits<double>::digits == 53);
  TEST_CHECK(std::numeric_limits<long double>::digits >= 64);
  auto target = hardware(1);
  target.processorCount = 4;
  target.synchronizationRegisterCount = 32;
  target.sharedMemoryRequestLatencyNs.assign(4, 2);
  target.sharedMemoryResponseLatencyNs.assign(4, 3);
  target.sharedMemoryReadLatencyNs = target.sharedMemoryWriteLatencyNs = 14;
  // Manhattan distances at (0,0), (0,11), (18,0), and (5,7).
  target.interProcessorLatencyNs = {0, 11, 18, 12, 11, 0, 29, 9, 18, 29, 0, 20, 12, 9, 20, 0};
  ComputeGraph graph;
  graph.predecessors.resize(5); graph.successors.resize(5);
  const Cost work[] = {163, 46, 69, 131, 79};
  for (uint32_t task = 0; task < 5; ++task) {
    ResidentWeight weight; weight.opaqueLane = task % 2;
    graph.nodes.push_back({{nullptr, task, 1}, work[task], {weight}, task});
    if (task) {
      graph.predecessors[task].push_back({task - 1, {64, 1}});
      graph.successors[task - 1].push_back({task, {64, 1}});
    }
  }
  std::vector<double> curve(5, std::numeric_limits<double>::infinity());
  auto primary = curve, alternative = curve;
  using Ratio = std::pair<uint64_t, uint64_t>;
  std::vector<std::optional<Ratio>> certified(5);
  auto certifiedPrimary = certified, certifiedAlternative = certified;
  auto less = [](Ratio a, Ratio b) { return a.first * b.second < b.first * a.second; };
  std::vector<std::vector<size_t>> winnerCores(5), winnerStages(5);
  std::vector<size_t> evaluations(5);
  for (size_t cuts = 0; cuts < 16; ++cuts) {
    std::vector<size_t> stages{0};
    for (size_t task = 1; task < 5; ++task)
      stages.push_back(stages.back() + ((cuts >> (task - 1)) & 1));
    size_t depth = stages.back() + 1;
    if (depth > 4) continue;
    for (size_t placement = 0; placement < 1024; ++placement) {
      MergeScheduleResult schedule;
      schedule.processorCount = 4;
      size_t encoded = placement;
      std::vector<size_t> assigned(4, depth), slots(4), idle;
      bool valid = true;
      for (size_t task = 0; task < 5; ++task) {
        size_t core = encoded % 4; encoded /= 4;
        if (assigned[core] != depth && assigned[core] != stages[task]) { valid = false; break; }
        assigned[core] = stages[task];
        schedule.computeToCpuMap[graph.nodes[task].instance] = core;
        schedule.computeToCpuSlotMap[graph.nodes[task].instance] = slots[core]++;
      }
      if (!valid) continue;
      // There are only two resident weights; both fit on every core.
      for (size_t core = 0; core < 4; ++core) if (assigned[core] == depth) idle.push_back(core);
      size_t maps = 1;
      for (size_t i = 0; i < idle.size(); ++i) maps *= depth;
      for (size_t map = 0; map < maps; ++map) {
        schedule.processorStages = assigned;
        size_t digits = map;
        for (size_t core : idle) { schedule.processorStages[core] = digits % depth; digits /= depth; }
        auto cost = queueCost(graph, schedule, target, true);
        Ratio worst{0, 1};
        for (auto [policy, proof] : llvm::enumerate(cost.recurrence)) {
          TEST_CHECK(proof.proved() && proof.periodGenerations > 0 && proof.periodGenerations <= 1024);
          double value = *proof.intervalNs;
          TEST_CHECK(value > 0 && value < (uint64_t{1} << 30));
          // Queue timestamps are checked integers. Adjacent binary64 values
          // enclose the final rounded duration/period division, independently of MCR.
          long double low = std::nextafter(value, -std::numeric_limits<double>::infinity());
          long double high = std::nextafter(value, std::numeric_limits<double>::infinity());
          uint64_t lower = std::ceil(low * proof.periodGenerations);
          uint64_t upper = std::floor(high * proof.periodGenerations);
          TEST_CHECK(lower == upper && upper < (uint64_t{1} << 40));
          uint64_t divisor = std::gcd(lower, proof.periodGenerations);
          Ratio exact{lower / divisor, proof.periodGenerations / divisor};
          auto &minimum = (policy ? certifiedAlternative : certifiedPrimary)[depth];
          if (!minimum || less(exact, *minimum)) minimum = exact;
          if (less(worst, exact)) worst = exact;
        }
        if (!certified[depth] || less(worst, *certified[depth])) certified[depth] = worst;
        double value = std::max(*cost.primary, *cost.alternative);
        if (value < curve[depth]) {
          curve[depth] = value;
          winnerStages[depth] = schedule.processorStages;
          winnerCores[depth].clear();
          for (const auto &node : graph.nodes)
            winnerCores[depth].push_back(schedule.computeToCpuMap.lookup(node.instance));
        }
        primary[depth] = std::min(primary[depth], *cost.primary);
        alternative[depth] = std::min(alternative[depth], *cost.alternative);
        ++evaluations[depth];
      }
    }
  }
  TEST_CHECK(curve[1] == 496 && curve[2] == 304 && curve[3] == 305 && curve[4] == 304.5);
  TEST_CHECK(primary == curve && alternative == curve);
  TEST_CHECK(certified[3] == Ratio(305, 1) && certified[4] == Ratio(609, 2));
  TEST_CHECK(certifiedPrimary == certified && certifiedAlternative == certified);
  TEST_CHECK(less(*certified[4], *certified[3]));
  TEST_CHECK(evaluations == std::vector<size_t>({0, 1024, 1728, 792, 96}));
  std::cout << "nonunimodal_physical_depth_reference";
  for (size_t depth = 1; depth <= 4; ++depth)
    std::cout << " p" << depth << "_ns=" << curve[depth] << " primary=" << primary[depth]
      << " alternative=" << alternative[depth] << " exact=" << certified[depth]->first
      << '/' << certified[depth]->second << " evaluations=" << evaluations[depth];
  std::cout << " reference_domain=all_chain_partitions_and_physical_assignments_with_restart\n";
  for (size_t depth = 1; depth <= 4; ++depth) {
    std::cout << "physical_depth_winner=" << depth << " core_stages=";
    for (size_t stage : winnerStages[depth]) std::cout << stage << ',';
    std::cout << " task_cores=";
    for (size_t core : winnerCores[depth]) std::cout << core << ',';
    std::cout << '\n';
  }
}

static void testWorseningIntermediate() {
  // A controlled cost oracle isolates the search's non-monotone landscape.
  // Hardware queue objectives are covered separately by runReference().
  auto target = hardware(0);
  target.processorCount = 8;
  target.interProcessorLatencyNs.assign(64, 0);
  ComputeGraph graph;
  graph.predecessors.resize(5); graph.successors.resize(5);
  for (uint32_t task = 0; task < 5; ++task)
    graph.nodes.push_back({{nullptr, task, 1}, task == 4 ? 0u : 10u, {}, task});
  auto evaluate = [&](const MergeScheduleResult &schedule) {
    bool firstStage = true;
    for (size_t i = 0; i < graph.nodes.size(); ++i)
      firstStage &= schedule.processorStages[schedule.computeToCpuMap.lookup(graph.nodes[i].instance)]
        == (i == 4 ? 1 : 0);
    size_t cores = llvm::count(schedule.processorStages, 0);
    double value = !firstStage ? 120 : cores == 1 ? 80 : cores >= 6 ? 100 : 110;
    WholeScheduleCost cost;
    cost.primary = cost.alternative = value;
    cost.loweringSupport = StaticLoweringSupport::Supported;
    cost.selectionMargin = kPipelineSelectionMargin;
    for (auto &proof : cost.recurrence)
      proof = {ScheduleRecurrenceEvidence::Proved, 1, value, {}};
    return cost;
  };
  std::string error;
  MergeScheduleResult initial, intermediate;
  const std::vector<size_t> stages{0, 0, 0, 0, 1};
  TEST_CHECK(mlir::succeeded(packPipelineSchedule(graph, stages, {7, 1}, target, initial, error)));
  TEST_CHECK(evaluate(initial).primary == 100);
  TEST_CHECK(mlir::succeeded(packPipelineSchedule(graph, stages, {3, 5}, target, intermediate, error)));
  TEST_CHECK(evaluate(intermediate).primary == 110);
  double reference = std::numeric_limits<double>::infinity();
  for (size_t cores = 1; cores < 8; ++cores) {
    MergeScheduleResult schedule;
    TEST_CHECK(mlir::succeeded(packPipelineSchedule(graph, stages, {cores, 8 - cores}, target, schedule, error)));
    reference = std::min(reference, *evaluate(schedule).primary);
  }
  TEST_CHECK(reference == 80);
  auto search = optimizePipelineSchedule(graph, initial, 2, target, evaluate, error);
  TEST_CHECK(mlir::succeeded(search));
  size_t winner = selectPipelineCandidate(search->costs, error);
  double selected = *search->costs[winner].primary;
  size_t depth = 0;
  for (size_t i = winner; i; i = search->parents[i]) ++depth;
  TEST_CHECK(depth >= 2);
  sawSuccessiveMoves = true;
  TEST_CHECK(selected >= reference && selected <= 100);
  TEST_CHECK(search->costs.size() <= kPipelineQueueEvaluationBudget);
  std::cout << "non_monotone_oracle_reference=" << reference << " selected=" << selected
    << " gap_percent=" << 100 * (selected / reference - 1)
    << " evaluations=" << search->costs.size()
    << " packing_attempts=" << search->packingAttempts << '\n';
}

static void testMultipleWorseningAllocations() {
  mlir::MLIRContext context;
  context.loadDialect<spatial::SpatialDialect>();
  auto module = mlir::ModuleOp::create(mlir::UnknownLoc::get(&context));
  mlir::OpBuilder builder(module.getBodyRegion());
  auto first = spatial::SpatGraphComputeBatch::create(builder, module.getLoc(),
    mlir::TypeRange{}, builder.getI32IntegerAttr(4), mlir::ValueRange{}, mlir::ValueRange{});
  auto last = spatial::SpatGraphComputeBatch::create(builder, module.getLoc(),
    mlir::TypeRange{}, builder.getI32IntegerAttr(1), mlir::ValueRange{}, mlir::ValueRange{});
  ComputeGraph graph;
  graph.predecessors.resize(5); graph.successors.resize(5);
  for (uint32_t task = 0; task < 5; ++task) {
    graph.nodes.push_back({{task < 4 ? first.getOperation() : last.getOperation(),
      task < 4 ? task : 0, 1}, task < 4 ? 10u : 0u, {}, task});
    if (task < 4) {
      graph.predecessors[4].push_back({task, {0, 0}});
      graph.successors[task].push_back({4, {0, 0}});
    }
  }
  auto target = hardware(0);
  target.processorCount = 32; target.interProcessorLatencyNs.assign(1024, 0);
  auto query = [&](const MergeScheduleResult &schedule) {
    size_t cores = llvm::count(schedule.processorStages, 0);
    double value = cores == 1 ? 80 : cores == 31 ? 100 : 110 + cores;
    WholeScheduleCost cost;
    cost.primary = cost.alternative = value;
    cost.loweringSupport = StaticLoweringSupport::Supported;
    cost.selectionMargin = kPipelineSelectionMargin;
    for (auto &proof : cost.recurrence) proof = {ScheduleRecurrenceEvidence::Proved, 1, value, {}};
    return cost;
  };
  std::string error;
  MergeScheduleResult initial;
  const std::vector<size_t> stages{0, 0, 0, 0, 1};
  double reference = std::numeric_limits<double>::infinity();
  for (size_t cores = 1; cores < 32; ++cores) {
    MergeScheduleResult schedule;
    TEST_CHECK(mlir::succeeded(packPipelineSchedule(graph, stages, {cores, 32 - cores}, target, schedule, error)));
    reference = std::min(reference, *query(schedule).primary);
  }
  TEST_CHECK(mlir::succeeded(packPipelineSchedule(graph, stages, {31, 1}, target, initial, error)));
  auto search = optimizePipelineSchedule(graph, initial, 2, target, query, error);
  auto repeat = optimizePipelineSchedule(graph, initial, 2, target, query, error);
  TEST_CHECK(mlir::succeeded(search) && mlir::succeeded(repeat));
  TEST_CHECK(search->parents == repeat->parents && search->costs.size() == repeat->costs.size());
  for (size_t i = 0; i < search->schedules.size(); ++i) {
    TEST_CHECK(search->schedules[i].computeToCpuMap == repeat->schedules[i].computeToCpuMap);
    TEST_CHECK(search->costs[i].primary == repeat->costs[i].primary);
  }
  size_t winner = selectPipelineCandidate(search->costs, error), path = 0;
  TEST_CHECK(reference == 80 && *search->costs[winner].primary == reference);
  for (size_t i = winner; i; i = search->parents[i]) {
    if (i != winner) TEST_CHECK(*search->costs[i].primary > 100);
    ++path;
  }
  TEST_CHECK(path >= 3);
  std::cout << "multiple_worsening_reference=" << reference << " selected=" << *search->costs[winner].primary
    << " gap_percent=0 path_length=" << path << " evaluations=" << search->costs.size()
    << " packing_attempts=" << search->packingAttempts
    << " reference_domain=all_31_core_budgets_for_fixed_operation_groups\n";
  module.erase();
}

static void runReference(size_t profile, size_t fixture) {
  auto target = hardware(profile);
  ComputeGraph graph;
  const std::vector<std::vector<Cost>> work{{900, 70, 800, 90, 700},
    {30, 700, 110, 600, 90}, {1000, 50, 40, 900, 60},
    {90, 80, 100, 110, 120}, {100, 800, 100, 800, 100},
    {900, 800, 40, 30, 20}};
  size_t count = work[fixture].size();
  graph.predecessors.resize(count); graph.successors.resize(count);
  for (size_t task = 0; task < count; ++task) {
    ResidentWeight weight; weight.opaqueLane = fixture == 4 ? task % 3 : task % 2;
    graph.nodes.push_back({{nullptr, static_cast<uint32_t>(task), 1}, work[fixture][task],
      {weight}, task});
    auto edge = [&](size_t source) {
      graph.predecessors[task].push_back({source, {64, 1}});
      graph.successors[source].push_back({task, {64, 1}});
      graph.publicationTransfers.push_back({source, task, {}, 0, {}, 64, 1});
    };
    if (fixture == 1) {
      if (task == 1 || task == 2) edge(0);
      if (task == 3) { edge(1); edge(2); }
      if (task == 4) edge(3);
    } else if (task) edge(task - 1);
    if (fixture == 2 && task == 4) edge(1);
  }
  if (fixture == 3) {
    target.sharedMemoryReadLatencyNs = 400;
    target.sharedMemoryWriteLatencyNs = 600;
  }
  std::string error;
  MergeScheduleResult initial;
  TEST_CHECK(mlir::succeeded(runPipelineScheduler(graph, initial, 3, target, error)));
  auto evaluate = [&](const MergeScheduleResult &schedule) { return queueCost(graph, schedule, target); };
  auto begin = std::chrono::steady_clock::now();
  auto search = optimizePipelineSchedule(graph, initial, 3, target, evaluate, error);
  double searchSeconds = std::chrono::duration<double>(std::chrono::steady_clock::now() - begin).count();
  TEST_CHECK(mlir::succeeded(search));
  auto repeat = optimizePipelineSchedule(graph, initial, 3, target, evaluate, error);
  TEST_CHECK(mlir::succeeded(repeat) && search->schedules.size() == repeat->schedules.size());
  TEST_CHECK(search->parents == repeat->parents && search->expansions == repeat->expansions);
  TEST_CHECK(search->packingAttempts == repeat->packingAttempts
    && search->consideredAssignments == repeat->consideredAssignments);
  for (size_t i = 0; i < search->schedules.size(); ++i) {
    TEST_CHECK(search->schedules[i].computeToCpuMap == repeat->schedules[i].computeToCpuMap);
    TEST_CHECK(search->schedules[i].computeToCpuSlotMap == repeat->schedules[i].computeToCpuSlotMap);
    TEST_CHECK(search->schedules[i].processorStages == repeat->schedules[i].processorStages);
    TEST_CHECK(search->costs.empty() || search->costs[i].primary == repeat->costs[i].primary);
    TEST_CHECK(search->costs.empty() || search->costs[i].alternative == repeat->costs[i].alternative);
  }
  double reference = std::numeric_limits<double>::infinity();
  size_t referenceEvaluations = 0;
  std::vector<size_t> stages(count), budgets(3, 1);
  std::function<void(size_t)> partitions = [&](size_t task) {
    if (task != count) {
      for (size_t stage = 0; stage <= (task ? 2 : 0); ++stage) {
        stages[task] = stage; partitions(task + 1);
      }
      return;
    }
    std::vector<bool> used(3);
    for (size_t stage : stages) used[stage] = true;
    if (!llvm::all_of(used, [](bool value) { return value; })) return;
    for (size_t core0 = 1; core0 + 1 < target.processorCount; ++core0)
      for (size_t core1 = 1; core0 + core1 < target.processorCount; ++core1) {
        budgets = {core0, core1, target.processorCount - core0 - core1};
        MergeScheduleResult schedule;
        if (mlir::failed(packPipelineSchedule(graph, stages, budgets, target, schedule, error))) continue;
        auto cost = evaluate(schedule); ++referenceEvaluations;
        reference = std::min(reference, std::max(*cost.primary, *cost.alternative));
      }
  };
  auto referenceBegin = std::chrono::steady_clock::now();
  partitions(0);
  double referenceSeconds = std::chrono::duration<double>(std::chrono::steady_clock::now() - referenceBegin).count();
  std::string reason;
  size_t winner = search->costs.empty() ? 0 : selectPipelineCandidate(search->costs, reason);
  auto selected = evaluate(search->schedules[winner]);
  double cost = std::max(*selected.primary, *selected.alternative);
  double bestEvaluated = cost;
  for (const auto &estimate : search->costs)
    bestEvaluated = std::min(bestEvaluated, std::max(*estimate.primary, *estimate.alternative));
  TEST_CHECK(std::isfinite(reference) && cost + 1e-6 >= reference);
  TEST_CHECK(bestEvaluated + 1e-6 >= reference);
  TEST_CHECK(cost <= reference * (1 + kPipelineSelectionMargin) + 1e-6);
  TEST_CHECK(search->costs.size() <= kPipelineQueueEvaluationBudget
    && search->packingAttempts <= 4 * kPipelineQueueEvaluationBudget
    && search->consideredAssignments <= 4 * kPipelineQueueEvaluationBudget);
  size_t depth = 0;
  for (size_t i = winner; i; i = search->parents[i]) {
    ++depth;
    const auto &current = search->schedules[i], &parent = search->schedules[search->parents[i]];
    bool cutChanged = false;
    for (const auto &node : graph.nodes)
      cutChanged |= current.processorStages[current.computeToCpuMap.lookup(node.instance)]
        != parent.processorStages[parent.computeToCpuMap.lookup(node.instance)];
    std::vector<size_t> currentBudgets(3), parentBudgets(3);
    for (size_t core = 0; core < target.processorCount; ++core) {
      ++currentBudgets[current.processorStages[core]];
      ++parentBudgets[parent.processorStages[core]];
    }
    sawJointMove |= cutChanged && currentBudgets != parentBudgets;
  }
  sawSuccessiveMoves |= depth >= 2;
  for (size_t i = 1; i < search->schedules.size(); ++i) {
    const auto &child = search->schedules[i], &parent = search->schedules[search->parents[i]];
    bool cut = false;
    for (const auto &node : graph.nodes)
      cut |= child.processorStages[child.computeToCpuMap.lookup(node.instance)]
        != parent.processorStages[parent.computeToCpuMap.lookup(node.instance)];
    std::vector<size_t> a(3), b(3);
    for (size_t stage : child.processorStages) ++a[stage];
    for (size_t stage : parent.processorStages) ++b[stage];
    sawJointMove |= cut && a != b;
  }
  std::cout << "profile=" << profile << " fixture=" << fixture
    << " reference_ns=" << reference << " selected_ns=" << cost
    << " gap_percent=" << 100 * (cost / reference - 1)
    << " best_evaluated_ns=" << bestEvaluated
    << " best_evaluated_gap_percent=" << 100 * (bestEvaluated / reference - 1)
    << " reference_scope=logical_with_deterministic_packing"
    << " reference_evaluations=" << referenceEvaluations
    << " search_evaluations=" << search->costs.size()
    << " packing_attempts=" << search->packingAttempts << " refinement_depth=" << depth
    << " search_seconds=" << searchSeconds << " reference_seconds=" << referenceSeconds
    << " total_seconds=" << std::chrono::duration<double>(std::chrono::steady_clock::now() - begin).count()
    << '\n';
  auto unavailable = [&](const MergeScheduleResult &) { WholeScheduleCost value; value.reason = "test unavailable timing"; return value; };
  auto unknown = optimizePipelineSchedule(graph, initial, 3, target, unavailable, error);
  TEST_CHECK(search->costs.empty() || (mlir::failed(unknown) && error.find("test unavailable timing") != std::string::npos));
  auto unproved = [&](const MergeScheduleResult &schedule) {
    auto cost = evaluate(schedule);
    cost.recurrence[0].status = ScheduleRecurrenceEvidence::Unproved;
    cost.reason = "test unproved recurrence";
    return cost;
  };
  auto invalid = optimizePipelineSchedule(graph, initial, 3, target, unproved, error);
  TEST_CHECK(search->costs.empty() || (mlir::failed(invalid) && error.find("test unproved recurrence") != std::string::npos));
  if (!profile && !fixture) {
    auto badGuidance = [&](const MergeScheduleResult &schedule) {
      auto cost = evaluate(schedule); cost.processorTaskWorkNs.resize(target.processorCount + 1);
      return cost;
    };
    TEST_CHECK(mlir::failed(optimizePipelineSchedule(graph, initial, 3, target, badGuidance, error)));
    TEST_CHECK(error.find("physical processor count") != std::string::npos);
    TEST_CHECK(mlir::failed(optimizePipelineSchedule(graph, initial, 3, target, evaluate, error, 0)));
    TEST_CHECK(mlir::failed(optimizePipelineSchedule(graph, initial, 3, target, evaluate, error, kPipelineQueueEvaluationBudget + 1)));
  }
}

// Exhaustive small-grid DP oracle. Physical placement is checked separately.
static void testSharedPartitionStates() {
  for (size_t profile = 0; profile < 2; ++profile) {
    ComputeGraph graph;
    const std::vector<Cost> work{10, 30, 5, 50, 20, 15};
    graph.predecessors.resize(work.size()); graph.successors.resize(work.size());
    for (uint32_t task = 0; task < work.size(); ++task) {
      ResidentWeight weight; weight.opaqueLane = task % (profile + 2);
      graph.nodes.push_back({{nullptr, task, 1}, work[task], {weight}, task});
      if (task) {
        graph.predecessors[task].push_back({task - 1, {}});
        graph.successors[task - 1].push_back({task, {}});
      }
    }
    if (profile) {
      graph.predecessors[4].push_back({1, {}});
      graph.successors[1].push_back({4, {}});
    }
    auto target = hardware(profile);
    MergeScheduleResult peft;
    std::string error;
    TEST_CHECK(mlir::succeeded(packPipelineSchedule(graph, std::vector<size_t>(work.size()),
      {target.processorCount}, target, peft, error)));
    auto equalCost = [](const MergeScheduleResult &) {
      WholeScheduleCost cost;
      cost.primary = cost.alternative = 100;
      cost.selectionMargin = kPipelineSelectionMargin;
      cost.loweringSupport = StaticLoweringSupport::Supported;
      for (auto &proof : cost.recurrence) proof = {ScheduleRecurrenceEvidence::Proved, 1, 100, {}};
      return cost;
    };
    auto search = optimizePipelineSchedule(graph, peft, 0, target, equalCost, error);
    if (mlir::failed(search)) std::cerr << error << '\n';
    TEST_CHECK(mlir::succeeded(search));
    for (size_t i = 1; i < search->initialSchedules; ++i) {
      size_t depth = getPipelineStageCount(search->schedules[i]);
      MergeScheduleResult independentlyInitialized;
      TEST_CHECK(mlir::succeeded(runPipelineScheduler(graph, independentlyInitialized, depth, target, error)));
      TEST_CHECK(search->schedules[i].processorStages == independentlyInitialized.processorStages);
      TEST_CHECK(search->schedules[i].computeToCpuMap == independentlyInitialized.computeToCpuMap);
      TEST_CHECK(search->schedules[i].computeToCpuSlotMap == independentlyInitialized.computeToCpuSlotMap);
    }
    IndexedResidentWeights weights(graph);
    for (const auto &entry : search->depthPriorities) {
      size_t depth = entry.first;
      Cost actual = entry.second;
      Cost reference = std::numeric_limits<Cost>::max();
      size_t evaluated = 0;
      for (size_t mask = 0; mask < (size_t{1} << (work.size() - 1)); ++mask) {
        std::vector<size_t> stages(work.size());
        for (size_t task = 1; task < work.size(); ++task)
          stages[task] = stages[task - 1] + ((mask >> (task - 1)) & 1);
        if (stages.back() + 1 != depth) continue;
        bool legal = true;
        for (size_t task = 0; task < work.size(); ++task)
          for (auto [before, transfer] : graph.predecessors[task])
            legal &= stages[task] <= stages[before] + 1;
        if (!legal) continue;
        std::vector<std::vector<size_t>> tasks(depth);
        for (size_t task = 0; task < work.size(); ++task) tasks[stages[task]].push_back(task);
        std::function<void(size_t, size_t, Cost)> allocate = [&](size_t stage, size_t remaining, Cost objective) {
          if (stage == depth) {
            if (!remaining) { reference = std::min(reference, objective); ++evaluated; }
            return;
          }
          const auto &members = tasks[stage];
          std::vector<Cost> ends(work.size());
          Cost span = 0, largest = 0;
          for (size_t task : members) {
            ends[task] = checkedAdd(ends[task], work[task]);
            span = std::max(span, ends[task]); largest = std::max(largest, work[task]);
            for (auto [next, transfer] : graph.successors[task]) ends[next] = std::max(ends[next], ends[task]);
          }
          Cost packed = std::numeric_limits<Cost>::max();
          for (size_t cores = 1; cores + depth - stage - 1 <= remaining; ++cores) {
            if (auto cost = weights.estimatePackCost(members, work, target.residentWeightCapacity, cores))
              packed = std::min(packed, *cost);
            if (packed != std::numeric_limits<Cost>::max())
              allocate(stage + 1, remaining - cores, std::max(objective, checkedAdd(span - largest, packed)));
          }
        };
        allocate(0, target.processorCount, 0);
      }
      TEST_CHECK(reference == actual);
      std::cout << "shared_dp_profile=" << profile << " depth=" << depth
        << " reference_priority=" << reference << " actual_priority=" << actual
        << " reference_allocations=" << evaluated << " dp_transitions=" << search->partitionTransitions << '\n';
    }
  }
}

static void testDepthDomain() {
  ComputeGraph graph;
  graph.predecessors.resize(12); graph.successors.resize(12);
  for (uint32_t i = 0; i < 12; ++i) {
    graph.nodes.push_back({{nullptr, i, 1}, 100, {}, i});
    if (i) {
      graph.predecessors[i].push_back({i - 1, {}});
      graph.successors[i - 1].push_back({i, {}});
    }
  }
  auto target = hardware(0);
  target.processorCount = 16;
  target.interProcessorLatencyNs.assign(256, 1);
  std::string error;
  for (size_t depth : {2, 4, 8}) {
    MergeScheduleResult schedule;
    TEST_CHECK(mlir::succeeded(runPipelineScheduler(graph, schedule, depth, target, error)));
    TEST_CHECK(mlir::succeeded(verifyPipelineStageCount(schedule, depth, error)));
    std::set<size_t> active;
    for (const auto &node : graph.nodes)
      active.insert(schedule.processorStages[schedule.computeToCpuMap.lookup(node.instance)]);
    TEST_CHECK(active.size() == depth && *active.begin() == 0 && *active.rbegin() == depth - 1);
    auto broken = schedule;
    for (const auto &[instance, core] : broken.computeToCpuMap)
      if (broken.processorStages[core] == depth - 1) broken.processorStages[core] = depth - 2;
    TEST_CHECK(mlir::failed(verifyPipelineStageCount(broken, depth, error)));
  }
  MergeScheduleResult invalid;
  TEST_CHECK(mlir::failed(runPipelineScheduler(graph, invalid, 13, target, error)));
  TEST_CHECK(error.find("compute operation groups") != std::string::npos);
  TEST_CHECK(mlir::failed(runPipelineScheduler(graph, invalid, 17, target, error)));

  MergeScheduleResult peft;
  peft.processorCount = target.processorCount;
  for (size_t i = 0; i < graph.nodes.size(); ++i) {
    peft.computeToCpuMap[graph.nodes[i].instance] = 0;
    peft.computeToCpuSlotMap[graph.nodes[i].instance] = i;
  }
  auto incomplete = peft;
  incomplete.processorStages = {0};
  TEST_CHECK(mlir::failed(verifyPipelineStageCount(incomplete, 1, error)));
  std::vector<size_t> winners;
  for (double boundaryCost : {1000., 300., 50., 0.}) {
    // A closed synthetic target contract isolates depth-domain selection.
    // Complete queue objectives and physical placements are tested separately.
    auto query = [&](const MergeScheduleResult &schedule) {
      size_t depth = getPipelineStageCount(schedule);
      double interval = 1200. / depth + boundaryCost * (depth - 1);
      WholeScheduleCost cost;
      cost.primary = cost.alternative = interval;
      cost.loweringSupport = StaticLoweringSupport::Supported;
      cost.selectionMargin = kPipelineSelectionMargin;
      for (auto &proof : cost.recurrence)
        proof = {ScheduleRecurrenceEvidence::Proved, 1, interval, {}};
      return cost;
    };
    auto search = optimizePipelineSchedule(graph, peft, 0, target, query, error);
    TEST_CHECK(mlir::succeeded(search));
    auto repeat = optimizePipelineSchedule(graph, peft, 0, target, query, error);
    TEST_CHECK(mlir::succeeded(repeat) && repeat->schedules.size() == search->schedules.size());
    TEST_CHECK(search->costs.size() <= kPipelineQueueEvaluationBudget);
    size_t winner = selectPipelineCandidate(search->costs, error);
    size_t second = selectPipelineCandidate(repeat->costs, error);
    TEST_CHECK(winner == second);
    for (size_t i = 0; i < search->schedules.size(); ++i) {
      TEST_CHECK(search->schedules[i].processorStages == repeat->schedules[i].processorStages);
      TEST_CHECK(search->schedules[i].computeToCpuMap == repeat->schedules[i].computeToCpuMap);
      TEST_CHECK(search->costs[i].primary == repeat->costs[i].primary);
      TEST_CHECK(mlir::succeeded(verifyPipelineStageCount(search->schedules[i],
        getPipelineStageCount(search->schedules[i]), error)));
    }
    winners.push_back(getPipelineStageCount(search->schedules[winner]));
    double reference = std::numeric_limits<double>::infinity();
    for (size_t depth = 1; depth <= 12; ++depth)
      reference = std::min(reference, 1200. / depth + boundaryCost * (depth - 1));
    std::cout << "auto_target_boundary_ns=" << boundaryCost << " depth=" << winners.back()
      << " reference_ns=" << reference << " selected_ns=" << *search->costs[winner].primary
      << " evaluations=" << search->costs.size() << '\n';
  }
  TEST_CHECK(winners == std::vector<size_t>({1, 2, 5, 12}));
  graph.predecessors[11].push_back({0, {}});
  graph.successors[0].push_back({11, {}});
  TEST_CHECK(mlir::failed(runPipelineScheduler(graph, invalid, 8, target, error)));
  auto unavailable = [](const MergeScheduleResult &) {
    WholeScheduleCost cost; cost.reason = "unavailable auto test cost"; return cost;
  };
  TEST_CHECK(mlir::failed(optimizePipelineSchedule(graph, peft, 0, target, unavailable, error)));
  TEST_CHECK(error.find("unavailable auto test cost") != std::string::npos);
}

int main() {
  testPackingCostCacheEviction();
  testPackingCurves();
  testSharedPartitionStates();
  testDepthDomain();
  testInitializerPackingWitness();
  testPackingFeasibility();
  for (uint32_t first : {19u, 21u}) testQuantizedFeasibility(first);
  testPhysicalReference();
  testAutoPhysicalReference();
  testNonunimodalPhysicalDepthCurve();
  testWorseningIntermediate();
  testMultipleWorseningAllocations();
  for (size_t profile = 0; profile < 2; ++profile)
    for (size_t fixture = 0; fixture < 6; ++fixture) runReference(profile, fixture);
  TEST_CHECK(sawJointMove && sawSuccessiveMoves);
}